use base64::{engine::general_purpose::STANDARD, Engine};
use bytes::Bytes;
use regex::Regex;
use serde_json::{json, Map, Value};
use sha2::{Digest, Sha256};
use std::{
    collections::{BTreeMap, BTreeSet},
    fs,
    path::Path,
};

#[derive(Clone, Debug, serde::Serialize)]
pub struct AssetFile {
    pub h: String,
    pub s: usize,
    pub g: String,
}

#[derive(Clone)]
pub struct AssetSnapshot {
    pub build_tag: String,
    pub manifest: Value,
    pub delivery_files: BTreeMap<String, AssetFile>,
    pub bytes: BTreeMap<String, Bytes>,
    pub etags: BTreeMap<String, String>,
}

fn hash10(bytes: &[u8]) -> String {
    hex::encode(Sha256::digest(bytes))[..10].to_string()
}
fn hash12(bytes: &[u8]) -> String {
    hex::encode(Sha256::digest(bytes))[..12].to_string()
}
fn etag(bytes: &[u8]) -> String {
    format!("\"{}\"", &hex::encode(Sha256::digest(bytes))[..32])
}
fn collect(dir: &Path, prefix: &str, out: &mut BTreeMap<String, Vec<u8>>) -> std::io::Result<()> {
    let mut entries = fs::read_dir(dir)?.collect::<Result<Vec<_>, _>>()?;
    entries.sort_by_key(|e| e.file_name());
    for entry in entries {
        let path = entry.path();
        let name = entry.file_name().to_string_lossy().to_string();
        let url = format!("{prefix}/{name}");
        if path.is_dir() {
            collect(&path, &url, out)?;
        } else {
            out.insert(url, fs::read(path)?);
        }
    }
    Ok(())
}
fn runtime(path: &str, include: &BTreeSet<String>) -> bool {
    if include.contains(path) {
        return true;
    }
    path != "/admin.html"
        && !path.to_ascii_lowercase().split('/').any(|p| p == "perf")
        && ![".test.js", ".test.mjs", ".md", ".txt", ".c", ".entry.js"]
            .iter()
            .any(|s| path.to_ascii_lowercase().ends_with(s))
}
fn embed(
    source: &str,
    files: &BTreeMap<String, AssetFile>,
    groups: &Value,
    tag: &str,
    asset_sync: bool,
    hashing: bool,
) -> String {
    let mut page_files = Map::new();
    if let Some(groups_obj) = groups.as_object() {
        for group in groups_obj.values() {
            if let Some(paths) = group.get("files").and_then(Value::as_array) {
                for p in paths.iter().filter_map(Value::as_str) {
                    if p == "/index.html" || p == "/sw.js" {
                        continue;
                    }
                    if let Some(file) = files.get(p) {
                        page_files.insert(p.to_string(), json!({"h":if hashing { file.h.as_str() } else { tag },"s":file.s,"g":file.g}));
                    }
                }
            }
        }
    }
    let assets = json!({"buildTag":tag,"groups":groups,"files":page_files,"appCache":asset_sync})
        .to_string()
        .replace('<', "\\u003c");
    let re = Regex::new(r#"(?s)(<script type="application/json" id="ytp-assets">).*?(</script>)"#)
        .unwrap();
    re.replace(source, format!("$1{assets}$2")).into_owned()
}
fn stamp(source: &str, files: &BTreeMap<String, AssetFile>, tag: &str, hashing: bool) -> String {
    if !hashing {
        let tag_re = Regex::new(r"(?is)<(?:script|link)\b[^>]*>").unwrap();
        let attr_re = Regex::new(r#"(?i)\b(href|src)=["']([^"']+\.(?:css|js))["']"#).unwrap();
        return tag_re
            .replace_all(
                &source.replace("__BUILD_TAG__", tag),
                |caps: &regex::Captures| {
                    attr_re
                        .replace_all(&caps[0], |a: &regex::Captures| {
                            let raw = &a[2];
                            if raw.starts_with('/') || raw.starts_with("//") || raw.contains(':') {
                                return a[0].to_string();
                            }
                            let separator = if raw.contains('?') { '&' } else { '?' };
                            format!("{}=\"{}{}v={}\"", &a[1], raw, separator, tag)
                        })
                        .into_owned()
                },
            )
            .into_owned();
    }
    let tag_re = Regex::new(r"(?is)<(?:script|link)\b[^>]*>").unwrap();
    let attr_re = Regex::new(r#"\b(href|src)=(?:"([^"]+)"|'([^']+)')"#).unwrap();
    let stamped = tag_re
        .replace_all(source, |caps: &regex::Captures| {
            let tag_text = caps.get(0).unwrap().as_str();
            attr_re
                .replace_all(tag_text, |a: &regex::Captures| {
                    let raw = a.get(2).or_else(|| a.get(3)).unwrap().as_str();
                    if raw.contains(":") || raw.starts_with("//") || raw.starts_with('#') {
                        return a[0].to_string();
                    }
                    let url = url::Url::parse(&format!("https://assets.invalid/{raw}"));
                    let Ok(parsed) = url else {
                        return a[0].to_string();
                    };
                    let path = parsed.path();
                    let Some(file) = files.get(path) else {
                        return a[0].to_string();
                    };
                    if !(path.ends_with(".css") || path.ends_with(".js")) || path == "/sw.js" {
                        return a[0].to_string();
                    }
                    let base = raw.split(['?', '#']).next().unwrap_or(raw);
                    let suffix = if let Some((_, query)) = raw.split_once('?') {
                        let rest = query
                            .split('#')
                            .next()
                            .unwrap_or("")
                            .split('&')
                            .filter(|p| !p.starts_with("v="))
                            .collect::<Vec<_>>()
                            .join("&");
                        if rest.is_empty() {
                            format!("?v={}", file.h)
                        } else {
                            format!("?v={}&{}", file.h, rest)
                        }
                    } else {
                        format!("?v={}", file.h)
                    };
                    let frag = raw
                        .split_once('#')
                        .map(|(_, f)| format!("#{f}"))
                        .unwrap_or_default();
                    let quote = if a.get(2).is_some() { "\"" } else { "'" };
                    format!("{}={}{}{}{}{}", &a[1], quote, base, suffix, frag, quote)
                })
                .into_owned()
        })
        .into_owned();
    stamped.replace("__BUILD_TAG__", tag)
}

pub fn load(
    root: &Path,
    hashing: bool,
    sync: bool,
    tag_override: Option<&str>,
) -> Result<AssetSnapshot, Box<dyn std::error::Error>> {
    let mut bytes = BTreeMap::new();
    collect(root, "", &mut bytes)?;
    let def: Value = bytes.get("/assets.json").map(|b| serde_json::from_slice(b)).transpose()?.unwrap_or_else(|| json!({"groups":{"core":{"contract":1,"eager":true,"files":bytes.keys().cloned().collect::<Vec<_>>()}}}));
    let include = def
        .get("include")
        .and_then(Value::as_array)
        .into_iter()
        .flatten()
        .filter_map(Value::as_str)
        .map(str::to_string)
        .collect::<BTreeSet<_>>();
    let mut groups = Map::new();
    let mut contracts = Map::new();
    let mut membership = BTreeMap::<String, String>::new();
    if let Some(group_obj) = def.get("groups").and_then(Value::as_object) {
        for name in group_obj.keys().cloned().collect::<BTreeSet<_>>() {
            let group = group_obj[&name].clone();
            let files = group
                .get("files")
                .and_then(Value::as_array)
                .ok_or("asset group files missing")?;
            let mut uniq = BTreeSet::new();
            for f in files.iter().filter_map(Value::as_str) {
                if !runtime(f, &include) {
                    return Err(format!("excluded runtime file in group {name}: {f}").into());
                }
                if !bytes.contains_key(f) {
                    return Err(format!("missing asset group file {f}").into());
                }
                if !uniq.insert(f) || membership.insert(f.into(), name.clone()).is_some() {
                    return Err(format!("duplicate asset membership: {f}").into());
                }
            }
            contracts.insert(
                name.clone(),
                group.get("contract").cloned().unwrap_or(json!(1)),
            );
            groups.insert(name, group);
        }
    }
    let groups = Value::Object(groups);
    let contracts = Value::Object(contracts);
    let mut files = BTreeMap::new();
    let mut delivery_files = BTreeMap::new();
    for (path, content) in &bytes {
        let f = AssetFile {
            h: hash10(content),
            s: content.len(),
            g: membership
                .get(path)
                .cloned()
                .unwrap_or_else(|| "core".into()),
        };
        delivery_files.insert(path.clone(), f.clone());
        if runtime(path, &include) {
            files.insert(path.clone(), f);
        }
    }
    let mut source = String::from_utf8(bytes.get("/index.html").cloned().unwrap_or_default())?;
    if let Some(app) = bytes.get("/app.js") {
        source = source.replace("__APP_SCRIPT_HASH__", &STANDARD.encode(Sha256::digest(app)));
    }
    let sync = hashing && sync;
    let embedded = embed(&source, &files, &groups, "__BUILD_TAG__", sync, hashing);
    let canonical_index = stamp(&embedded, &files, "__BUILD_TAG__", hashing);
    let mut canonical_files = Map::new();
    for (path, file) in &files {
        let mut obj = Map::new();
        obj.insert("h".into(), json!(file.h));
        obj.insert("s".into(), json!(file.s));
        obj.insert("g".into(), json!(file.g));
        canonical_files.insert(path.clone(), Value::Object(obj));
    }
    if hashing && files.contains_key("/index.html") {
        let original = &files["/index.html"];
        let mut obj = Map::new();
        obj.insert("h".into(), json!(hash10(canonical_index.as_bytes())));
        obj.insert("s".into(), json!(canonical_index.len()));
        obj.insert("g".into(), json!(original.g));
        obj.insert("source".into(), json!(original.h));
        canonical_files.insert("/index.html".into(), Value::Object(obj));
    }
    let canonical =
        json!({"assetSync":sync,"files":canonical_files,"groups":groups,"contracts":contracts});
    let mut legacy_hasher = Sha256::new();
    for (path, content) in &bytes {
        if runtime(path, &include) {
            legacy_hasher.update(format!("./public{path}").as_bytes());
            legacy_hasher.update(content);
        }
    }
    let legacy_tag = hex::encode(legacy_hasher.finalize())[..12].to_string();
    let build_tag = if hashing {
        hash12(serde_json::to_string(&canonical)?.as_bytes())
    } else {
        tag_override.unwrap_or(&legacy_tag).to_string()
    };
    let index = embed(
        &stamp(&source, &files, &build_tag, hashing),
        &files,
        &groups,
        &build_tag,
        sync,
        hashing,
    );
    let mut sw =
        String::from_utf8(bytes.get("/sw.js").cloned().unwrap_or_default()).unwrap_or_default();
    let inject =
        Regex::new(r"typeof __BUILD_TAG__ !== 'undefined' \? __BUILD_TAG__ : '[^']*'").unwrap();
    sw = inject.replace(&sw, format!("\"{build_tag}\"")).into_owned();
    sw = sw.replace(
        "typeof __ASSET_SYNC__ !== 'undefined' ? __ASSET_SYNC__ : true",
        if sync { "true" } else { "false" },
    );
    if let Some(core) = files.get("/asset-sync-core.js") {
        sw = sw.replace(
            "importScripts('/asset-sync-core.js')",
            &format!("importScripts('/asset-sync-core.js?v={}')", core.h),
        );
    }
    if !index.is_empty() {
        files.insert(
            "/index.html".into(),
            AssetFile {
                h: hash10(index.as_bytes()),
                s: index.len(),
                g: "core".into(),
            },
        );
    }
    if !sw.is_empty() {
        files.insert(
            "/sw.js".into(),
            AssetFile {
                h: hash10(sw.as_bytes()),
                s: sw.len(),
                g: "core".into(),
            },
        );
    }
    if !index.is_empty() {
        bytes.insert("/index.html".into(), index.as_bytes().to_vec());
    }
    if !sw.is_empty() {
        bytes.insert("/sw.js".into(), sw.as_bytes().to_vec());
    }
    delivery_files.extend(files.clone());
    let manifest =
        json!({"buildTag":build_tag,"files":files,"groups":groups,"contracts":contracts});
    let etags = bytes
        .iter()
        .map(|(path, content)| (path.clone(), etag(content)))
        .collect();
    let bytes = bytes
        .into_iter()
        .map(|(path, content)| (path, Bytes::from(content)))
        .collect();
    Ok(AssetSnapshot {
        build_tag,
        manifest,
        delivery_files,
        bytes,
        etags,
    })
}

pub fn cache_control(
    path: &str,
    version: Option<&str>,
    current: Option<&AssetFile>,
    hashing: bool,
    build_tag: &str,
) -> &'static str {
    if !hashing && path == "/index.html" {
        return "no-cache";
    }
    if let Some(version) = version {
        if current.map(|f| f.h.as_str()) == Some(version) || (!hashing && version == build_tag) {
            return "public, max-age=31536000, immutable";
        }
        return if hashing { "no-store" } else { "no-cache" };
    }
    if path.starts_with("/fonts/") || path.starts_with("/icons/") {
        "public, max-age=2592000"
    } else {
        "no-cache"
    }
}

#[cfg(test)]
mod tests {
    use super::{cache_control, embed, stamp, AssetFile};
    use std::{collections::BTreeMap, path::Path};

    #[test]
    fn repository_frontend_snapshot_has_hashed_assets_and_injected_tags() {
        let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../frontend");
        let snapshot = super::load(&root, true, true, None).expect("frontend asset manifest");
        assert_eq!(snapshot.build_tag.len(), 12);
        assert!(
            snapshot.manifest["files"]["/app.js"]["h"]
                .as_str()
                .unwrap()
                .len()
                == 10
        );
        assert!(snapshot.bytes["/index.html"]
            .windows(snapshot.build_tag.len())
            .any(|w| w == snapshot.build_tag.as_bytes()));
        let sw = String::from_utf8(snapshot.bytes["/sw.js"].to_vec()).unwrap();
        assert!(sw.contains(&snapshot.build_tag));
        let index = String::from_utf8(snapshot.bytes["/index.html"].to_vec()).unwrap();
        assert!(index.contains(&format!(
            "styles.css?v={}",
            snapshot.manifest["files"]["/styles.css"]["h"]
                .as_str()
                .unwrap()
        )));
        assert!(index.contains(&format!(
            r#""/app.js":{{"h":"{}""#,
            snapshot.manifest["files"]["/app.js"]["h"].as_str().unwrap()
        )));
    }

    #[test]
    fn rollback_snapshot_uses_legacy_single_tag() {
        let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../frontend");
        let snapshot =
            super::load(&root, false, false, Some("legacy-build")).expect("rollback snapshot");
        assert_eq!(snapshot.build_tag, "legacy-build");
        let index = String::from_utf8(snapshot.bytes["/index.html"].to_vec()).unwrap();
        assert!(index.contains("styles.css?v=legacy-build"));
        assert!(index.contains(r#""/app.js":{"h":"legacy-build""#));
        assert!(index.contains(r#""appCache":false"#));
    }

    #[test]
    fn legacy_stamping_only_changes_relative_urls_in_resource_tags() {
        let input = r#"<link rel="stylesheet" href="styles.css"><script src="/root.js"></script><script src="https://cdn.test/app.js"></script><a href="other.js">link</a>"#;
        let output = stamp(input, &BTreeMap::new(), "legacy-tag", false);
        assert!(output.contains(r#"href="styles.css?v=legacy-tag""#));
        assert!(output.contains(r#"src="/root.js""#));
        assert!(output.contains(r#"src="https://cdn.test/app.js""#));
        assert!(output.contains(r#"href="other.js""#));
    }

    #[test]
    fn legacy_index_embeds_single_build_tag_for_each_asset() {
        let file = AssetFile {
            h: "file-hash".into(),
            s: 10,
            g: "core".into(),
        };
        let files = BTreeMap::from([("/app.js".to_string(), file)]);
        let groups = serde_json::json!({"core":{"files":["/app.js"]}});
        let source = r#"<script type="application/json" id="ytp-assets">{}</script>"#;
        let output = embed(source, &files, &groups, "legacy-tag", false, false);
        assert!(output.contains(r#""h":"legacy-tag""#));
        assert!(!output.contains("file-hash"));
    }

    #[test]
    fn cache_header_matrix_uses_legacy_build_tag_and_file_hashes() {
        let file = AssetFile {
            h: "abc123".into(),
            s: 10,
            g: "core".into(),
        };
        assert_eq!(
            cache_control(
                "/app.js",
                Some("legacy-tag"),
                Some(&file),
                false,
                "legacy-tag"
            ),
            "public, max-age=31536000, immutable"
        );
        assert_eq!(
            cache_control("/app.js", Some("old-tag"), Some(&file), false, "legacy-tag"),
            "no-cache"
        );
        assert_eq!(
            cache_control("/app.js", Some("abc123"), Some(&file), true, "build"),
            "public, max-age=31536000, immutable"
        );
        assert_eq!(
            cache_control("/app.js", Some("stale"), Some(&file), true, "build"),
            "no-store"
        );
    }
}