use std::collections::HashMap;
use std::path::Path;

pub trait GlbProperties: Default + Clone {
    fn from_extras(extras: &serde_json::Value) -> Self;
}

impl GlbProperties for () {
    fn from_extras(_extras: &serde_json::Value) -> Self {}
}

/// Reads a boolean glTF extra by key, defaulting to `false` when absent or non-boolean.
pub fn extra_bool(extras: &serde_json::Value, key: &str) -> bool {
    extras
        .get(key)
        .and_then(serde_json::Value::as_bool)
        .unwrap_or(false)
}

/// Reads a string glTF extra by key, returning `None` when absent or non-string.
pub fn extra_str(extras: &serde_json::Value, key: &str) -> Option<String> {
    extras
        .get(key)
        .and_then(serde_json::Value::as_str)
        .map(str::to_string)
}

#[derive(Clone)]
pub struct LoadedMesh<P: GlbProperties = ()> {
    pub vertices: Vec<MeshVertex>,
    pub indices: Vec<u32>,
    pub material: Option<MeshMaterial>,
    pub transform: [[f32; 4]; 4],
    pub properties: P,
    pub texture_index: Option<usize>,
    pub normal_texture_index: Option<usize>,
    pub metallic_roughness_texture_index: Option<usize>,
    pub occlusion_texture_index: Option<usize>,
    pub alpha_mode: AlphaMode,
}

impl<P: GlbProperties> LoadedMesh<P> {
    /// Axis-aligned bounds `[min, max]` over all vertex positions.
    ///
    /// An empty mesh yields `[[+∞; 3], [-∞; 3]]`.
    pub fn bounds(&self) -> [[f32; 3]; 2] {
        self.vertices.iter().fold(
            [[f32::INFINITY; 3], [f32::NEG_INFINITY; 3]],
            |[min, max], vertex| {
                let [x, y, z] = vertex.position;
                [
                    [min[0].min(x), min[1].min(y), min[2].min(z)],
                    [max[0].max(x), max[1].max(y), max[2].max(z)],
                ]
            },
        )
    }
}

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

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum AlphaMode {
    Opaque,
    Mask(f32),
    Blend,
}

#[derive(Clone)]
pub struct MeshMaterial {
    pub name: String,
    pub texture_filename: Option<String>,
    pub base_color: [f32; 4],
    pub alpha_mode: AlphaMode,
}

#[derive(Clone)]
pub struct DecodedTexture {
    pub pixels: Vec<u8>,
    pub width: u32,
    pub height: u32,
}

#[derive(Clone)]
pub struct LevelNode<P: GlbProperties = ()> {
    pub transform: [[f32; 4]; 4],
    pub properties: P,
}

#[derive(Clone)]
pub struct LoadedLevel<P: GlbProperties = ()> {
    pub meshes: Vec<LoadedMesh<P>>,
    pub textures: Vec<DecodedTexture>,
    pub spawn: Option<[f32; 3]>,
    pub nodes: Vec<LevelNode<P>>,
}

impl<P: GlbProperties> LoadedLevel<P> {
    /// Axis-aligned bounds `[min, max]` over all meshes.
    ///
    /// A level with no meshes yields `[[+∞; 3], [-∞; 3]]`.
    pub fn bounds(&self) -> [[f32; 3]; 2] {
        self.meshes.iter().fold(
            [[f32::INFINITY; 3], [f32::NEG_INFINITY; 3]],
            |[min, max], mesh| {
                let [mesh_min, mesh_max] = mesh.bounds();
                [
                    [
                        min[0].min(mesh_min[0]),
                        min[1].min(mesh_min[1]),
                        min[2].min(mesh_min[2]),
                    ],
                    [
                        max[0].max(mesh_max[0]),
                        max[1].max(mesh_max[1]),
                        max[2].max(mesh_max[2]),
                    ],
                ]
            },
        )
    }

    pub fn load(path: &Path) -> Result<Self, String> {
        let cache = HashMap::new();
        let mut loading = std::collections::HashSet::new();
        Self::load_impl(path, &cache, &mut loading)
    }

    fn load_impl(
        path: &Path,
        cache: &HashMap<std::path::PathBuf, Self>,
        loading: &mut std::collections::HashSet<std::path::PathBuf>,
    ) -> Result<Self, String> {
        let canonical = path
            .canonicalize()
            .map_err(|error| format!("failed to resolve {}: {error}", path.display()))?;

        if let Some(cached) = cache.get(&canonical) {
            return Ok(cached.clone());
        }

        if !loading.insert(canonical.clone()) {
            return Err(format!("prefab cycle detected: {}", canonical.display()));
        }

        let gltf::Glb { json, bin, .. } = gltf::Glb::from_reader(
            std::fs::File::open(path)
                .map_err(|error| format!("failed to open {}: {error}", path.display()))?,
        )
        .map_err(|error| format!("failed to parse {}: {error}", path.display()))?;

        Self::from_glb(&json, bin, path.parent().unwrap_or_else(|| Path::new(".")))
    }

    /// Loads a level directly from in-memory `.glb` bytes.
    ///
    /// Unlike [`load`](Self::load) this resolves no external files: textures
    /// must be embedded and prefab references are not expanded. Suitable for
    /// embedded or fetched assets (e.g. on `wasm32` where there is no filesystem).
    pub fn load_from_slice(bytes: &[u8]) -> Result<Self, String> {
        let gltf::Glb { json, bin, .. } = gltf::Glb::from_slice(bytes)
            .map_err(|error| format!("failed to parse glb: {error}"))?;
        Self::from_glb(&json, bin, Path::new("."))
    }

    fn from_glb(
        json: &[u8],
        bin: Option<std::borrow::Cow<'_, [u8]>>,
        base_dir: &Path,
    ) -> Result<Self, String> {
        let mut json_value: serde_json::Value = serde_json::from_slice(json)
            .map_err(|error| format!("failed to parse glTF JSON: {error}"))?;
        if let Some(object) = json_value.as_object_mut() {
            object.remove("extensionsRequired");
        }
        let json_bytes = serde_json::to_vec(&json_value)
            .map_err(|error| format!("failed to re-encode JSON: {error}"))?;

        let root: gltf::json::Root = gltf::json::deserialize::from_slice(&json_bytes)
            .map_err(|error| format!("failed to parse glTF root: {error}"))?;
        let document = gltf::Document::from_json_without_validation(root);
        let buffers = bin
            .map(|data| vec![gltf::buffer::Data(data.into_owned())])
            .unwrap_or_default();

        let textures = extract_textures(&document, &buffers, base_dir);
        let world_transforms = compute_world_transforms(&document);

        let mut meshes = Vec::new();
        let mut all_nodes = Vec::new();
        let mut spawn = None;

        for node in document.nodes() {
            let transform = world_transforms[node.index()];
            let properties = Self::parse_properties(&node);
            all_nodes.push(LevelNode {
                transform,
                properties: properties.clone(),
            });

            if let Some(extras) = node.extras().as_ref()
                && let Ok(json_obj) = serde_json::from_str::<serde_json::Value>(extras.get())
                && json_obj
                    .as_object()
                    .and_then(|m| m.get("spawn"))
                    .and_then(serde_json::Value::as_bool)
                    .unwrap_or(false)
            {
                let world = &world_transforms[node.index()];
                spawn = Some([world[3][0], world[3][1], world[3][2]]);
            }

            let Some(mesh) = node.mesh() else {
                continue;
            };

            let transform = world_transforms[node.index()];

            for primitive in mesh.primitives() {
                if primitive.mode() != gltf::mesh::Mode::Triangles {
                    continue;
                }

                let material = extract_material(primitive.material());
                let texture_index = primitive
                    .material()
                    .pbr_metallic_roughness()
                    .base_color_texture()
                    .map(|info| info.texture().source().index());
                let normal_texture_index = primitive
                    .material()
                    .normal_texture()
                    .map(|info| info.texture().source().index());
                let metallic_roughness_texture_index = primitive
                    .material()
                    .pbr_metallic_roughness()
                    .metallic_roughness_texture()
                    .map(|info| info.texture().source().index());
                let occlusion_texture_index = primitive
                    .material()
                    .occlusion_texture()
                    .map(|info| info.texture().source().index());

                let reader = primitive.reader(|buffer| Some(&buffers[buffer.index()]));

                let positions: Vec<[f32; 3]> = reader
                    .read_positions()
                    .map(Iterator::collect)
                    .unwrap_or_default();

                let normals: Vec<[f32; 3]> = reader
                    .read_normals()
                    .map_or_else(|| vec![[0.0, 1.0, 0.0]; positions.len()], Iterator::collect);

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

                let vertices: Vec<MeshVertex> = (0..positions.len())
                    .map(|index| MeshVertex {
                        position: transform_point(&transform, positions[index]),
                        normal: transform_normal(&transform, normals[index]),
                        uv: uvs[index],
                    })
                    .collect();

                if vertices.is_empty() {
                    continue;
                }

                let indices: Vec<u32> = reader.read_indices().map_or_else(
                    || (0..vertices.len() as u32).collect(),
                    |index_reader| index_reader.into_u32().collect(),
                );

                let mesh_alpha_mode = material
                    .as_ref()
                    .map_or(AlphaMode::Opaque, |m| m.alpha_mode);

                meshes.push(LoadedMesh {
                    vertices,
                    indices,
                    material,
                    transform,
                    properties: properties.clone(),
                    texture_index,
                    normal_texture_index,
                    metallic_roughness_texture_index,
                    occlusion_texture_index,
                    alpha_mode: mesh_alpha_mode,
                });
            }
        }

        Ok(Self {
            meshes,
            textures,
            spawn,
            nodes: all_nodes,
        })
    }

    fn parse_properties(node: &gltf::Node) -> P {
        let Some(extras) = node.extras().as_ref() else {
            return P::default();
        };

        let Ok(json) = serde_json::from_str::<serde_json::Value>(extras.get()) else {
            return P::default();
        };

        P::from_extras(&json)
    }
}

fn extract_textures(
    document: &gltf::Document,
    buffers: &[gltf::buffer::Data],
    base_dir: &Path,
) -> Vec<DecodedTexture> {
    let mut textures = Vec::new();
    for image in document.images() {
        let (data, format) = match image.source() {
            gltf::image::Source::View { view, mime_type } => {
                let buffer = &buffers[view.buffer().index()];
                let offset = view.offset();
                let length = view.length();
                let format = match mime_type {
                    "image/png" => image::ImageFormat::Png,
                    "image/jpeg" => image::ImageFormat::Jpeg,
                    _ => continue,
                };
                (buffer[offset..offset + length].to_vec(), format)
            }
            gltf::image::Source::Uri { uri, mime_type } => {
                let path = base_dir.join(uri);
                let Ok(data) = std::fs::read(&path) else {
                    continue;
                };
                let format = match mime_type {
                    Some("image/png") => image::ImageFormat::Png,
                    Some("image/jpeg") => image::ImageFormat::Jpeg,
                    _ => match path.extension().and_then(|e| e.to_str()) {
                        Some("png") => image::ImageFormat::Png,
                        Some("jpg" | "jpeg") => image::ImageFormat::Jpeg,
                        _ => continue,
                    },
                };
                (data, format)
            }
        };
        let Ok(decoded) = image::load_from_memory_with_format(&data, format) else {
            continue;
        };
        textures.push(DecodedTexture {
            width: decoded.width(),
            height: decoded.height(),
            pixels: decoded.to_rgba8().into_raw(),
        });
    }
    textures
}

fn compute_world_transforms(document: &gltf::Document) -> Vec<[[f32; 4]; 4]> {
    let node_count = document.nodes().count();
    let mut world = vec![IDENTITY_MATRIX; node_count];

    fn visit(node: &gltf::Node, parent_transform: &[[f32; 4]; 4], world: &mut [[[f32; 4]; 4]]) {
        let local = node.transform().matrix();
        let current = multiply_matrices(parent_transform, &local);
        world[node.index()] = current;
        for child in node.children() {
            visit(&child, &current, world);
        }
    }

    for scene in document.scenes() {
        for root in scene.nodes() {
            visit(&root, &IDENTITY_MATRIX, &mut world);
        }
    }

    world
}

const IDENTITY_MATRIX: [[f32; 4]; 4] = [
    [1.0, 0.0, 0.0, 0.0],
    [0.0, 1.0, 0.0, 0.0],
    [0.0, 0.0, 1.0, 0.0],
    [0.0, 0.0, 0.0, 1.0],
];

fn multiply_matrices(a: &[[f32; 4]; 4], b: &[[f32; 4]; 4]) -> [[f32; 4]; 4] {
    let mut result = [[0.0; 4]; 4];
    for column in 0..4 {
        for row in 0..4 {
            for k in 0..4 {
                result[column][row] += a[k][row] * b[column][k];
            }
        }
    }
    result
}

fn transform_point(matrix: &[[f32; 4]; 4], point: [f32; 3]) -> [f32; 3] {
    let [x, y, z] = point;
    [
        matrix[2][0].mul_add(z, matrix[0][0].mul_add(x, matrix[1][0] * y)) + matrix[3][0],
        matrix[2][1].mul_add(z, matrix[0][1].mul_add(x, matrix[1][1] * y)) + matrix[3][1],
        matrix[2][2].mul_add(z, matrix[0][2].mul_add(x, matrix[1][2] * y)) + matrix[3][2],
    ]
}

fn transform_normal(matrix: &[[f32; 4]; 4], normal: [f32; 3]) -> [f32; 3] {
    let [x, y, z] = normal;
    let nx = matrix[2][0].mul_add(z, matrix[0][0].mul_add(x, matrix[1][0] * y));
    let ny = matrix[2][1].mul_add(z, matrix[0][1].mul_add(x, matrix[1][1] * y));
    let nz = matrix[2][2].mul_add(z, matrix[0][2].mul_add(x, matrix[1][2] * y));
    let length = nz.mul_add(nz, nx.mul_add(nx, ny * ny)).sqrt();
    if length > 0.0 {
        [nx / length, ny / length, nz / length]
    } else {
        [0.0, 1.0, 0.0]
    }
}

fn extract_material(material: gltf::Material) -> Option<MeshMaterial> {
    let name = material.name().unwrap_or("").to_string();
    let pbr = material.pbr_metallic_roughness();
    let base_color = pbr.base_color_factor();

    let texture_filename = pbr.base_color_texture().and_then(|info| {
        let image = info.texture().source();
        match image.source() {
            gltf::image::Source::Uri { uri, .. } => Some(uri.to_string()),
            gltf::image::Source::View { .. } => image.name().map(str::to_string),
        }
    });

    let alpha_mode = match material.alpha_mode() {
        gltf::material::AlphaMode::Opaque => AlphaMode::Opaque,
        gltf::material::AlphaMode::Mask => AlphaMode::Mask(material.alpha_cutoff().unwrap_or(0.5)),
        gltf::material::AlphaMode::Blend => AlphaMode::Blend,
    };

    if name.is_empty() && texture_filename.is_none() && base_color == [1.0, 1.0, 1.0, 1.0] {
        return None;
    }

    Some(MeshMaterial {
        name,
        texture_filename,
        base_color,
        alpha_mode,
    })
}

#[cfg(test)]
mod tests {
    use super::{
        AlphaMode, IDENTITY_MATRIX, LoadedMesh, MeshVertex, extra_bool, extra_str,
        multiply_matrices, transform_point,
    };

    fn vertex(position: [f32; 3]) -> MeshVertex {
        MeshVertex {
            position,
            normal: [0.0, 1.0, 0.0],
            uv: [0.0, 0.0],
        }
    }

    #[test]
    fn mesh_bounds_span_positions() {
        let mesh = LoadedMesh::<()> {
            vertices: vec![vertex([-1.0, 2.0, 0.0]), vertex([3.0, -1.0, 5.0])],
            indices: vec![],
            material: None,
            transform: IDENTITY_MATRIX,
            properties: (),
            texture_index: None,
            normal_texture_index: None,
            metallic_roughness_texture_index: None,
            occlusion_texture_index: None,
            alpha_mode: AlphaMode::Opaque,
        };
        assert_eq!(mesh.bounds(), [[-1.0, -1.0, 0.0], [3.0, 2.0, 5.0]]);
    }

    #[test]
    fn extras_read_typed_values() {
        let json = serde_json::json!({ "door": "east", "spawn": true });
        assert!(extra_bool(&json, "spawn"));
        assert!(!extra_bool(&json, "missing"));
        assert_eq!(extra_str(&json, "door").as_deref(), Some("east"));
        assert_eq!(extra_str(&json, "missing"), None);
    }

    #[test]
    fn parent_rotation_applies_to_child_translation() {
        let parent_rotation_z_90 = [
            [0.0, 1.0, 0.0, 0.0],
            [-1.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 1.0, 0.0],
            [0.0, 0.0, 0.0, 1.0],
        ];
        let child_translation = [
            [1.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 0.0, 0.0],
            [0.0, 0.0, 1.0, 0.0],
            [1.0, 0.0, 0.0, 1.0],
        ];
        let world = multiply_matrices(&parent_rotation_z_90, &child_translation);
        let [x, y, z] = transform_point(&world, [0.0, 0.0, 0.0]);
        assert!(x.abs() < 1e-6);
        assert!((y - 1.0).abs() < 1e-6);
        assert!(z.abs() < 1e-6);
    }
}
