//! Build script: reads all files from `dist/`, compresses them (brotli + gzip),
//! computes etags via SHA-256, and generates a `frontend` Rust module for `dist.rs`.

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

use sha2::{Digest, Sha256};
use flate2::write::GzEncoder;
use flate2::Compression;

// ── Compression helpers ──────────────────────────────────────────────────

fn gzip_compress(data: &[u8]) -> Vec<u8> {
    let mut encoder = GzEncoder::new(Vec::new(), Compression::new(6));
    encoder.write_all(data).unwrap();
    encoder.finish().unwrap()
}

fn brotli_compress(data: &[u8]) -> Option<Vec<u8>> {
    use brotli::CompressorWriter;
    let buf = Vec::new();
    let mut writer = CompressorWriter::new(buf, 4096, 6, 16);
    if writer.write_all(data).is_ok() && writer.flush().is_ok() {
        Some(writer.into_inner())
    } else {
        None
    }
}

// ── ETag computation ─────────────────────────────────────────────────────

fn compute_etag(data: &[u8]) -> String {
    let mut hasher = Sha256::new();
    hasher.update(data);
    let hash = hasher.finalize();
    // First 32 hex chars for a compact etag
    format!("{:x}", hash)[..32].to_string()
}

// ── Asset collection ─────────────────────────────────────────────────────

struct Asset {
    path: String,
    data: Vec<u8>,
    etag: String,
    brotli: Option<Vec<u8>>,
    gzip: Vec<u8>,
}

fn collect_assets(dist_dir: &Path) -> BTreeMap<String, Asset> {
    let mut assets = BTreeMap::new();

    for entry in walkdir(dist_dir) {
        let rel = entry.strip_prefix(dist_dir).unwrap();
        let path_str = rel.to_string_lossy().replace('\\', "/");
        if path_str.is_empty() {
            continue;
        }

        let data = fs::read(&entry).unwrap_or_else(|e| {
            panic!("Failed to read dist file {}: {}", path_str, e)
        });

        let etag = compute_etag(&data);
        let brotli_data = brotli_compress(&data);
        let gzip_data = gzip_compress(&data);

        assets.insert(
            path_str.clone(),
            Asset {
                path: path_str,
                data,
                etag,
                brotli: brotli_data,
                gzip: gzip_data,
            },
        );
    }

    assets
}

fn walkdir(dir: &Path) -> Vec<PathBuf> {
    let mut files = Vec::new();
    if let Ok(entries) = fs::read_dir(dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                files.extend(walkdir(&path));
            } else {
                files.push(path);
            }
        }
    }
    files
}

// ── Code generation ──────────────────────────────────────────────────────

fn rust_byte_literal(data: &[u8]) -> String {
    if data.len() < 200 {
        let bytes: Vec<String> = data.iter().map(|b| b.to_string()).collect();
        format!("[{}]", bytes.join(", "))
    } else {
        let lines: Vec<String> = data
            .chunks(80)
            .map(|chunk| {
                chunk.iter().map(|b| b.to_string()).collect::<Vec<_>>().join(", ")
            })
            .collect();
        format!("[\n{}\n]", lines.join(",\n"))
    }
}

fn path_to_ident(path: &str) -> String {
    let s = path.replace('-', "_").replace('.', "_").replace('/', "_").to_uppercase();
    format!("ASSET_{s}")
}

fn etag_ident(path: &str) -> String {
    format!("ETAG_{}", path_to_ident(path))
}

fn br_ident(path: &str) -> String {
    format!("{}_BR", path_to_ident(path))
}

fn gz_ident(path: &str) -> String {
    format!("{}_GZ", path_to_ident(path))
}

fn generate_frontend_module(assets: &BTreeMap<String, Asset>, out_dir: &Path) {
    let mut code = String::new();

    code += "// AUTO-GENERATED by build.rs — DO NOT EDIT.\n";
    code += "// Frontend assets from dist/ embedded as static byte arrays.\n\n";

    // Generate static byte arrays for each asset
    for (path, asset) in assets {
        let ident = path_to_ident(path);
        let etag_id = etag_ident(path);
        let br_id = br_ident(path);
        let gz_id = gz_ident(path);

        code += &format!("static {}: &[u8] = &{};\n", ident, rust_byte_literal(&asset.data));
        code += &format!("static {}: &str = \"{}\";\n", etag_id, asset.etag);
        if let Some(ref br) = asset.brotli {
            code += &format!("static {}: &[u8] = &{};\n", br_id, rust_byte_literal(br));
        }
        code += &format!("static {}: &[u8] = &{};\n", gz_id, rust_byte_literal(&asset.gzip));
        code += "\n";
    }

    // Uncompressed lookup
    code += "/// Get an uncompressed frontend asset by path, returning (data, etag).\n";
    code += "pub fn get_frontend_asset_with_etag(path: &str) -> Option<(&'static [u8], &'static str)> {\n";
    code += "    match path {\n";
    for (path, _asset) in assets {
        let ident = path_to_ident(path);
        let etag_id = etag_ident(path);
        code += &format!("        \"{path}\" => Some((&{ident}, &{etag_id})),\n");
    }
    code += "        _ => None,\n";
    code += "    }\n";
    code += "}\n\n";

    // Compressed lookup (prefers brotli, falls back to gzip)
    code += "/// Get a pre-compressed frontend asset by path.\n";
    code += "/// Returns (data, encoding, etag) — prefers brotli over gzip.\n";
    code += "pub fn get_frontend_asset_compressed(path: &str) -> Option<(&'static [u8], &'static str, &'static str)> {\n";
    code += "    match path {\n";
    for (path, asset) in assets {
        let etag_id = etag_ident(path);
        if asset.brotli.is_some() {
            let br_id = br_ident(path);
            code += &format!("        \"{path}\" => Some((&{br_id}, \"br\", &{etag_id})),\n");
        } else {
            let gz_id = gz_ident(path);
            code += &format!("        \"{path}\" => Some((&{gz_id}, \"gzip\", &{etag_id})),\n");
        }
    }
    code += "        _ => None,\n";
    code += "    }\n";
    code += "}\n";

    let out_path = out_dir.join("frontend.rs");
    fs::write(&out_path, code).unwrap_or_else(|e| {
        panic!("Failed to write generated frontend.rs: {}", e)
    });
}

fn main() {
    let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
    let workspace_root = Path::new(&manifest_dir)
        .parent()
        .unwrap()
        .parent()
        .unwrap();
    let dist_dir = workspace_root.join("dist");

    if !dist_dir.exists() {
        println!("cargo:warning=dist/ directory not found — frontend assets will not be embedded");
        let out_dir = env::var("OUT_DIR").unwrap();
        let out_path = Path::new(&out_dir).join("frontend.rs");
        fs::write(
            &out_path,
            "//! No dist/ directory found — frontend assets not embedded.\n\
             pub fn get_frontend_asset_with_etag(_path: &str) -> Option<(&'static [u8], &'static str)> { None }\n\
             pub fn get_frontend_asset_compressed(_path: &str) -> Option<(&'static [u8], &'static str, &'static str)> { None }\n",
        ).unwrap();
        return;
    }

    println!("cargo:rerun-if-changed=dist/");

    let assets = collect_assets(&dist_dir);
    println!("cargo:warning=Collected {} frontend assets from dist/", assets.len());

    let out_dir = env::var("OUT_DIR").unwrap();
    generate_frontend_module(&assets, Path::new(&out_dir));
}