use std::collections::HashMap;
use wgpu::util::DeviceExt;

use winit::{
    event::*,
    event_loop::{ControlFlow, EventLoop},
    window::WindowBuilder,
};

mod camera;

use camera::*;

struct Renderer {
    pub device: wgpu::Device,
    pub mesh_pool: MeshPool,
    pub size: winit::dpi::PhysicalSize<u32>,
    surface: wgpu::Surface,
    queue: wgpu::Queue,
    config: wgpu::SurfaceConfiguration,
    camera_uniform: CameraUniform,
    camera_buffer: wgpu::Buffer,
    camera_bind_group: wgpu::BindGroup,
    meshes_buffer: wgpu::Buffer,
    meshes_bind_group: wgpu::BindGroup,
    render_pipeline: wgpu::RenderPipeline,
}

impl Renderer {
    pub async fn new(window: &winit::window::Window) -> Self {
        let size = window.inner_size();

        let instance = wgpu::Instance::new(wgpu::Backends::all());
        let surface = unsafe { instance.create_surface(window) };
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                compatible_surface: Some(&surface),
                force_fallback_adapter: false,
            })
            .await
            .unwrap();

        let (device, queue) = adapter
            .request_device(
                &wgpu::DeviceDescriptor {
                    features: wgpu::Features::empty(),
                    limits: wgpu::Limits::default(),
                    label: None,
                },
                None,
            )
            .await
            .unwrap();

        let config = wgpu::SurfaceConfiguration {
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            format: surface.get_preferred_format(&adapter).unwrap(),
            width: size.width,
            height: size.height,
            present_mode: wgpu::PresentMode::Fifo,
        };
        surface.configure(&device, &config);

        let mesh_pool = MeshPool::default();

        let camera_uniform = CameraUniform::new();

        let camera_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Camera Buffer"),
            contents: bytemuck::cast_slice(&[camera_uniform]),
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        });

        let camera_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                entries: &[wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                }],
                label: Some("Camera Bind Group Layout"),
            });

        let camera_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            layout: &camera_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: camera_buffer.as_entire_binding(),
            }],
            label: Some("Camera Bind Group"),
        });

        let meshes_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Meshes Buffer"),
            size: 65536, // TODO resizable meshes buffer/gpu vectors
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let meshes_bind_group_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                entries: &[wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::VERTEX,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                }],
                label: Some("Meshes Bind Group Layout"),
            });

        let meshes_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            layout: &meshes_bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: meshes_buffer.as_entire_binding(),
            }],
            label: Some("Meshes Bind Group"),
        });

        let render_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("Render Pipeline Layout"),
                bind_group_layouts: &[&camera_bind_group_layout, &meshes_bind_group_layout],
                push_constant_ranges: &[],
            });

        let shader = device.create_shader_module(&wgpu::include_wgsl!("shader.wgsl"));

        let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Render Pipeline"),
            layout: Some(&render_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: "vs_main",
                buffers: &[Vertex::desc()],
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: "fs_main",
                targets: &[wgpu::ColorTargetState {
                    format: config.format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                }],
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: None,
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState {
                count: 1,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            multiview: None,
        });

        Self {
            size,
            surface,
            device,
            queue,
            config,
            mesh_pool,
            camera_uniform,
            camera_buffer,
            camera_bind_group,
            meshes_buffer,
            meshes_bind_group,
            render_pipeline,
        }
    }

    pub fn resize(&mut self, new_size: winit::dpi::PhysicalSize<u32>) {
        if new_size.width > 0 && new_size.height > 0 {
            self.size = new_size;
            self.config.width = new_size.width;
            self.config.height = new_size.height;
            self.surface.configure(&self.device, &self.config);
        }
    }

    pub fn render(
        &mut self,
        camera: &impl Camera,
        meshes: &MeshCommands,
    ) -> Result<(), wgpu::SurfaceError> {
        self.camera_uniform.update(camera);
        self.queue.write_buffer(
            &self.camera_buffer,
            0,
            bytemuck::cast_slice(&[self.camera_uniform]),
        );

        let mut sorted_meshes = HashMap::<usize, Vec<MeshInstance>>::new();
        for mesh in meshes.iter() {
            let group_id = mesh.handle.group_id;
            if let Some(by_group) = sorted_meshes.get_mut(&group_id) {
                by_group.push(*mesh);
            } else {
                let new_list = vec![*mesh];
                sorted_meshes.insert(group_id, new_list);
            }
        }

        let mut mesh_transforms = Vec::<[f32; 16]>::new();
        let mut transform_ranges = Vec::<(usize, std::ops::Range<u32>)>::new();

        // this code assumes MeshHandle only uses group_id (which it does for now)
        // TODO bucket by sub_id too before MeshHandle supports it
        for (group_id, instances) in sorted_meshes.iter() {
            let start_idx = mesh_transforms.len() as u32;
            let transforms = instances.iter().map(|i| i.transform.to_cols_array());
            mesh_transforms.extend(transforms);
            let end_idx = mesh_transforms.len() as u32;
            transform_ranges.push((*group_id, start_idx..end_idx));
        }

        self.queue.write_buffer(
            &self.meshes_buffer,
            0,
            bytemuck::cast_slice(&mesh_transforms),
        );

        let output = self.surface.get_current_texture()?;
        let view = output
            .texture
            .create_view(&wgpu::TextureViewDescriptor::default());
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("Render Encoder"),
            });

        {
            let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("Render Pass"),
                color_attachments: &[wgpu::RenderPassColorAttachment {
                    view: &view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color {
                            r: 0.1,
                            g: 0.2,
                            b: 0.3,
                            a: 1.0,
                        }),
                        store: true,
                    },
                }],
                depth_stencil_attachment: None,
            });

            render_pass.set_pipeline(&self.render_pipeline);
            render_pass.set_bind_group(0, &self.camera_bind_group, &[]);
            render_pass.set_bind_group(1, &self.meshes_bind_group, &[]);

            // TODO one group per mesh, still...
            // TODO this could be implemented without accessing private members
            for (group_id, meshes_range) in transform_ranges.iter() {
                let group = self.mesh_pool.groups.get(*group_id).unwrap();
                render_pass.set_vertex_buffer(0, group.vertices.slice(..));
                render_pass.set_index_buffer(group.indices.slice(..), wgpu::IndexFormat::Uint32);
                let indices = 0..(group.index_capacity as u32);
                render_pass.draw_indexed(indices, 0, meshes_range.to_owned());
            }
        }

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

        Ok(())
    }
}

struct MeshGroup {
    vertices: wgpu::Buffer,
    vertex_capacity: usize,
    indices: wgpu::Buffer,
    index_capacity: usize,
}

impl MeshGroup {
    pub fn new(device: &wgpu::Device, data: &MeshData) -> Self {
        let vertex_capacity = data.vertices.len();
        let vertices = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Vertex Buffer"),
            contents: bytemuck::cast_slice(&data.vertices),
            usage: wgpu::BufferUsages::VERTEX,
        });

        let index_capacity = data.indices.len();
        let indices = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Index Buffer"),
            contents: bytemuck::cast_slice(&data.indices),
            usage: wgpu::BufferUsages::INDEX,
        });

        Self {
            vertex_capacity,
            vertices,
            index_capacity,
            indices,
        }
    }
}

#[derive(Default)]
struct MeshPool {
    groups: slab::Slab<MeshGroup>,
}

impl MeshPool {
    pub fn allocate(&mut self, device: &wgpu::Device, data: &MeshData) -> MeshHandle {
        let group = MeshGroup::new(device, data);
        let group_id = self.groups.insert(group);
        let sub_id = 0;
        MeshHandle { group_id, sub_id }
    }
}

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

impl CameraUniform {
    pub fn new() -> Self {
        Self {
            vp: glam::Mat4::IDENTITY.to_cols_array_2d(),
        }
    }

    pub fn update(&mut self, camera: &impl Camera) {
        self.vp = camera.get_vp();
    }
}

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct Vertex {
    position: [f32; 3],
}

impl Vertex {
    pub fn desc<'a>() -> wgpu::VertexBufferLayout<'a> {
        wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &[wgpu::VertexAttribute {
                offset: 0,
                shader_location: 0,
                format: wgpu::VertexFormat::Float32x3,
            }],
        }
    }
}

type Index = u32;

struct MeshData {
    vertices: Vec<Vertex>,
    indices: Vec<Index>,
}

#[repr(C)]
#[derive(Copy, Clone, Eq, Hash, PartialEq)]
struct MeshHandle {
    group_id: usize,
    // unused for now, since each group contains only one mesh
    sub_id: usize,
}

#[derive(Copy, Clone, PartialEq)]
struct MeshInstance {
    pub handle: MeshHandle,
    pub transform: glam::Mat4,
}

type MeshCommands = Vec<MeshInstance>;

fn load_model() -> MeshData {
    use tobj::*;

    let mut model_data = include_bytes!("viking_room.obj").to_vec();
    let mut model_data = &mut model_data.as_slice();

    let load_options = LoadOptions {
        triangulate: true,
        single_index: true,
        ..Default::default()
    };

    let (models, _mats) = load_obj_buf(model_data, &load_options, |_| unimplemented!()).unwrap();

    let mut vertices = Vec::new();
    let mut indices = Vec::new();
    for m in models {
        let index_base = vertices.len() as u32;
        for i in 0..m.mesh.positions.len() / 3 {
            let t = i * 3;
            vertices.push(Vertex {
                position: [
                    m.mesh.positions[t],
                    m.mesh.positions[t + 2],
                    -m.mesh.positions[t + 1],
                ],
            });
        }

        indices.extend(m.mesh.indices.iter().map(|i| i + index_base));
    }

    MeshData { vertices, indices }
}

fn main() {
    let mesh_data = load_model();
    let event_loop = EventLoop::new();
    let window = WindowBuilder::new().build(&event_loop).unwrap();
    let mut camera = Flycam::new(10.0, 0.002);
    let mut is_grabbed = false;
    let mut ren = pollster::block_on(Renderer::new(&window));

    let mesh = ren.mesh_pool.allocate(&ren.device, &mesh_data);
    let mut commands = Vec::new();
    for x in -5..5 {
        for y in -5..5 {
            let translation = glam::Vec3::new(x as f32, 0.0, y as f32) * 3.0;
            let transform = glam::Mat4::from_translation(translation);
            commands.push(MeshInstance {
                handle: mesh,
                transform,
            });
        }
    }

    event_loop.run(move |event, _, control_flow| match event {
        Event::RedrawRequested(_) => match ren.render(&camera, &commands) {
            Ok(_) => {}
            Err(wgpu::SurfaceError::Lost) => ren.resize(ren.size),
            Err(wgpu::SurfaceError::OutOfMemory) => *control_flow = ControlFlow::Exit,
            Err(e) => println!("error: {:?}", e),
        },
        Event::MainEventsCleared => {
            camera.update();
            window.request_redraw();
        }
        Event::DeviceEvent { ref event, .. } => match event {
            DeviceEvent::MouseMotion { delta } => {
                if is_grabbed {
                    camera.process_mouse(delta.0, delta.1);
                }
            }
            _ => {}
        },
        Event::WindowEvent {
            ref event,
            window_id,
        } if window_id == window.id() => match event {
            WindowEvent::KeyboardInput {
                input:
                    KeyboardInput {
                        virtual_keycode: Some(key),
                        state,
                        ..
                    },
                ..
            } => {
                if *state == ElementState::Pressed && *key == VirtualKeyCode::Escape {
                    if is_grabbed {
                        window.set_cursor_grab(false).unwrap();
                        window.set_cursor_visible(true);
                        is_grabbed = false;
                    }
                } else {
                    camera.process_keyboard(*key, *state);
                }
            }
            WindowEvent::MouseInput {
                button: MouseButton::Left,
                state: ElementState::Pressed,
                ..
            } => {
                if !is_grabbed {
                    window.set_cursor_grab(true).unwrap();
                    window.set_cursor_visible(false);
                    is_grabbed = true;
                }
            }
            WindowEvent::CloseRequested => *control_flow = ControlFlow::Exit,
            WindowEvent::Resized(physical_size) => {
                ren.resize(*physical_size);
                camera.resize(physical_size.width, physical_size.height);
            }
            WindowEvent::ScaleFactorChanged { new_inner_size, .. } => {
                ren.resize(**new_inner_size);
                camera.resize(new_inner_size.width, new_inner_size.height);
            }
            _ => {}
        },
        _ => {}
    });
}