#![feature(random)]
use std::{
    collections::{BTreeMap, BTreeSet},
    env::args,
    fs::File,
    io::BufReader,
};
use unity_tools::{
    serialized_file::{TypeTreeNode, read_serialized_file},
    unityfs::UnityFS,
};

fn main() -> anyhow::Result<()> {
    env_logger::init_from_env("LOG");
    let file = BufReader::new(File::open(args().nth(1).unwrap())?);
    let mut fs = UnityFS::open(file)?;
    let filter_prims = args().any(|a| a == "no_primitives");

    let mut edges = BTreeSet::new();
    let node = fs.find_main_file().unwrap().to_owned();
    let mut cab = fs.read(&node)?;
    let file = read_serialized_file(&mut cab)?;

    for ob in file.objects {
        let typetree = if ob.type_id < 0 {
            unimplemented!()
        } else {
            &file.types[ob.type_id as usize]
        };
        fn print_types(
            edges: &mut BTreeSet<(String, String)>,
            filter_prims: bool,
            tt: TypeTreeNode,
        ) {
            for c in tt.children {
                let mut c = vec![c];
                loop {
                    let mut nc = Vec::new();
                    let mut f = false;
                    for mut c in c {
                        if let Some(inner) = c.type_string.strip_prefix("PPtr<") {
                            c.type_string = inner.strip_suffix(">").unwrap().to_owned();
                            f = true;
                        } else if matches!(c.type_string.as_str(), "Array" | "pair" | "map") {
                            nc.extend(c.children);
                            f = true
                        } else {
                            nc.push(c);
                        };
                    }
                    c = nc;
                    if !f {
                        break;
                    }
                }
                for c in c {
                    if filter_prims && is_primitive(&c.type_string) {
                        continue;
                    }
                    edges.insert((tt.type_string.to_owned(), c.type_string.to_owned()));
                    print_types(edges, filter_prims, c);
                }
            }
        }

        if let Some(tree) = &typetree.type_tree {
            print_types(&mut edges, filter_prims, tree.clone());
        }
    }
    let nodes = edges
        .iter()
        .flat_map(|(k, v)| [k, v])
        .collect::<BTreeSet<_>>();
    let ids = nodes
        .into_iter()
        .enumerate()
        .map(|(i, n)| (n, i))
        .collect::<BTreeMap<_, _>>();

    println!("digraph {{");
    for (name, id) in &ids {
        println!("t{id} [label={name:?}]");
    }
    for (a, b) in &edges {
        println!("t{} -> t{}", ids[a], ids[b]);
    }
    println!("}}");

    Ok(())
}

fn is_primitive(s: &str) -> bool {
    matches!(
        s,
        "Type*"
            | "int"
            | "unsigned int"
            | "UInt8"
            | "UInt16"
            | "UInt32"
            | "UInt64"
            | "SInt8"
            | "SInt16"
            | "SInt32"
            | "SInt64"
            | "bool"
            | "float"
            | "vector"
            | "string"
            | "float3"
            | "float4"
            | "float2"
            | "Vector2f"
            | "Vector3f"
            | "Vector4f"
            | "Quaternionf"
            | "xform"
    )
}