/*
    wearechat - generic multiplayer game with voip
    Copyright (C) 2025 metamuffin

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published by
    the Free Software Foundation, version 3 of the License only.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/
use crate::{
    Args, TextureCache, animation::import_animation, mesh::import_mesh, physics::import_physics,
    transform_to_affine, vrm::extract_vrm_data,
};
use anyhow::{Context, Result, anyhow};
use glam::{Affine3A, Vec3, Vec3A, vec3a};
use gltf::{Gltf, Node, import_buffers, scene::Transform};
use log::{debug, info};
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use serde::Deserialize;
use serde_json::Value;
use std::{
    borrow::Cow,
    collections::BTreeMap,
    fs::{File, read_to_string},
    io::Read,
    path::Path,
    sync::Mutex,
};
use weareshared::{
    packets::Resource,
    resources::{
        ArmaturePart, AvatarInfoPart, EnvironmentPart, Image, LightPart, ParticlesPart, Prefab,
    },
    store::ResourceStore,
};

pub fn import_prefab(
    store: &ResourceStore,
    texture_cache: &TextureCache,
    scenepath: &Path,
    args: &Args,
) -> Result<Resource<Prefab>> {
    let path_base = scenepath.parent().unwrap();
    let mut gltf =
        Gltf::from_reader_without_validation(File::open(scenepath)?).context("gltf parsing")?;
    let blob = gltf.blob.take();
    let buffers = import_buffers(&gltf, Some(path_base), blob).context("importing buffers")?;

    let root = gltf.default_scene().ok_or(anyhow!("no default scene"))?;

    // gltf.as_json().asset.copyright
    // eprintln!("{:?}", gltf.extensions_used());
    // eprintln!("{:?}", gltf.extensions());
    // eprintln!("{:?}", root.extensions());
    // eprintln!("{:?}", root.extras());

    let mut nodes = Vec::new();
    fn traverse<'a>(out: &mut Vec<(Affine3A, Node<'a>)>, node: Node<'a>, trans: Affine3A) {
        let trans = trans * transform_to_affine(node.transform());
        for c in node.children() {
            traverse(out, c, trans);
        }
        out.push((trans, node));
    }

    let mut root_affine = Affine3A::IDENTITY;
    root_affine.matrix3 *= args.scale.unwrap_or(1.);
    root_affine.translation *= args.scale.unwrap_or(1.);

    for node in root.nodes() {
        traverse(&mut nodes, node, root_affine);
    }

    let vrm = extract_vrm_data(&gltf)?;

    let mut skin_index_to_arm_index = BTreeMap::new();
    let mut joint_index_to_arm_index = BTreeMap::new();
    let joint_name_to_arm_index: BTreeMap<String, usize>;
    let armature = {
        let mut name = Vec::new();
        let mut parent_pre_map = Vec::new();
        let mut transform = Vec::new();
        let mut inverse_bind_transform = Vec::new();

        for skin in gltf.skins() {
            let mut inverse_bind_mat = skin
                .reader(|buf| Some(&buffers[buf.index()]))
                .read_inverse_bind_matrices();
            for (j_ind, j) in skin.joints().enumerate() {
                let ibm = inverse_bind_mat.as_mut().map(|x| x.next().unwrap());
                let a_ind = match joint_index_to_arm_index.get(&j.index()) {
                    Some(i) => *i,
                    None => {
                        let a_ind = name.len();
                        // name.push(j.name().unwrap_or("").to_string());
                        name.push(format!("bone{}", a_ind));
                        transform.push(transform_to_affine(j.transform()));
                        parent_pre_map.push(
                            gltf.nodes()
                                .find(|n| n.children().any(|c| c.index() == j.index()))
                                .map(|n| n.index()),
                        );
                        inverse_bind_transform.push(
                            ibm.map(|a| transform_to_affine(Transform::Matrix { matrix: a }))
                                .unwrap_or(Affine3A::IDENTITY),
                        );

                        joint_index_to_arm_index.insert(j.index(), a_ind);
                        a_ind
                    }
                };
                skin_index_to_arm_index.insert((skin.index(), j_ind as u16), a_ind as u32);
            }
        }

        let parent = parent_pre_map
            .into_iter()
            .enumerate()
            .map(|(i, p)| {
                p.and_then(|i| joint_index_to_arm_index.get(&i).copied())
                    .unwrap_or(i) as u16
            })
            .collect::<Vec<_>>();

        for (node, bname) in &vrm.bone_node_names {
            let ind = joint_index_to_arm_index[node];
            name[ind] = bname.to_owned();
        }

        joint_name_to_arm_index = name
            .iter()
            .cloned()
            .enumerate()
            .map(|(a, b)| (b, a))
            .collect();

        ArmaturePart {
            name: Some(name),
            parent: Some(parent),
            transform: Some(transform),
            inverse_bind_transform: Some(inverse_bind_transform),
        }
    };

    let head_bones = {
        let pa = armature.parent.clone().unwrap_or_default();
        let na = armature.name.clone().unwrap_or_default();
        (0..pa.len())
            .filter(|&(mut i)| {
                let mut f = false;
                while pa[i] as usize != i {
                    f |= na[i] == "head";
                    i = pa[i] as usize;
                }
                f
            })
            .map(|e| e as u32)
            .collect::<Vec<_>>()
    };

    let prefab = Mutex::new(Prefab::default());
    let node_to_meshes = Mutex::new(BTreeMap::<usize, Vec<usize>>::new());

    nodes
        .par_iter()
        .map(|(trans, node)| {
            // if node.name().unwrap_or_default() == "particles" {
            //     eprintln!("{:?}", node.transform());
            //     eprintln!("{:?}", node.extensions());
            //     eprintln!("{:?}", node.extras());
            // }
            let name = node.name().unwrap_or_default();
            let extras: Value = node
                .extras()
                .to_owned()
                .map(|v| serde_json::from_str(v.get()).unwrap())
                .unwrap_or(Value::Object(Default::default()));

            if !name.ends_with("-collider") {
                if let Some(mesh) = node.mesh() {
                    let meshes = import_mesh(
                        mesh,
                        *trans,
                        &buffers,
                        &store,
                        path_base,
                        node,
                        &args,
                        &texture_cache,
                        &skin_index_to_arm_index,
                        &vrm,
                        &head_bones,
                    )?;
                    for (node, trans, mesh) in meshes {
                        let mut k = prefab.lock().unwrap();
                        let i = k.mesh.len();
                        k.mesh.push((trans, mesh));
                        node_to_meshes
                            .lock()
                            .unwrap()
                            .entry(node)
                            .or_default()
                            .push(i)
                    }
                }
            }
            if extras.get("particles") == Some(&Value::Bool(true)) && !args.no_particles {
                #[derive(Deserialize)]
                struct ParticlesAttr {
                    density: Option<f32>,
                    lifetime: Option<f32>,
                    lifetime_spread: Option<f32>,
                    velocity: Option<Vec3A>,
                    velocity_spread: Option<Vec3A>,
                }
                // let sprite = extras
                //     .get("sprite")
                //     .ok_or(anyhow!("particle volume is missing sprite"))?;

                let attr: ParticlesAttr =
                    serde_json::from_value(extras).context("particles attributes")?;

                info!("adding particles part");

                let part = store.set(&ParticlesPart {
                    sprite: None,
                    density: attr.density,
                    lifetime: attr.lifetime,
                    lifetime_spread: attr.lifetime_spread,
                    velocity: attr.velocity,
                    velocity_spread: attr.velocity_spread,
                })?;

                prefab
                    .lock()
                    .unwrap()
                    .particles
                    .push((transform_to_affine(node.transform()), part));
            }

            if let Some(light) = node.light() {
                let name = node.name().map(|e| e.to_owned());
                if let Some(name) = &name {
                    info!("adding light {name:?}");
                } else {
                    info!("adding light");
                }
                let emission = Some(Vec3A::from_array(light.color()) * light.intensity());
                if let Some(e) = emission {
                    debug!("emission is {e}");
                }
                let (position, _, _) = node.transform().decomposed();
                let part = store.set(&LightPart {
                    emission,
                    name,
                    radius: None,
                })?;
                prefab
                    .lock()
                    .unwrap()
                    .light
                    .push((Vec3A::from_array(position), part));
            }
            {
                let collider = import_physics(&gltf, *trans, node, &store, &buffers)?;
                prefab.lock().unwrap().collision.extend(collider);
            }

            Ok::<_, anyhow::Error>(())
        })
        .reduce(
            || Ok(()),
            |a, b| match (a, b) {
                (Ok(()), Ok(())) => Ok(()),
                (Ok(()), a) => a,
                (a, _) => a,
            },
        )?;

    let mut prefab = prefab.into_inner().unwrap();
    let node_to_meshes = node_to_meshes.into_inner().unwrap();

    if let Some(apath) = args.animation.clone() {
        let path_base = apath.parent().unwrap();
        let mut gltf =
            Gltf::from_reader_without_validation(File::open(&apath)?).context("gltf parsing")?;
        let blob = gltf.blob.take();
        let buffers = import_buffers(&gltf, Some(path_base), blob).context("importing buffers")?;

        let anim_name_map = if let Some(ref map_path) = args.animation_bone_map {
            let mut map = BTreeMap::new();
            for l in read_to_string(map_path)?.lines() {
                if !l.trim().is_empty() && !l.starts_with(";") {
                    let (a, b) = l.split_once("=").unwrap();
                    map.insert(a.to_string(), b.to_string());
                }
            }
            Some(map)
        } else {
            None
        };

        let mut anim_joint_index_to_arm_index = BTreeMap::new();
        let mut joint_index_to_ibm = BTreeMap::new();
        for n in gltf.nodes() {
            if let Some(name) = n.name() {
                let Some(vrm_name) = (if let Some(map) = &anim_name_map {
                    map.get(name).cloned()
                } else {
                    Some(name.to_string())
                }) else {
                    continue;
                };
                let anim_node_index = n.index();
                if let Some(scene_arm_index) = joint_name_to_arm_index.get(&vrm_name) {
                    debug!(
                        "mapping {name:?} (node={}) -> {vrm_name:?} (bone={})",
                        anim_node_index, scene_arm_index
                    );
                    anim_joint_index_to_arm_index.insert(anim_node_index, *scene_arm_index);
                }
            }
        }
        if args.animation_apply_ibm {
            for s in gltf.skins() {
                let reader = s.reader(|buf| Some(&buffers[buf.index()]));
                if let Some(ibms) = reader.read_inverse_bind_matrices() {
                    for (jn, ibm) in s.joints().zip(ibms) {
                        joint_index_to_ibm.insert(
                            jn.index(),
                            transform_to_affine(Transform::Matrix { matrix: ibm }),
                        );
                    }
                }
            }
            debug!("{} joint IBMs found", joint_index_to_ibm.len());
        }
        let transform = args
            .animation_rotation_y
            .map(Affine3A::from_rotation_y)
            .unwrap_or_default()
            * args
                .animation_scale
                .map(Vec3::splat)
                .map(Affine3A::from_scale)
                .unwrap_or_default();
        for a in gltf.animations() {
            prefab.animation.push(import_animation(
                a,
                &store,
                transform,
                &joint_index_to_ibm,
                &anim_joint_index_to_arm_index,
                &BTreeMap::new(),
                &buffers,
            )?);
        }
    }

    for a in gltf.animations() {
        if !args.no_animations {
            prefab.animation.push(import_animation(
                a,
                &store,
                Affine3A::IDENTITY,
                &BTreeMap::new(),
                &joint_index_to_arm_index,
                &node_to_meshes,
                &buffers,
            )?);
        };
    }

    if vrm.camera_mount.is_some()
        || vrm.camera_mount_offset.is_some()
        || !vrm.bone_node_names.is_empty()
    {
        info!("avatar info enabled");
        prefab.avatar_info = Some(
            store.set(&AvatarInfoPart {
                armature: Some(0), // TODO
                camera_mount: vrm
                    .camera_mount
                    .map(|e| joint_index_to_arm_index[&e] as u32),
                camera_mount_offset: vrm.camera_mount_offset,
                ..Default::default()
            })?,
        );
    }

    prefab.armature = if armature.parent.as_ref().is_some_and(|a| !a.is_empty()) {
        vec![store.set(&armature)?]
    } else {
        vec![]
    };

    let skybox = if let Some(skybox) = &args.skybox {
        let mut buf = Vec::new();
        File::open(skybox)?.read_to_end(&mut buf)?;
        Some(store.set(&Image(Cow::Owned(buf)))?)
    } else {
        None
    };
    let sun = if args.with_default_sun {
        Some((
            vec3a(1., -5., 1.).normalize(),
            vec3a(1., 1., 1.).normalize() * 100_000.,
        ))
    } else {
        None
    };
    prefab.environment = if skybox.is_some() || sun.is_some() {
        Some(store.set(&EnvironmentPart { skybox, sun })?)
    } else {
        None
    };

    prefab.name = args.name.clone().or(gltf
        .default_scene()
        .and_then(|n| n.name())
        .map(|n| n.to_owned()));

    if args.debug_light {
        prefab.light.push((
            vec3a(5., 5., 5.),
            store.set(&LightPart {
                name: Some("debug light".to_owned()),
                emission: Some(vec3a(10., 5., 15.)),
                radius: Some(0.3),
            })?,
        ));
    }

    Ok(store.set(&prefab)?)
}