//! Image build-dir resolution and reproducible checksums.
//!
//! `sm121 build` has to work from anywhere: a launcher workspace
//! (`docker-images/<dir>` + `docker_image_sources`), or a published
//! bundle clone (git.kilork.org/sm121/recipes, flat `<name>/` layout).
//! The checksum over the build dir is what gets stamped as the
//! `sm121.checksum` label and compared by `docker/stale.rs`.

use anyhow::{Context, Result};
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, HashMap};
use std::io::Read;
use std::path::{Path, PathBuf};

/// Compute a SHA-256 checksum over all files in `dir` (recursive).
/// Sorted by relative path to ensure reproducibility.
/// Flatten an image overlay into `dst` and rewrite its Dockerfile:
/// the `patches/` subtree plus every sub-directory the Dockerfile
/// COPY/ADD-references is flattened to the bundle root (gists cannot
/// carry directories).  Name collisions rename the flattened file to
/// `{dir}-{name}` and the COPY line is rewritten to the new name; a
/// collision under a glob source aborts the flatten for that directory
/// (it stays nested, `nested = true` — the gist flow refuses it).
/// Unreferenced sub-directories are always preserved nested.
pub fn flatten_overlay(src_dir: &Path, dst: &Path, dockerfile: &str) -> Result<(String, bool)> {
    use std::collections::BTreeSet;

    // 1) Sub-directories referenced by COPY/ADD sources (incl. continuations).
    let mut referenced: BTreeSet<String> = BTreeSet::new();
    let mut in_copy = false;
    for line in dockerfile.lines() {
        let t = line.trim_start();
        let is_copy = in_copy || t.starts_with("COPY ") || t.starts_with("ADD ");
        in_copy = is_copy && t.ends_with('\\');
        if !is_copy {
            continue;
        }
        let body = t.trim_end_matches('\\');
        let toks: Vec<&str> = body.split_whitespace().skip(1).collect();
        for tok in &toks[..toks.len().saturating_sub(1)] {
            if tok.starts_with('-') {
                continue; // --from=… etc.
            }
            let parts: Vec<&str> = tok.split('/').collect();
            // parts[0].is_empty() => absolute path (/tmp/…) — not a build-dir ref
            if parts.len() >= 2 && !parts[0].is_empty() && src_dir.join(parts[0]).is_dir() {
                referenced.insert(parts[0].to_string());
            }
        }
    }
    let flatten_dirs: BTreeSet<String> = referenced
        .iter()
        .cloned()
        .chain(std::iter::once("patches".to_string()))
        .collect();

    // 2) Root files first (they define the collision set), except the driver.
    for entry in std::fs::read_dir(src_dir)? {
        let entry = entry?;
        if entry.file_type()?.is_file() {
            let name = entry.file_name().to_string_lossy().to_string();
            if name != "build.sh" && name != "Dockerfile" {
                std::fs::copy(entry.path(), dst.join(&name))?;
            }
        }
    }

    // 3) Flatten the referenced dirs (+ patches/); everything else nests.
    let mut nested = false;
    let mut flattened: BTreeSet<String> = BTreeSet::new(); // dirs actually flattened
    let mut renames: BTreeMap<String, String> = BTreeMap::new(); // "dir/file" -> new name
    let mut globs: BTreeSet<String> = BTreeSet::new(); // dirs whose COPY source is a glob
    for line in dockerfile.lines() {
        let t = line.trim_start();
        if t.starts_with("COPY ") || t.starts_with("ADD ") {
            for tok in t.split_whitespace().skip(1) {
                if tok.starts_with('-') {
                    continue;
                }
                let parts: Vec<&str> = tok.split('/').collect();
                if parts.len() >= 2 && flatten_dirs.contains(parts[0]) && tok.contains('*') {
                    globs.insert(parts[0].to_string());
                }
            }
        }
    }
    for dir in flatten_dirs.iter() {
        let dir_path = src_dir.join(dir);
        if !dir_path.is_dir() {
            continue; // referenced but absent — COPY will fail in docker, not here
        }
        let files: Vec<PathBuf> = std::fs::read_dir(&dir_path)?
            .filter_map(|e| e.ok())
            .filter(|e| e.file_type().map(|t| t.is_file()).unwrap_or(false))
            .map(|e| e.path())
            .collect();
        let collides = files
            .iter()
            .any(|f| dst.join(f.file_name().unwrap()).exists());
        if collides && globs.contains(dir) {
            nested = true;
            std::fs::create_dir_all(dst.join(dir))?;
            copy_tree(&dir_path, &dst.join(dir))?; // rename under a glob is impossible
            continue;
        }
        for f in files {
            let name = f.file_name().unwrap().to_string_lossy().to_string();
            let key = format!("{dir}/{name}");
            let new = if dst.join(&name).exists() {
                let renamed = format!("{dir}-{name}");
                renames.insert(key, renamed.clone());
                renamed
            } else {
                name
            };
            std::fs::copy(f, dst.join(&new))?;
        }
        flattened.insert(dir.clone());
    }
    for entry in std::fs::read_dir(src_dir)? {
        let entry = entry?;
        let name = entry.file_name().to_string_lossy().to_string();
        if !entry.file_type()?.is_dir() || flatten_dirs.contains(&name) {
            continue;
        }
        // profiles/ is launcher metadata — its pins are already merged into
        // the bundle's versions.env, so it is not re-shipped.
        if name == "profiles" {
            continue;
        }
        nested = true;
        std::fs::create_dir_all(dst.join(&name))?;
        copy_tree(&entry.path(), &dst.join(&name))?;
    }

    // 4) Rewrite COPY/ADD sources for the flattened dirs (incl. continuations).
    let mut in_copy = false;
    let mut dst_done = false; // destination token already seen on an earlier line
    let rewritten = dockerfile
        .lines()
        .map(|line| {
            let t = line.trim_start();
            let is_copy = in_copy || t.starts_with("COPY ") || t.starts_with("ADD ");
            in_copy = is_copy && t.ends_with('\\');
            if !is_copy {
                return line.to_string();
            }
            let head = &line[..line.len() - t.len()]; // preserve indentation
            let is_first = t.starts_with("COPY ") || t.starts_with("ADD ");
            let toks: Vec<&str> = t.split_whitespace().collect();
            let n = toks.len();
            if is_first {
                // command + at least one src + a dest on this line -> all
                // continuation tokens are sources
                dst_done = toks[1..]
                    .iter()
                    .filter(|t| **t != "\\" && !t.starts_with('-'))
                    .count()
                    >= 2;
            }
            let rewritten_toks: Vec<String> = toks
                .iter()
                .enumerate()
                .map(|(i, tok)| {
                    let is_dest = !is_first && !dst_done && i == n.saturating_sub(1);
                    if (is_first && i == 0) || is_dest || tok.starts_with('-') {
                        return tok.to_string();
                    }
                    let parts: Vec<&str> = tok.split('/').collect();
                    if parts.len() < 2 || !flattened.contains(parts[0]) {
                        return tok.to_string();
                    }
                    if tok.contains('*') {
                        return parts[1..].join("/"); // glob: drop the dir prefix
                    }
                    let key = (*tok).to_string();
                    renames
                        .get(&key)
                        .cloned()
                        .unwrap_or_else(|| parts[1].to_string())
                })
                .collect();
            format!("{}{}", head, rewritten_toks.join(" "))
        })
        .collect::<Vec<_>>()
        .join("\n")
        + if dockerfile.ends_with('\n') { "\n" } else { "" };
    Ok((rewritten, nested))
}

fn copy_tree(src_dir: &Path, dst_dir: &Path) -> Result<()> {
    for entry in std::fs::read_dir(src_dir)? {
        let entry = entry?;
        let ft = entry.file_type()?;
        let target = dst_dir.join(entry.file_name());
        if ft.is_dir() {
            std::fs::create_dir_all(&target)?;
            copy_tree(&entry.path(), &target)?;
        } else if ft.is_file() {
            std::fs::copy(entry.path(), &target)?;
        }
    }
    Ok(())
}

pub fn compute_dir_checksum(dir: &Path) -> Result<String> {
    let mut hasher = Sha256::new();
    let mut entries: Vec<_> = Vec::new();

    for entry in walkdir::WalkDir::new(dir).sort_by_file_name() {
        let entry = entry?;
        if entry.file_type().is_dir() {
            continue;
        }
        let rel = entry
            .path()
            .strip_prefix(dir)
            .unwrap_or(entry.path())
            .to_string_lossy()
            .to_string();
        entries.push((rel, entry.path().to_path_buf()));
    }
    entries.sort_by(|a, b| a.0.cmp(&b.0));

    for (rel, path) in &entries {
        let mut file = std::fs::File::open(path)?;
        let mut buf = Vec::new();
        file.read_to_end(&mut buf)?;
        hasher.update(rel.as_bytes());
        hasher.update(b"\0");
        hasher.update(&buf);
    }

    let hash = hasher.finalize();
    Ok(hash.iter().map(|b| format!("{:02x}", b)).collect())
}

/// Parse a `versions.env` file into a sorted map. Only `KEY=VAL` lines are
/// captured; comments (`#`) and blank lines are dropped. The value is
/// everything after the first `=` (so values may contain `=`), with one
/// level of matching `"` or `'` quotes stripped — quoted and unquoted forms
/// of the same value must resolve (and checksum) identically.
pub fn parse_versions_env(raw: &str) -> BTreeMap<String, String> {
    let mut pins = BTreeMap::new();
    for line in raw.lines() {
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }
        if let Some((k, v)) = line.split_once('=') {
            let k = k.trim();
            if !k.is_empty() {
                pins.insert(k.to_string(), unquote(v.trim()));
            }
        }
    }
    pins
}

/// Strip one level of matching single/double quotes (build.sh does the
/// same; keeping the two parsers in agreement keeps checksums stable).
fn unquote(v: &str) -> String {
    let b = v.as_bytes();
    let quoted = b.len() >= 2
        && (b[0] == b'"' && *b.last().unwrap() == b'"'
            || b[0] == b'\'' && *b.last().unwrap() == b'\'');
    if quoted {
        v[1..v.len() - 1].to_string()
    } else {
        v.to_string()
    }
}

/// Serialize a pin map back to `versions.env` text (sorted by key; one
/// `KEY=VAL` per line). Comments are not preserved — this is a generated
/// artifact written to the remote build dir, not the source of truth.
pub fn serialize_versions_env(pins: &BTreeMap<String, String>) -> String {
    let mut out = String::from(
        "# Merged versions.env — generated by sm121 (profile + recipe [build.versions] applied)\n",
    );
    for (k, v) in pins {
        out.push_str(k);
        out.push('=');
        out.push_str(v);
        out.push('\n');
    }
    out
}

/// Compute the build checksum, folding in the version profile and recipe
/// `[build.versions]` overrides.
///
/// Precedence: image `versions.env` < profile file < recipe overrides.
/// Returns `(merged_versions_env, checksum)`:
/// - No profile and no overrides -> `(None, dir_checksum)`: the image's
///   `versions.env` is uploaded as-is and the checksum is unchanged
///   (existing images skip rebuild, backward-compatible).
/// - Profile and/or overrides -> `(Some(merged), combined)`: `merged` is the
///   image's `versions.env` with the profile pins and recipe pins applied
///   (recipe wins) for the launcher to write to the remote build dir;
///   `combined` hashes the dir checksum together with the merged env so any
///   pin change forces a rebuild. Profile files live inside the build dir, so
///   editing one changes the dir checksum too (double-protected).
pub(crate) fn build_checksum(
    local_dir: &Path,
    profile: Option<&str>,
    overrides: &HashMap<String, String>,
) -> Result<(Option<String>, String)> {
    let dir_checksum = compute_dir_checksum(local_dir)?;
    if profile.is_none() && overrides.is_empty() {
        return Ok((None, dir_checksum));
    }
    let env_path = local_dir.join("versions.env");
    let raw = std::fs::read_to_string(&env_path).context(format!(
        "a version profile / [build.versions] override requires '{}' to exist",
        env_path.display()
    ))?;
    let mut pins = parse_versions_env(&raw);
    if let Some(name) = profile {
        let profile_path = local_dir.join("profiles").join(format!("{name}.env"));
        let raw = std::fs::read_to_string(&profile_path).with_context(|| {
            let available = std::fs::read_dir(local_dir.join("profiles"))
                .map(|d| {
                    d.filter_map(|e| {
                        e.ok()
                            .filter(|e| e.path().extension().is_some_and(|x| x == "env"))
                            .map(|e| e.file_name().to_string_lossy().into_owned())
                    })
                    .collect::<Vec<_>>()
                    .join(", ")
                })
                .unwrap_or_default();
            format!(
                "build profile '{}' not found at {} (available: [{}])",
                name,
                profile_path.display(),
                available
            )
        })?;
        for (k, v) in parse_versions_env(&raw) {
            pins.insert(k, v);
        }
    }
    for (k, v) in overrides {
        pins.insert(k.clone(), v.clone());
    }
    let merged = serialize_versions_env(&pins);
    let mut h = Sha256::new();
    h.update(dir_checksum.as_bytes());
    h.update(b"\0");
    h.update(merged.as_bytes());
    let combined: String = h.finalize().iter().map(|b| format!("{:02x}", b)).collect();
    Ok((Some(merged), combined))
}

/// Where a recipe's image is built from.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BuildSource {
    /// Full build context: `build.sh` + `Dockerfile` + versions.env.
    Workspace,
    /// Published bundle (flattened; `build.sh`/`Dockerfile` at bundle root).
    Bundle,
}

/// Resolve a recipe's image build dir. Search order:
/// 1. `./docker-images/<dir>` (launcher workspace) — needs `build.sh`
/// 2. each `docker_image_sources` entry `S`: `S/docker-images/<dir>`, then
///    `S/<dir>` (a repo root containing `docker-images/`, or a dir of
///    image dirs)
/// 3. `./recipes/<name>/` then `./<name>/` carrying a bundled `build.sh`
///    (published bundles from the recipe registry are self-contained)
///
/// Errors name every path tried.
pub fn resolve_build_dir(
    base: &Path,
    recipe_name: &str,
    image_dir: &str,
    sources: &[String],
) -> Result<(std::path::PathBuf, BuildSource)> {
    let mut tried: Vec<String> = Vec::new();
    let workspace = base.join("docker-images").join(image_dir);
    if workspace.join("build.sh").is_file() {
        return Ok((workspace, BuildSource::Workspace));
    }
    tried.push(workspace.to_string_lossy().into());
    for src in sources {
        let root = crate::config::local_path(src);
        for cand in [
            root.join("docker-images").join(image_dir),
            root.join(image_dir),
        ] {
            if cand.join("build.sh").is_file() {
                return Ok((cand, BuildSource::Workspace));
            }
            tried.push(cand.to_string_lossy().into());
        }
    }
    for bun in [
        base.join("recipes").join(recipe_name),
        base.join(recipe_name),
    ] {
        if bun.join("build.sh").is_file() && bun.join("Dockerfile").is_file() {
            return Ok((bun, BuildSource::Bundle));
        }
        tried.push(bun.to_string_lossy().into());
    }
    anyhow::bail!(
        "no build context for image dir '{}' of recipe '{}'; searched (each needs build.sh; bundles also need a Dockerfile): {}",
        image_dir,
        recipe_name,
        tried.join(", ")
    )
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_versions_env_quotes_are_normalized() {
        let raw = "A=\"x:tag@sha256:abc\"\nB='y'\nC=z\nE=\"half\nF='weird\"\n";
        let pins = parse_versions_env(raw);
        assert_eq!(pins.get("A"), Some(&"x:tag@sha256:abc".to_string()));
        assert_eq!(pins.get("B"), Some(&"y".to_string()));
        assert_eq!(pins.get("C"), Some(&"z".to_string()));
        assert_eq!(pins.get("E"), Some(&"\"half".to_string()));
        assert_eq!(pins.get("F"), Some(&"'weird\"".to_string()));
    }

    #[test]
    fn parse_versions_env_skips_comments_and_blanks() {
        let raw = "# comment\n\nKEY=val\nOTHER=a=b\n  SPACED  =  x  \nnoeq\n";
        let pins = parse_versions_env(raw);
        assert_eq!(pins.get("KEY"), Some(&"val".to_string()));
        assert_eq!(pins.get("OTHER"), Some(&"a=b".to_string()));
        assert_eq!(pins.get("SPACED"), Some(&"x".to_string()));
        assert!(!pins.contains_key("noeq"));
        assert_eq!(pins.len(), 3);
    }

    #[test]
    fn resolve_falls_back_to_bundle_layout() {
        let base = std::env::temp_dir().join(format!("sm121-resolve-test-{}", std::process::id()));
        let bundle = base.join("recipes").join("test-resolve-recipe-x");
        std::fs::create_dir_all(&bundle).unwrap();
        std::fs::write(bundle.join("build.sh"), "").unwrap();
        std::fs::write(bundle.join("Dockerfile"), "").unwrap();
        let (dir, src) =
            resolve_build_dir(&base, "test-resolve-recipe-x", "test-resolve-recipe-x", &[])
                .unwrap();
        assert_eq!(src, BuildSource::Bundle);
        assert_eq!(dir, bundle);
        let _ = std::fs::remove_dir_all(&base);
    }
}