use bytemuck::{Pod, Zeroable};
use glam::{Mat4, Quat, Vec3};
use std::path::Path;
use std::rc::Rc;

use crate::utility::transform::Transform;

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

impl Vertex
{
    pub fn desc() -> wgpu::VertexBufferLayout<'static>
    {
        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,
                },
                wgpu::VertexAttribute {
                    offset: std::mem::size_of::<[f32; 3]>() as wgpu::BufferAddress,
                    shader_location: 1,
                    format: wgpu::VertexFormat::Float32x3,
                },
                wgpu::VertexAttribute {
                    offset: (std::mem::size_of::<[f32; 3]>() * 2) as wgpu::BufferAddress,
                    shader_location: 2,
                    format: wgpu::VertexFormat::Float32x2,
                },
            ],
        }
    }
}

#[derive(Debug, Clone)]
pub struct InstanceData
{
    pub position: Vec3,
    pub rotation: Quat,
    pub scale: Vec3,
    pub dissolve_amount: f32,
}

impl InstanceData
{
    pub fn to_raw(&self) -> InstanceRaw
    {
        let model = Mat4::from_scale_rotation_translation(self.scale, self.rotation, self.position);
        InstanceRaw {
            model: model.to_cols_array_2d(),
            dissolve_amount: self.dissolve_amount,
            _padding: [0.0; 3],
        }
    }
}

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
pub struct InstanceRaw
{
    pub model: [[f32; 4]; 4],
    pub dissolve_amount: f32,
    pub _padding: [f32; 3],
}

impl InstanceRaw
{
    pub fn desc() -> wgpu::VertexBufferLayout<'static>
    {
        wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<InstanceRaw>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Instance,
            attributes: &[
                wgpu::VertexAttribute {
                    offset: 0,
                    shader_location: 3,
                    format: wgpu::VertexFormat::Float32x4,
                },
                wgpu::VertexAttribute {
                    offset: std::mem::size_of::<[f32; 4]>() as wgpu::BufferAddress,
                    shader_location: 4,
                    format: wgpu::VertexFormat::Float32x4,
                },
                wgpu::VertexAttribute {
                    offset: (std::mem::size_of::<[f32; 4]>() * 2) as wgpu::BufferAddress,
                    shader_location: 5,
                    format: wgpu::VertexFormat::Float32x4,
                },
                wgpu::VertexAttribute {
                    offset: (std::mem::size_of::<[f32; 4]>() * 3) as wgpu::BufferAddress,
                    shader_location: 6,
                    format: wgpu::VertexFormat::Float32x4,
                },
                wgpu::VertexAttribute {
                    offset: (std::mem::size_of::<[f32; 4]>() * 4) as wgpu::BufferAddress,
                    shader_location: 7,
                    format: wgpu::VertexFormat::Float32,
                },
            ],
        }
    }
}

pub struct Mesh
{
    pub vertex_buffer: wgpu::Buffer,
    pub index_buffer: wgpu::Buffer,
    pub num_indices: u32,
}

impl Mesh
{
    pub fn new(device: &wgpu::Device, vertices: &[Vertex], indices: &[u32]) -> Self
    {
        use wgpu::util::DeviceExt;

        let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Vertex Buffer"),
            contents: bytemuck::cast_slice(vertices),
            usage: wgpu::BufferUsages::VERTEX,
        });

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

        Self {
            vertex_buffer,
            index_buffer,
            num_indices: indices.len() as u32,
        }
    }

    pub fn create_cube_mesh(device: &wgpu::Device) -> Mesh
    {
        let vertices = vec![
            Vertex {
                position: [-0.5, -0.5, 0.5],
                normal: [0.0, 0.0, 1.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [0.5, -0.5, 0.5],
                normal: [0.0, 0.0, 1.0],
                uv: [1.0, 0.0],
            },
            Vertex {
                position: [0.5, 0.5, 0.5],
                normal: [0.0, 0.0, 1.0],
                uv: [1.0, 1.0],
            },
            Vertex {
                position: [-0.5, 0.5, 0.5],
                normal: [0.0, 0.0, 1.0],
                uv: [0.0, 1.0],
            },
            Vertex {
                position: [0.5, -0.5, 0.5],
                normal: [1.0, 0.0, 0.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [0.5, -0.5, -0.5],
                normal: [1.0, 0.0, 0.0],
                uv: [1.0, 0.0],
            },
            Vertex {
                position: [0.5, 0.5, -0.5],
                normal: [1.0, 0.0, 0.0],
                uv: [1.0, 1.0],
            },
            Vertex {
                position: [0.5, 0.5, 0.5],
                normal: [1.0, 0.0, 0.0],
                uv: [0.0, 1.0],
            },
            Vertex {
                position: [0.5, -0.5, -0.5],
                normal: [0.0, 0.0, -1.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [-0.5, -0.5, -0.5],
                normal: [0.0, 0.0, -1.0],
                uv: [1.0, 0.0],
            },
            Vertex {
                position: [-0.5, 0.5, -0.5],
                normal: [0.0, 0.0, -1.0],
                uv: [1.0, 1.0],
            },
            Vertex {
                position: [0.5, 0.5, -0.5],
                normal: [0.0, 0.0, -1.0],
                uv: [0.0, 1.0],
            },
            Vertex {
                position: [-0.5, -0.5, -0.5],
                normal: [-1.0, 0.0, 0.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [-0.5, -0.5, 0.5],
                normal: [-1.0, 0.0, 0.0],
                uv: [1.0, 0.0],
            },
            Vertex {
                position: [-0.5, 0.5, 0.5],
                normal: [-1.0, 0.0, 0.0],
                uv: [1.0, 1.0],
            },
            Vertex {
                position: [-0.5, 0.5, -0.5],
                normal: [-1.0, 0.0, 0.0],
                uv: [0.0, 1.0],
            },
            Vertex {
                position: [-0.5, 0.5, 0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [0.5, 0.5, 0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [1.0, 0.0],
            },
            Vertex {
                position: [0.5, 0.5, -0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [1.0, 1.0],
            },
            Vertex {
                position: [-0.5, 0.5, -0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [0.0, 1.0],
            },
            Vertex {
                position: [-0.5, -0.5, -0.5],
                normal: [0.0, -1.0, 0.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [0.5, -0.5, -0.5],
                normal: [0.0, -1.0, 0.0],
                uv: [1.0, 0.0],
            },
            Vertex {
                position: [0.5, -0.5, 0.5],
                normal: [0.0, -1.0, 0.0],
                uv: [1.0, 1.0],
            },
            Vertex {
                position: [-0.5, -0.5, 0.5],
                normal: [0.0, -1.0, 0.0],
                uv: [0.0, 1.0],
            },
        ];

        let indices = vec![
            0, 1, 2, 2, 3, 0, 4, 5, 6, 6, 7, 4, 8, 9, 10, 10, 11, 8, 12, 13, 14, 14, 15, 12, 16,
            17, 18, 18, 19, 16, 20, 21, 22, 22, 23, 20,
        ];

        Mesh::new(device, &vertices, &indices)
    }

    pub fn create_plane_mesh(
        device: &wgpu::Device,
        width: f32,
        height: f32,
        subdivisions_x: u32,
        subdivisions_y: u32,
    ) -> Mesh
    {
        let mut vertices = Vec::new();
        let mut indices = Vec::new();

        for y in 0..=subdivisions_y
        {
            for x in 0..=subdivisions_x
            {
                let fx = x as f32 / subdivisions_x as f32;
                let fy = y as f32 / subdivisions_y as f32;

                let px = (fx - 0.5) * width;
                let py = 0.0;
                let pz = (fy - 0.5) * height;

                vertices.push(Vertex {
                    position: [px, py, pz],
                    normal: [0.0, 1.0, 0.0],
                    uv: [fx, fy],
                });
            }
        }

        for y in 0..subdivisions_y
        {
            for x in 0..subdivisions_x
            {
                let row_stride = subdivisions_x + 1;
                let i0 = y * row_stride + x;
                let i1 = i0 + 1;
                let i2 = i0 + row_stride;
                let i3 = i2 + 1;

                indices.push(i0);
                indices.push(i2);
                indices.push(i1);

                indices.push(i1);
                indices.push(i2);
                indices.push(i3);
            }
        }

        Mesh::new(device, &vertices, &indices)
    }

    pub fn create_wireframe_box(device: &wgpu::Device) -> Mesh
    {
        let vertices = vec![
            Vertex {
                position: [-0.5, -0.5, -0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [0.5, -0.5, -0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [0.5, -0.5, 0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [-0.5, -0.5, 0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [-0.5, 0.5, -0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [0.5, 0.5, -0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [0.5, 0.5, 0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [0.0, 0.0],
            },
            Vertex {
                position: [-0.5, 0.5, 0.5],
                normal: [0.0, 1.0, 0.0],
                uv: [0.0, 0.0],
            },
        ];

        let indices = vec![
            0, 1, 1, 2, 2, 3, 3, 0, 4, 5, 5, 6, 6, 7, 7, 4, 0, 4, 1, 5, 2, 6, 3, 7,
        ];

        Mesh::new(device, &vertices, &indices)
    }

    pub fn load_gltf_mesh(
        device: &wgpu::Device,
        path: impl AsRef<Path>,
    ) -> Result<Mesh, Box<dyn std::error::Error>>
    {
        let (gltf, buffers, _images) = gltf::import(path)?;

        let mut all_vertices = Vec::new();
        let mut all_indices = Vec::new();

        for scene in gltf.scenes()
        {
            for node in scene.nodes()
            {
                Self::process_node(
                    &node,
                    &buffers,
                    Mat4::IDENTITY,
                    &mut all_vertices,
                    &mut all_indices,
                )?;
            }
        }

        Ok(Mesh::new(device, &all_vertices, &all_indices))
    }

    fn process_node(
        node: &gltf::Node,
        buffers: &[gltf::buffer::Data],
        parent_transform: Mat4,
        all_vertices: &mut Vec<Vertex>,
        all_indices: &mut Vec<u32>,
    ) -> Result<(), Box<dyn std::error::Error>>
    {
        let local_transform = Mat4::from_cols_array_2d(&node.transform().matrix());
        let global_transform = parent_transform * local_transform;

        if let Some(mesh) = node.mesh()
        {
            for primitive in mesh.primitives()
            {
                let reader =
                    primitive.reader(|buffer| buffers.get(buffer.index()).map(|data| &data[..]));

                let positions = reader
                    .read_positions()
                    .ok_or("Missing position data")?
                    .collect::<Vec<[f32; 3]>>();

                let normals = reader
                    .read_normals()
                    .ok_or("Missing normal data")?
                    .collect::<Vec<[f32; 3]>>();

                let uvs = reader
                    .read_tex_coords(0)
                    .map(|iter| iter.into_f32().collect::<Vec<[f32; 2]>>())
                    .unwrap_or_else(|| vec![[0.0, 0.0]; positions.len()]);

                let base_index = all_vertices.len() as u32;

                let normal_matrix = global_transform.inverse().transpose();

                for ((pos, normal), uv) in positions.iter().zip(normals.iter()).zip(uvs.iter())
                {
                    let pos_vec3 = Vec3::from(*pos);
                    let normal_vec3 = Vec3::from(*normal);

                    let transformed_pos = global_transform.transform_point3(pos_vec3);
                    let transformed_normal =
                        normal_matrix.transform_vector3(normal_vec3).normalize();

                    all_vertices.push(Vertex {
                        position: transformed_pos.into(),
                        normal: transformed_normal.into(),
                        uv: *uv,
                    });
                }

                if let Some(indices_reader) = reader.read_indices()
                {
                    all_indices.extend(indices_reader.into_u32().map(|i| i + base_index));
                }
            }
        }

        for child in node.children()
        {
            Self::process_node(&child, buffers, global_transform, all_vertices, all_indices)?;
        }

        Ok(())
    }

    pub fn load_mesh(path: impl AsRef<Path>) -> Result<Mesh, Box<dyn std::error::Error>>
    {
        crate::render::with_device(|device| Mesh::load_gltf_mesh(device, path))
    }

    pub fn load_gltf_with_instances(
        device: &wgpu::Device,
        path: impl AsRef<Path>,
    ) -> anyhow::Result<Vec<(Mesh, Vec<InstanceData>)>>
    {
        let path = path.as_ref();
        let gltf_str = std::fs::read_to_string(path)?;
        let gltf_json: serde_json::Value = serde_json::from_str(&gltf_str)?;

        let (document, buffers, _images) = gltf::import(path)?;

        let mut result = Vec::new();

        let nodes = gltf_json["nodes"]
            .as_array()
            .ok_or_else(|| anyhow::anyhow!("Missing nodes array"))?;

        for (node_index, json_node) in nodes.iter().enumerate()
        {
            let node = document
                .nodes()
                .nth(node_index)
                .ok_or_else(|| anyhow::anyhow!("Node index mismatch"))?;

            if let Some(mesh_data) = node.mesh()
            {
                let has_instancing = json_node
                    .get("extensions")
                    .and_then(|ext| ext.get("EXT_mesh_gpu_instancing"))
                    .is_some();

                if has_instancing
                {
                    let extensions = json_node.get("extensions").unwrap();
                    let instancing_ext = extensions.get("EXT_mesh_gpu_instancing").unwrap();
                    let mut mesh_vertices = Vec::new();
                    let mut mesh_indices = Vec::new();

                    for primitive in mesh_data.primitives()
                    {
                        let reader = primitive
                            .reader(|buffer| buffers.get(buffer.index()).map(|data| &data[..]));

                        let positions = reader
                            .read_positions()
                            .ok_or_else(|| anyhow::anyhow!("Missing position data"))?
                            .collect::<Vec<[f32; 3]>>();

                        let normals = reader
                            .read_normals()
                            .ok_or_else(|| anyhow::anyhow!("Missing normal data"))?
                            .collect::<Vec<[f32; 3]>>();

                        let uvs = reader
                            .read_tex_coords(0)
                            .map(|iter| iter.into_f32().collect::<Vec<[f32; 2]>>())
                            .unwrap_or_else(|| vec![[0.0, 0.0]; positions.len()]);

                        let base_index = mesh_vertices.len() as u32;

                        for ((pos, normal), uv) in
                            positions.iter().zip(normals.iter()).zip(uvs.iter())
                        {
                            mesh_vertices.push(Vertex {
                                position: *pos,
                                normal: *normal,
                                uv: *uv,
                            });
                        }

                        if let Some(indices_reader) = reader.read_indices()
                        {
                            mesh_indices.extend(indices_reader.into_u32().map(|i| i + base_index));
                        }
                    }

                    let attributes = instancing_ext
                        .get("attributes")
                        .and_then(|v| v.as_object())
                        .ok_or_else(|| {
                            anyhow::anyhow!("Missing attributes in EXT_mesh_gpu_instancing")
                        })?;

                    let translation_accessor_index = attributes
                        .get("TRANSLATION")
                        .and_then(|v| v.as_u64())
                        .ok_or_else(|| {
                            anyhow::anyhow!("Missing TRANSLATION in instancing extension")
                        })? as usize;

                    let rotation_accessor_index = attributes
                        .get("ROTATION")
                        .and_then(|v| v.as_u64())
                        .ok_or_else(|| {
                            anyhow::anyhow!("Missing ROTATION in instancing extension")
                        })? as usize;

                    let scale_accessor_index = attributes
                        .get("SCALE")
                        .and_then(|v| v.as_u64())
                        .ok_or_else(|| anyhow::anyhow!("Missing SCALE in instancing extension"))?
                        as usize;

                    let translations =
                        Self::read_vec3_accessor(&document, &buffers, translation_accessor_index)?;
                    let rotations =
                        Self::read_quat_accessor(&document, &buffers, rotation_accessor_index)?;
                    let scales =
                        Self::read_vec3_accessor(&document, &buffers, scale_accessor_index)?;

                    let instances: Vec<InstanceData> = translations
                        .into_iter()
                        .zip(rotations.into_iter())
                        .zip(scales.into_iter())
                        .map(|((position, rotation), scale)| InstanceData {
                            position,
                            rotation,
                            scale,
                            dissolve_amount: 0.0,
                        })
                        .collect();

                    let mesh = Mesh::new(device, &mesh_vertices, &mesh_indices);
                    result.push((mesh, instances));
                }
                else
                {
                    let mut mesh_vertices = Vec::new();
                    let mut mesh_indices = Vec::new();

                    for primitive in mesh_data.primitives()
                    {
                        let reader = primitive
                            .reader(|buffer| buffers.get(buffer.index()).map(|data| &data[..]));

                        let positions = reader
                            .read_positions()
                            .ok_or_else(|| anyhow::anyhow!("Missing position data"))?
                            .collect::<Vec<[f32; 3]>>();

                        let normals = reader
                            .read_normals()
                            .ok_or_else(|| anyhow::anyhow!("Missing normal data"))?
                            .collect::<Vec<[f32; 3]>>();

                        let uvs = reader
                            .read_tex_coords(0)
                            .map(|iter| iter.into_f32().collect::<Vec<[f32; 2]>>())
                            .unwrap_or_else(|| vec![[0.0, 0.0]; positions.len()]);

                        let base_index = mesh_vertices.len() as u32;

                        for ((pos, normal), uv) in
                            positions.iter().zip(normals.iter()).zip(uvs.iter())
                        {
                            mesh_vertices.push(Vertex {
                                position: *pos,
                                normal: *normal,
                                uv: *uv,
                            });
                        }

                        if let Some(indices_reader) = reader.read_indices()
                        {
                            mesh_indices.extend(indices_reader.into_u32().map(|i| i + base_index));
                        }
                    }

                    let mesh = Mesh::new(device, &mesh_vertices, &mesh_indices);
                    result.push((mesh, Vec::new()));
                }
            }
        }

        Ok(result)
    }

    fn read_vec3_accessor(
        document: &gltf::Document,
        buffers: &[gltf::buffer::Data],
        accessor_index: usize,
    ) -> anyhow::Result<Vec<Vec3>>
    {
        let accessor = document
            .accessors()
            .nth(accessor_index)
            .ok_or_else(|| anyhow::anyhow!("Invalid accessor index"))?;

        let buffer_view = accessor
            .view()
            .ok_or_else(|| anyhow::anyhow!("Missing buffer view"))?;
        let buffer = &buffers[buffer_view.buffer().index()];
        let start = buffer_view.offset() + accessor.offset();
        let stride = buffer_view.stride().unwrap_or(12);

        let mut result = Vec::new();
        for i in 0..accessor.count()
        {
            let offset = start + i * stride;
            let x = f32::from_le_bytes([
                buffer[offset],
                buffer[offset + 1],
                buffer[offset + 2],
                buffer[offset + 3],
            ]);
            let y = f32::from_le_bytes([
                buffer[offset + 4],
                buffer[offset + 5],
                buffer[offset + 6],
                buffer[offset + 7],
            ]);
            let z = f32::from_le_bytes([
                buffer[offset + 8],
                buffer[offset + 9],
                buffer[offset + 10],
                buffer[offset + 11],
            ]);
            result.push(Vec3::new(x, y, z));
        }

        Ok(result)
    }

    fn read_quat_accessor(
        document: &gltf::Document,
        buffers: &[gltf::buffer::Data],
        accessor_index: usize,
    ) -> anyhow::Result<Vec<Quat>>
    {
        let accessor = document
            .accessors()
            .nth(accessor_index)
            .ok_or_else(|| anyhow::anyhow!("Invalid accessor index"))?;

        let buffer_view = accessor
            .view()
            .ok_or_else(|| anyhow::anyhow!("Missing buffer view"))?;
        let buffer = &buffers[buffer_view.buffer().index()];
        let start = buffer_view.offset() + accessor.offset();
        let stride = buffer_view.stride().unwrap_or(16);

        let mut result = Vec::new();
        for i in 0..accessor.count()
        {
            let offset = start + i * stride;
            let x = f32::from_le_bytes([
                buffer[offset],
                buffer[offset + 1],
                buffer[offset + 2],
                buffer[offset + 3],
            ]);
            let y = f32::from_le_bytes([
                buffer[offset + 4],
                buffer[offset + 5],
                buffer[offset + 6],
                buffer[offset + 7],
            ]);
            let z = f32::from_le_bytes([
                buffer[offset + 8],
                buffer[offset + 9],
                buffer[offset + 10],
                buffer[offset + 11],
            ]);
            let w = f32::from_le_bytes([
                buffer[offset + 12],
                buffer[offset + 13],
                buffer[offset + 14],
                buffer[offset + 15],
            ]);
            result.push(Quat::from_xyzw(x, y, z, w));
        }

        Ok(result)
    }
}