/*
    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/>.
*/
#![feature(iter_array_chunks)]
#![allow(clippy::too_many_arguments, clippy::type_complexity)]
pub mod animation;
pub mod mesh;
pub mod physics;
pub mod vrm;

use animation::import_animation;
use anyhow::{Context, Result, anyhow};
use clap::Parser;
use glam::Vec3;
use gltf::{Gltf, Node, image::Source, import_buffers, scene::Transform};
use humansize::BINARY;
use image::{ImageReader, codecs::webp::WebPEncoder};
use log::{debug, info};
use mesh::import_mesh;
use physics::import_physics;
use rand::random;
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use serde::Deserialize;
use serde_json::Value;
use std::{
    borrow::Cow,
    collections::{BTreeMap, HashMap},
    fs::{File, read_to_string},
    io::{Cursor, Read, Write},
    marker::PhantomData,
    net::{SocketAddr, TcpStream},
    path::{Path, PathBuf},
    sync::{Arc, Mutex},
    thread::{self, sleep},
    time::Duration,
};
use vrm::extract_vrm_data;
use weareshared::{
    Affine3A, Vec3A,
    helper::ReadWrite,
    packets::{Data, Object, Packet, Resource},
    resources::{
        ArmaturePart, AvatarInfoPart, EnvironmentPart, Image, LightPart, ParticlesPart, Prefab,
    },
    store::ResourceStore,
    vec3a,
};

#[derive(Parser)]
pub struct Args {
    address: SocketAddr,
    /// Path(s) to a glTF file, binary or json format
    scene: Vec<PathBuf>,
    /// Send all resources to the server then quit
    #[arg(short, long)]
    push: bool,

    /// Remove all other object from the world
    #[arg(short, long)]
    clear: bool,
    /// Add the object to the world
    #[arg(short, long)]
    add: bool,
    /// Transcode all textures to WebP
    #[arg(short, long)]
    webp: bool,
    /// Add skybox
    #[arg(long)]
    skybox: Option<PathBuf>,
    /// Override prefab name
    #[arg(short, long)]
    name: Option<String>,
    #[arg(short, long)]
    scale: Option<f32>,
    #[arg(short, long)]
    dry_run: bool,
    #[arg(short, long)]
    line_up: bool,

    #[arg(long)]
    use_cache: bool,

    #[arg(short = 'S', long)]
    with_default_sun: bool,

    #[arg(long)]
    animation: Option<PathBuf>,
    #[arg(long)]
    animation_bone_map: Option<PathBuf>,
    #[arg(long)]
    animation_rotation_y: Option<f32>,
    #[arg(long)]
    animation_scale: Option<f32>,
    #[arg(long)]
    animation_apply_ibm: bool,

    /// Spins the object
    #[arg(long)]
    debug_spin: bool,
    /// Adds a light
    #[arg(long)]
    debug_light: bool,
    /// Wiggles joint index 6
    #[arg(long)]
    debug_armature: bool,
    #[arg(long)]
    no_particles: bool,
    #[arg(long)]
    no_animations: bool,
}

fn main() -> Result<()> {
    env_logger::init_from_env("LOG");
    let mut args = Args::parse();

    let store = if args.use_cache {
        ResourceStore::new_env()?
    } else {
        ResourceStore::new_memory()
    };

    let mut prefabs = Vec::new();
    let texture_cache = Arc::new(Mutex::new(HashMap::new()));

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

    let mut debug_bone = (0, Affine3A::IDENTITY);

    for scenepath in &args.scene {
        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));
        }

        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();
            }

            for (i, (name, tr)) in name.iter().zip(transform.iter()).enumerate() {
                if name == "head" {
                    debug_bone = (i, *tr);
                }
            }

            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.take() {
            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),
                })?,
            ));
        }

        prefabs.push(store.set(&prefab)?);
    }

    let mut size = 0;
    store.iter(|_k, len| size += len).unwrap();
    info!(
        "prefab has network size of {}",
        humansize::format_size(size, BINARY)
    );

    if args.dry_run {
        return Ok(());
    }

    let mut sock = TcpStream::connect(args.address)?;
    Packet::Connect(random()).write(&mut sock)?;
    for p in &prefabs {
        Packet::AnnouncePrefab(p.clone()).write(&mut sock)?;
    }
    sock.flush()?;

    let mut obs = Vec::new();
    if args.add {
        for (i, p) in prefabs.iter().enumerate() {
            let ob = Object::new();
            info!("adding object {ob}");
            Packet::Add(ob, p.clone()).write(&mut sock)?;
            if args.line_up {
                Packet::Position(ob, vec3a(i as f32 * 1.2, 0., i as f32 * 0.3), Vec3A::ZERO)
                    .write(&mut sock)?;
            }
            obs.push(ob);
        }
        sock.flush()?;
    }

    if args.debug_spin {
        let ob = obs[0];
        let mut sock2 = sock.try_clone().unwrap();
        thread::spawn(move || {
            let mut x = 0.;
            loop {
                Packet::Position(ob, Vec3A::ZERO, vec3a(0., x * 0.1, 0.))
                    .write(&mut sock2)
                    .unwrap();
                sock2.flush().unwrap();
                x += 0.1;
                sleep(Duration::from_millis(50));
            }
        });
    }
    if args.debug_armature {
        let ob = obs[0];
        let mut sock2 = sock.try_clone().unwrap();
        thread::spawn(move || {
            let mut x = 0f32;
            loop {
                let a = Affine3A::IDENTITY
                    * Affine3A::from_rotation_x(x.cos() * 0.3)
                    * Affine3A::from_rotation_y(x.sin() * 0.5)
                    * debug_bone.1;
                Packet::Pose(ob, vec![(debug_bone.0 as u16, a)])
                    .write(&mut sock2)
                    .unwrap();
                sock2.flush().unwrap();
                x += 0.2;
                sleep(Duration::from_millis(50));
            }
        });
    }

    if args.push {
        if args.use_cache {
            return Ok(());
        }
        store.iter(|k, _| {
            Packet::RespondResource(k, Data(store.get_raw(k).unwrap().unwrap()))
                .write(&mut sock)
                .unwrap();
        })?;
        sock.flush()?;
    } else {
        loop {
            let packet = Packet::read(&mut sock)?;
            match packet {
                Packet::RequestResource(hash) => {
                    if let Some(d) = store.get_raw(hash)? {
                        Packet::RespondResource(hash, Data(d)).write(&mut sock)?;
                        sock.flush()?;
                    }
                }
                Packet::Add(ob_a, _) => {
                    if args.clear && !obs.contains(&ob_a) {
                        info!("removing object {ob_a}");
                        Packet::Remove(ob_a).write(&mut sock)?;
                        sock.flush()?;
                    }
                }
                _ => (),
            }
        }
    }
    Ok(())
}

pub type TextureCache = Arc<Mutex<HashMap<String, Resource<Image<'static>>>>>;
fn load_texture(
    name: &str,
    store: &ResourceStore,
    path: &Path,
    buffers: &[gltf::buffer::Data],
    source: &Source,
    webp: bool,
    texture_cache: &TextureCache,
) -> Result<Resource<Image<'static>>> {
    let (mut image, uri) = match source {
        gltf::image::Source::View { view, mime_type } => {
            debug!("{name} texture is embedded and of type {mime_type:?}");
            let buf =
                &buffers[view.buffer().index()].0[view.offset()..view.offset() + view.length()];
            (Image(Cow::Borrowed(buf)), None)
        }
        gltf::image::Source::Uri {
            uri,
            mime_type: Some(mime_type),
        } => {
            debug!("{name} texture is {uri:?} and of type {mime_type:?}");
            if let Some(res) = texture_cache.lock().unwrap().get(*uri) {
                return Ok(res.to_owned());
            }
            let path = path.join(uri);
            let mut buf = Vec::new();
            File::open(path)?.read_to_end(&mut buf)?;
            (Image(buf.into()), Some(uri.to_string()))
        }
        gltf::image::Source::Uri {
            uri,
            mime_type: None,
        } => {
            debug!("{name} texture is {uri:?} and has no type");
            if let Some(res) = texture_cache.lock().unwrap().get(*uri) {
                return Ok(res.to_owned());
            }
            let path = path.join(uri);
            let mut buf = Vec::new();
            File::open(path)?.read_to_end(&mut buf)?;
            (Image(buf.into()), Some(uri.to_string()))
        }
    };

    if webp {
        let mut image_out = Vec::new();

        let len = image.0.len();
        ImageReader::new(Cursor::new(image.0))
            .with_guessed_format()?
            .decode()?
            .write_with_encoder(WebPEncoder::new_lossless(&mut image_out))?;
        debug!("webp encode: {len} -> {}", image_out.len());
        image = Image(Cow::Owned(image_out));
    }
    let res = Resource(store.set(&image)?.0, PhantomData);
    if let Some(uri) = uri {
        texture_cache.lock().unwrap().insert(uri, res.clone());
    }
    Ok(res)
}

pub fn transform_to_affine(trans: Transform) -> Affine3A {
    let mat = trans.matrix();
    Affine3A::from_cols_array_2d(&[
        [mat[0][0], mat[0][1], mat[0][2]],
        [mat[1][0], mat[1][1], mat[1][2]],
        [mat[2][0], mat[2][1], mat[2][2]],
        [mat[3][0], mat[3][1], mat[3][2]],
    ])
}