//! Cyborg is a high-performance, modern, experimental rendering engine written
//! in Rust.
//!
//! # Features
//!
//! - `legion`: Enables [Legion](https://github.com/amethyst/legion) integration
//!   for Cyborg. Disabled by default.

use parking_lot::Mutex;
use rayon::prelude::*;
use std::sync::Arc;

pub mod camera;
pub mod pass;
pub mod phase;
pub mod scene;
pub mod shader;
pub mod staging;
pub mod storage;
pub mod viewport;

#[cfg(feature = "legion")]
pub mod legion;

use camera::Camera;
use pass::*;
use phase::*;
use viewport::Viewport;

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
pub struct ViewportUniform {
    pub eye: [f32; 4],
    pub vp: [[f32; 4]; 4],
}

impl ViewportUniform {
    pub fn from_camera(camera: &Camera) -> Self {
        Self {
            eye: camera.eye,
            vp: camera.vp,
        }
    }
}

impl ViewportUniform {
    pub fn binding_size() -> wgpu::BufferSize {
        let size = std::mem::size_of::<Self>() as u64;
        size.try_into().unwrap()
    }
}

pub struct RenderLayouts {
    pub device: Arc<wgpu::Device>,
    pub bind_viewport: wgpu::BindGroupLayout,
}

impl RenderLayouts {
    pub fn new_arc(device: Arc<wgpu::Device>) -> Arc<Self> {
        let bind_viewport = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("viewport"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: Some(ViewportUniform::binding_size()),
                },
                count: None,
            }],
        });

        Arc::new(Self {
            device,
            bind_viewport,
        })
    }
}

pub struct FrameData {
    pub frame_index: usize,
    pub bind_viewport: wgpu::BindGroup,
    pub viewport_uniform: wgpu::Buffer,
}

impl FrameData {
    pub fn new(frame_index: usize, layouts: &RenderLayouts) -> Self {
        let viewport_uniform = layouts.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("viewport"),
            size: ViewportUniform::binding_size().into(),
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let bind_viewport = layouts
            .device
            .create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("viewport"),
                layout: &layouts.bind_viewport,
                entries: &[wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                        buffer: &viewport_uniform,
                        offset: 0,
                        size: None,
                    }),
                }],
            });

        Self {
            frame_index,
            bind_viewport,
            viewport_uniform,
        }
    }

    pub fn make_phase_data<'a>(
        &'a self,
        phase: Phase,
        viewport_data: &'a ViewportData,
    ) -> IndexedPhaseData<'a> {
        IndexedPhaseData {
            phase,
            frame_data: self.frame_index,
            viewport_data,
            bind_viewport: &self.bind_viewport,
        }
    }
}

pub struct Renderer {
    device: Arc<wgpu::Device>,
    queue: Arc<wgpu::Queue>,
    layouts: Arc<RenderLayouts>,
    frames_in_flight: usize,
    frame_datas: Vec<FrameData>,
    frame_index: usize,
    render_passes: Vec<Box<dyn RenderPassBoxTrait>>,
}

impl Renderer {
    pub fn new(device: Arc<wgpu::Device>, queue: Arc<wgpu::Queue>) -> Self {
        let layouts = RenderLayouts::new_arc(device.clone());

        let frames_in_flight = 1;

        let mut frame_datas = Vec::with_capacity(frames_in_flight);
        for frame_index in 0..frames_in_flight {
            frame_datas.push(FrameData::new(frame_index, &layouts));
        }

        Self {
            device,
            layouts,
            queue,
            frames_in_flight,
            frame_datas,
            frame_index: 0,
            render_passes: Vec::new(),
        }
    }

    pub fn get_device(&self) -> &Arc<wgpu::Device> {
        &self.device
    }

    pub fn get_layouts(&self) -> &Arc<RenderLayouts> {
        &self.layouts
    }

    pub fn get_frames_in_flight(&self) -> usize {
        self.frames_in_flight
    }

    pub fn render<'a>(
        &mut self,
        passes: &mut [&mut dyn RenderPassBoxTrait],
        target: &dyn Viewport,
        camera: &Camera,
    ) {
        puffin::profile_function!();

        self.frame_index += 1;
        if self.frame_index >= self.frame_datas.len() {
            self.frame_index = 0;
        }

        let frame_index = self.frame_index;
        let frame_data = &self.frame_datas[frame_index];

        let viewport_uniform = ViewportUniform::from_camera(camera);
        self.queue.write_buffer(
            &frame_data.viewport_uniform,
            0,
            bytemuck::cast_slice(&[viewport_uniform]),
        );

        let phase_passes = multimap::MultiMap::<Phase, usize>::new();
        let phase_passes = std::sync::Mutex::new(phase_passes);

        passes
            .par_iter_mut()
            .enumerate()
            .for_each(|(pass_index, rp)| {
                let mut phases_buf = Vec::new();
                phases_buf.clear();
                rp.begin_frame(frame_index, &mut phases_buf, &self.queue);

                let mut passes = phase_passes.lock().unwrap();
                for phase in phases_buf.into_iter() {
                    passes.insert(phase, pass_index);
                }
            });

        let phase_passes = phase_passes.into_inner().unwrap();
        let viewport = ViewportData;

        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("Render Encoder"),
            });

        if let Some(upload) = phase_passes.get_vec(&Phase::Upload) {
            upload.iter().for_each(|pass_index| {
                let phase_data = frame_data.make_phase_data(Phase::Upload, &viewport);
                let pass = &passes[*pass_index];
                pass.record_commands(phase_data, &mut encoder);
            });
        }

        if let Some(skinning) = phase_passes.get_vec(&Phase::Skinning) {
            let mut cmds = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("Skinning Phase"),
            });

            skinning.iter().for_each(|pass_index| {
                let phase_data = frame_data.make_phase_data(Phase::Skinning, &viewport);
                let pass = &passes[*pass_index];
                pass.record_compute(phase_data, &mut cmds);
            })
        }

        let record_render = |phase| {
            let cmds = Mutex::new(Vec::new());
            if let Some(phase_passes) = phase_passes.get_vec(&phase) {
                phase_passes.par_iter().for_each(|pass_index| {
                    let phase_data = frame_data.make_phase_data(phase, &viewport);
                    let pass = &passes[*pass_index];

                    if let Some(cmd) = pass.record_render(phase_data) {
                        cmds.lock().push(cmd);
                    }
                })
            }

            cmds.into_inner()
        };

        // TODO: parallelize each phase's record_render
        let depth_cmds = record_render(Phase::Depth);
        let opaque_cmds = record_render(Phase::Opaque);
        let overlay_cmds = record_render(Phase::Overlay);

        {
            let target_views = target.get_views();

            let mut rp = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("Render Pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: target_views.output,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                        store: true,
                    },
                })],
                depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
                    view: target_views.depth,
                    depth_ops: Some(wgpu::Operations {
                        load: wgpu::LoadOp::Clear(1.0),
                        store: true,
                    }),
                    stencil_ops: None,
                }),
            });

            rp.execute_bundles(depth_cmds.iter());
            rp.execute_bundles(opaque_cmds.iter());
            rp.execute_bundles(overlay_cmds.iter());
        }

        self.queue.submit(std::iter::once(encoder.finish()));
    }
}