use anyhow::{Context, Result};
use cid::Cid;
use iroh_car::CarReader;
use star::{RepoMstNode, StarCommit, StarMstEntry, StarMstNode, StarSerializer, calculate_height};
use std::collections::HashMap;
use std::env;

#[tokio::main]
async fn main() -> Result<()> {
    let args: Vec<String> = env::args().collect();
    if args.len() != 3 {
        eprintln!("Usage: car-to-star <input.car> <output.star>");
        std::process::exit(1);
    }

    let car_path = &args[1];
    let star_path = &args[2];

    let reader = tokio::fs::File::open(car_path).await?;
    let reader = tokio::io::BufReader::new(reader);

    println!("Reading CAR file...");
    let mut car = CarReader::new(reader).await?;
    let roots = car.header().roots();
    assert_eq!(roots.len(), 1);

    let commit_cid = *roots.first().expect("a root to be present");

    let mut blocks: HashMap<Cid, Vec<u8>> = HashMap::new();

    while let Some((cid, data)) = car.next_block().await? {
        blocks.insert(cid, data);
    }
    println!("Loaded {} blocks.", blocks.len());

    let output_file = std::fs::File::create(star_path)?;

    let commit_bytes = blocks.get(&commit_cid).context("Commit block not found")?;

    #[derive(serde::Deserialize)]
    struct RepoCommit {
        did: String,
        version: i64,
        data: Cid,
        rev: String,
        prev: Option<Cid>,
        sig: Option<serde_bytes::ByteBuf>,
    }

    let repo_commit: RepoCommit = serde_ipld_dagcbor::from_slice(commit_bytes)?;

    let root_bytes = blocks
        .get(&repo_commit.data)
        .context("repo data cannot be null")?;
    let root_node: RepoMstNode =
        serde_ipld_dagcbor::from_slice(root_bytes).context("root must be an mst node")?;

    let star_data = if root_node.l.is_none() && root_node.e.is_empty() {
        None
    } else {
        Some(repo_commit.data)
    };

    let star_commit = StarCommit {
        did: repo_commit.did,
        version: repo_commit.version,
        data: star_data,
        rev: repo_commit.rev,
        prev: repo_commit.prev,
        sig: repo_commit.sig,
    };

    let mut serializer = StarSerializer::new(output_file);

    serializer.write_header(&star_commit)?;
    println!("wrote header. Root: {}", repo_commit.data);

    if let Some(root_cid) = star_commit.data {
        println!("writing tree...");
        let (nodes, records) = write_tree(root_cid, &blocks, &mut serializer)?;
        println!("wrote {nodes} nodes and {records} records.");
    } else {
        println!("empty MST, no tree written.");
    }

    serializer.finish()?;
    println!("Done!");
    Ok(())
}

fn write_tree(
    node_cid: Cid,
    blocks: &HashMap<Cid, Vec<u8>>,
    serializer: &mut StarSerializer<std::fs::File>,
) -> Result<(usize, usize)> {
    // println!("writing tree under {node_cid:?}...");

    let mut nodes_written = 0;
    let mut records_written = 0;

    let block_bytes = blocks
        .get(&node_cid)
        .with_context(|| format!("Missing block {}", node_cid))?;

    let repo_node: RepoMstNode = serde_ipld_dagcbor::from_slice(block_bytes)?;

    let height = if let Some(first_entry) = repo_node.e.first() {
        calculate_height(&first_entry.k)
    } else {
        0
    };

    let star_node = StarMstNode {
        l: repo_node.l,
        l_archived: repo_node.l.map(|_| true),
        e: repo_node
            .e
            .iter()
            .map(|e| {
                let v = if height == 0 { None } else { Some(e.v) };
                StarMstEntry {
                    p: e.p,
                    k: e.k.clone(),
                    v,
                    v_archived: Some(true),
                    t: e.t,
                    t_archived: e.t.map(|_| true),
                }
            })
            .collect(),
    };

    serializer.write_node(&star_node)?;
    nodes_written += 1;

    if let Some(l_cid) = repo_node.l {
        let (n, r) = write_tree(l_cid, blocks, serializer)?;
        nodes_written += n;
        records_written += r;
    }

    for e in repo_node.e {
        let record_bytes = blocks
            .get(&e.v)
            .with_context(|| format!("Missing record {}", e.v))?;
        // eprintln!("writing record {:?} (<= {node_cid:?})", e.v);
        serializer.write_record(record_bytes)?;
        records_written += 1;

        if let Some(t_cid) = e.t {
            let (n, r) = write_tree(t_cid, blocks, serializer)?;
            nodes_written += n;
            records_written += r;
        }
    }

    Ok((nodes_written, records_written))
}