/*
    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)]
use anyhow::{Result, bail};
use clap::Parser;
use gltf::image::Source;
use log::{debug, info};
use rand::random;
use std::{
    fs::File,
    io::{Read, Write},
    net::{SocketAddr, TcpStream},
    path::{Path, PathBuf},
    thread::{self, sleep},
    time::Duration,
};
use weareshared::{
    Affine3A, Vec3A,
    packets::{Data, Object, Packet, ReadWrite, Resource},
    resources::{AttributeArray, IndexArray, MeshPart, Prefab},
    store::ResourceStore,
    vec3a,
};

#[derive(Parser)]
struct Args {
    address: SocketAddr,
    scene: PathBuf,
    #[arg(short, long)]
    push: bool,
    #[arg(short, long)]
    spin: bool,
}

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

    let mut sock = TcpStream::connect(args.address)?;
    let store = ResourceStore::new_memory();

    Packet::Connect(random()).write(&mut sock)?;

    let (gltf, buffers, _) = gltf::import(&args.scene)?;
    let path_base = args.scene.parent().unwrap();

    let mut prefab = Prefab::default();
    for node in gltf.nodes() {
        if let Some(mesh) = node.mesh() {
            for p in mesh.primitives() {
                let reader = p.reader(|buf| Some(&buffers[buf.index()]));

                let va_position = reader
                    .read_positions()
                    .map(|iter| {
                        let mut pos_x = vec![];
                        let mut pos_y = vec![];
                        let mut pos_z = vec![];
                        for p in iter {
                            pos_x.push(p[0]);
                            pos_y.push(p[1]);
                            pos_z.push(p[2]);
                        }
                        info!("{} vertex positions", pos_x.len());
                        Ok::<_, anyhow::Error>([
                            store.set(&AttributeArray(pos_x).write_alloc())?,
                            store.set(&AttributeArray(pos_y).write_alloc())?,
                            store.set(&AttributeArray(pos_z).write_alloc())?,
                        ])
                    })
                    .transpose()?;

                let va_normal = reader
                    .read_normals()
                    .map(|iter| {
                        let mut normal_x = vec![];
                        let mut normal_y = vec![];
                        let mut normal_z = vec![];
                        for p in iter {
                            normal_x.push(p[0]);
                            normal_y.push(p[1]);
                            normal_z.push(p[2]);
                        }
                        info!("{} vertex normals", normal_x.len());
                        Ok::<_, anyhow::Error>([
                            store.set(&AttributeArray(normal_x).write_alloc())?,
                            store.set(&AttributeArray(normal_y).write_alloc())?,
                            store.set(&AttributeArray(normal_z).write_alloc())?,
                        ])
                    })
                    .transpose()?;

                let va_texcoord = reader
                    .read_tex_coords(0)
                    .map(|iter| {
                        let mut texcoord_u = vec![];
                        let mut texcoord_v = vec![];
                        for p in iter.into_f32() {
                            texcoord_u.push(p[0]);
                            texcoord_v.push(p[1]);
                        }
                        info!("{} vertex texture coordinates", texcoord_u.len());
                        Ok::<_, anyhow::Error>([
                            store.set(&AttributeArray(texcoord_u).write_alloc())?,
                            store.set(&AttributeArray(texcoord_v).write_alloc())?,
                        ])
                    })
                    .transpose()?;

                let va_albedo = reader
                    .read_colors(0)
                    .map(|iter| {
                        let mut color_r = vec![];
                        let mut color_g = vec![];
                        let mut color_b = vec![];
                        for p in iter.into_rgb_f32() {
                            color_r.push(p[0]);
                            color_g.push(p[1]);
                            color_b.push(p[2]);
                        }
                        info!("{} vertex colors", color_r.len());
                        Ok::<_, anyhow::Error>([
                            store.set(&AttributeArray(color_r).write_alloc())?,
                            store.set(&AttributeArray(color_g).write_alloc())?,
                            store.set(&AttributeArray(color_b).write_alloc())?,
                        ])
                    })
                    .transpose()?;

                let va_transmission = reader
                    .read_colors(0)
                    .map(|iter| {
                        let mut color_a = vec![];
                        for p in iter.into_rgba_f32() {
                            color_a.push(p[3]);
                        }
                        let o = if color_a.iter().any(|x| *x != 1.) {
                            info!("{} vertex transmissions", color_a.len());
                            Some(store.set(&AttributeArray(color_a).write_alloc())?)
                        } else {
                            debug!("vertex transmission pruned");
                            None
                        };
                        Ok::<_, anyhow::Error>(o)
                    })
                    .transpose()?
                    .flatten();

                let index = reader
                    .read_indices()
                    .unwrap()
                    .into_u32()
                    .map(|e| e as u16)
                    .array_chunks::<3>()
                    .collect::<Vec<_>>();
                info!("{} indecies", index.len() * 3);
                let index = Some(store.set(&IndexArray(index).write_alloc())?);

                let mut tex_albedo = None;
                let mut tex_transmission = None;
                if let Some(tex) = p.material().pbr_metallic_roughness().base_color_texture() {
                    let r = load_texture(
                        "albedo",
                        &store,
                        path_base,
                        &buffers,
                        &tex.texture().source().source(),
                    )?;
                    tex_albedo = Some(r);
                    tex_transmission = Some(r);
                }
                let mut tex_normal = None;
                if let Some(tex) = p.material().normal_texture() {
                    tex_normal = Some(load_texture(
                        "normal",
                        &store,
                        path_base,
                        &buffers,
                        &tex.texture().source().source(),
                    )?);
                }
                let mut tex_emission = None;
                if let Some(tex) = p.material().emissive_texture() {
                    tex_emission = Some(load_texture(
                        "emission",
                        &store,
                        path_base,
                        &buffers,
                        &tex.texture().source().source(),
                    )?);
                }
                let mut tex_roughness = None;
                let mut tex_metallic = None;
                if let Some(tex) = p
                    .material()
                    .pbr_metallic_roughness()
                    .metallic_roughness_texture()
                {
                    let r = load_texture(
                        "metallic+roughness",
                        &store,
                        path_base,
                        &buffers,
                        &tex.texture().source().source(),
                    )?;
                    tex_roughness = Some(r);
                    tex_metallic = Some(r);
                }

                let g_metallic = Some(p.material().pbr_metallic_roughness().metallic_factor());
                let g_roughness = Some(p.material().pbr_metallic_roughness().roughness_factor());

                let base_color = p.material().pbr_metallic_roughness().base_color_factor();

                let g_albedo = if base_color[0] != 1. || base_color[1] != 1. || base_color[2] != 1.
                {
                    info!(
                        "global albedo is r={},g={},b={}",
                        base_color[0], base_color[1], base_color[2]
                    );
                    Some(Vec3A::new(base_color[0], base_color[1], base_color[2]))
                } else {
                    debug!("global albedo pruned");
                    None
                };
                let g_transmission = if base_color[3] != 1. {
                    info!("global transmission is {}", base_color[3]);
                    Some(base_color[3])
                } else {
                    debug!("global transmission pruned");
                    None
                };

                let emission = p.material().emissive_factor();
                let g_emission = if emission[0] != 0. || emission[1] != 0. || emission[2] != 0. {
                    info!(
                        "global emission is r={},g={},b={}",
                        base_color[0], base_color[1], base_color[2]
                    );
                    Some(Vec3A::new(emission[0], emission[1], emission[2]))
                } else {
                    debug!("global emission pruned");
                    None
                };

                let mesh = store.set(
                    &MeshPart {
                        g_albedo,
                        g_transmission,
                        g_metallic,
                        g_roughness,
                        g_emission,
                        va_position,
                        va_normal,
                        va_texcoord,
                        va_albedo,
                        va_transmission,
                        tex_albedo,
                        tex_normal,
                        tex_roughness,
                        tex_metallic,
                        tex_transmission,
                        tex_emission,
                        index,
                        va_emission: None,  // not supported by gltf
                        va_metallic: None,  // not supported by gltf
                        va_roughness: None, // not supported by gltf
                    }
                    .write_alloc(),
                )?;
                let mat = node.transform().matrix();
                let aff = 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]],
                ]);
                prefab.mesh.push((aff, mesh))
            }
        }
    }

    let ob = Object::new();
    Packet::Add(ob, store.set(&prefab.write_alloc())?).write(&mut sock)?;

    if args.spin {
        let mut sock2 = sock.try_clone().unwrap();
        thread::spawn(move || {
            let mut x = 0.;
            loop {
                Packet::Position(ob, Vec3A::ZERO, vec3a(x, x * 0.3, x * 0.1))
                    .write(&mut sock2)
                    .unwrap();
                sock2.flush().unwrap();
                x += 0.1;
                sleep(Duration::from_millis(50));
            }
        });
    }

    if args.push {
        store.iter(|d| {
            Packet::RespondResource(Data(d.to_vec()))
                .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(hash)? {
                        Packet::RespondResource(Data(d)).write(&mut sock)?;
                        sock.flush()?;
                    }
                }
                _ => (),
            }
        }
    }
    Ok(())
}

fn load_texture(
    name: &str,
    store: &ResourceStore,
    path: &Path,
    buffers: &[gltf::buffer::Data],
    source: &Source,
) -> Result<Resource> {
    match source {
        gltf::image::Source::View { view, mime_type } => {
            info!("{name} texture is embedded and of type {mime_type:?}");
            let buf =
                &buffers[view.buffer().index()].0[view.offset()..view.offset() + view.length()];
            return Ok(store.set(buf)?);
        }
        gltf::image::Source::Uri {
            uri,
            mime_type: Some(mime_type),
        } => {
            info!("{name} texture is at {uri:?} and of type {mime_type:?}");
            let path = path.join(uri);
            let mut buf = Vec::new();
            File::open(path)?.read_to_end(&mut buf)?;
            return Ok(store.set(&buf)?);
        }
        _ => {
            bail!("texture is external and has no mime type")
        }
    }
}