use std::collections::{BTreeMap, BTreeSet};
use std::env;
use std::fs;
use std::path::{Path, PathBuf};

#[derive(Debug)]
struct AssetFile {
    virtual_path: String,
    disk_path: PathBuf,
}

fn main() {
    let manifest_dir =
        PathBuf::from(env::var("CARGO_MANIFEST_DIR").expect("missing CARGO_MANIFEST_DIR"));
    let asset_root = manifest_dir.join("src/asset");
    let out_file = PathBuf::from(env::var("OUT_DIR").expect("missing OUT_DIR")).join("asset.rs");

    println!("cargo::rerun-if-changed={}", asset_root.display());

    let mut directory_entries = BTreeMap::<String, BTreeSet<String>>::new();
    directory_entries
        .entry(String::new())
        .or_default()
        .insert(String::from("asset"));

    let mut asset_files = Vec::new();
    collect_assets(
        &asset_root,
        "asset",
        &mut directory_entries,
        &mut asset_files,
    );

    let generated = render_assets(&asset_files, &directory_entries);

    fs::write(out_file, generated).expect("failed to write generated assets module");
}

fn collect_assets(
    disk_dir: &Path,
    virtual_dir: &str,
    directory_entries: &mut BTreeMap<String, BTreeSet<String>>,
    asset_files: &mut Vec<AssetFile>,
) {
    let mut entries = fs::read_dir(disk_dir)
        .unwrap_or_else(|err| panic!("failed to read {}: {err}", disk_dir.display()))
        .map(|entry| entry.expect("failed to read directory entry"))
        .collect::<Vec<_>>();

    entries.sort_by_key(|entry| entry.file_name());

    for entry in entries {
        let file_name = entry
            .file_name()
            .into_string()
            .unwrap_or_else(|_| panic!("non-utf8 asset name in {}", disk_dir.display()));
        let child_virtual_path = format!("{virtual_dir}/{file_name}");
        let path = entry.path();

        directory_entries
            .entry(String::from(virtual_dir))
            .or_default()
            .insert(file_name);

        if path.is_dir() {
            collect_assets(&path, &child_virtual_path, directory_entries, asset_files);
        } else if path.is_file() {
            asset_files.push(AssetFile {
                virtual_path: child_virtual_path,
                disk_path: path,
            });
        }
    }
}

fn render_assets(
    asset_files: &[AssetFile],
    directory_entries: &BTreeMap<String, BTreeSet<String>>,
) -> String {
    let mut output = String::new();

    output.push_str(
        "pub fn load_asset(path: &str) -> gpui::Result<Option<std::borrow::Cow<'static, [u8]>>> {\n",
    );
    output.push_str("    match path {\n");
    for file in asset_files {
        output.push_str("        ");
        output.push_str(&string_literal(&file.virtual_path));
        output.push_str(" => Ok(Some(std::borrow::Cow::Borrowed(include_bytes!(");
        output.push_str(&string_literal(&file.disk_path.to_string_lossy()));
        output.push_str(")))),\n");
    }
    output.push_str("        _ => Err(anyhow::anyhow!(\"asset not found: {path}\")),\n");
    output.push_str("    }\n");
    output.push_str("}\n\n");

    output.push_str("pub fn list_assets(path: &str) -> gpui::Result<Vec<gpui::SharedString>> {\n");
    output.push_str("    let normalized = path.trim_end_matches('/');\n");
    output.push_str("    let normalized = if normalized == \".\" { \"\" } else { normalized };\n");
    output.push_str("    let entries: &[&str] = match normalized {\n");
    for (directory, entries) in directory_entries {
        output.push_str("        ");
        output.push_str(&string_literal(directory));
        output.push_str(" => &[\n");
        for entry in entries {
            output.push_str("            ");
            output.push_str(&string_literal(entry));
            output.push_str(",\n");
        }
        output.push_str("        ],\n");
    }
    output.push_str("        _ => &[],\n");
    output.push_str("    };\n");
    output.push_str("    Ok(entries.iter().copied().map(gpui::SharedString::from).collect())\n");
    output.push_str("}\n");

    output
}

fn string_literal(value: &str) -> String {
    format!("{value:?}")
}