pub mod spatial;

use anyhow::Result;
use glam::{Affine3A, DAffine3, DVec3};
use indicatif::ProgressIterator;
use log::{debug, info};
use osm_pbf_reader::{Blobs, data::primitives::Primitive};
use rayon::iter::{ParallelBridge, ParallelIterator};
use spatial::{SpatialTree, octtree::Octtree};
use std::{collections::BTreeMap, env::args, f64::consts::PI, fs::File, io::BufWriter};
use weareshared::{
    graphics::GraphicsPart,
    helper::AABB,
    packets::Resource,
    resources::{Prefab, RespackEntry, SpatialIndex},
    respack::save_full_respack,
    store::ResourceStore,
};

fn main() -> Result<()> {
    env_logger::init_from_env("LOG");
    let inpath = args().nth(1).unwrap();
    let outpath = args().nth(2).unwrap();

    info!("Indexing node positions...");
    let node_positions = Blobs::from_path(&inpath)?
        .progress_count(1_500)
        .par_bridge()
        .map(|e| {
            let mut e = e.unwrap();
            let e = e.decode().unwrap();
            let mut npos = Vec::new();
            for p in e.primitives() {
                if let Primitive::Node(n) = p {
                    let lat = n.nano_lat as f64 / 1_000_000_000. / 180. * PI;
                    let lon = n.nano_lon as f64 / 1_000_000_000. / 180. * PI;
                    let pos = DVec3::new(lat.sin() * lon.cos(), lon.sin(), lat.cos() * lon.cos());
                    npos.push((map_nodeid(n.id), pos));
                };
            }
            npos
        })
        .fold(
            || BTreeMap::new(),
            |mut a, b| {
                a.extend(b);
                a
            },
        )
        .reduce(
            || BTreeMap::new(),
            |mut a, b| {
                if !a.is_empty() && !b.is_empty() {
                    debug!("merge {} + {}", a.len(), b.len());
                }
                a.extend(b);
                a
            },
        );

    let mut tree = Octtree::new(DVec3::splat(1.5));
    info!("Building tree over node positions...");
    for (id, pos) in node_positions.into_iter().progress() {
        tree.insert(pos, id);
    }
    info!("Done, depth={}", tree.depth());

    let store = ResourceStore::new_memory();

    info!("Exporting resources...");
    let (_, root) = export_level(&tree, &store, &mut Vec::new())?;
    let entry = Some(store.set(&RespackEntry {
        c_spatial_index: vec![root],
        ..Default::default()
    })?);
    info!("Done, N={}", store.count()?);

    let output = BufWriter::new(File::create(outpath)?);
    save_full_respack(output, &store, entry)?;

    Ok(())
}

fn export_level(
    node: &Octtree,
    store: &ResourceStore,
    important_features: &mut Vec<DVec3>,
) -> Result<(AABB, Resource<SpatialIndex>)> {
    let mut lod_features = Vec::new();
    let child = node
        .children
        .iter()
        .flat_map(|e| e.as_slice())
        .flatten()
        .flatten()
        .map(|o| export_level(o, store, &mut lod_features))
        .collect::<Result<Vec<_>>>()?;

    let prefab = {
        let points = if node.elems.is_empty() {
            lod_features
        } else {
            node.elems.iter().map(|(e, _id)| *e).collect()
        };

        let local_origin = node.center;
        let mut g = GraphicsPart::default();

        use weareshared::graphics::GraphicsCommand::*;
        g.push(StrokeWidth(0.001));
        g.push(Stroke(5));
        for (i, p) in points.into_iter().enumerate() {
            g.push(Point((p - local_origin).as_vec3()));
            if i % 8 == 0 {
                important_features.push(p);
            }
        }

        Some(store.set(&Prefab {
            transform: Some(DAffine3::from_translation(local_origin)),
            graphics: vec![(Affine3A::IDENTITY, store.set(&g)?)],
            ..Default::default()
        })?)
    };

    Ok((
        AABB {
            min: (node.center - node.size / 2.).as_vec3(),
            max: (node.center + node.size / 2.).as_vec3(),
        },
        store.set(&SpatialIndex {
            level: None,
            prefab,
            child,
        })?,
    ))
}

type NodeID = u64;

fn map_nodeid(id: i64) -> u64 {
    // fn xorshift(mut x: u64) -> u64 {
    //     x ^= x << 13;
    //     x ^= x >> 7;
    //     x ^= x << 17;
    //     x
    // }
    // xorshift(xorshift(xorshift(id as u64)))
    id as u64
}