sandcage/src/docker.rs

#![allow(unused_assignments)]
use std::collections::HashMap;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::Command;

use miette::Diagnostic;
use sha2::{Digest, Sha256};
use thiserror::Error;

use crate::config::SandcageConfig;

// ---------------------------------------------------------------------------
// Bundled Dockerfiles
// ---------------------------------------------------------------------------

const DOCKERFILE_BASE: &str = include_str!("../images/base/Dockerfile");
const DOCKERFILE_CLAUDE: &str = include_str!("../images/claude/Dockerfile");
const DOCKERFILE_CODEX: &str = include_str!("../images/codex/Dockerfile");

const COMPOSE_YAML: &str = include_str!("../compose/docker-compose.yml");

#[derive(Debug, Error, Diagnostic)]
pub enum DockerError {
    #[error("docker executable not found in PATH")]
    #[diagnostic(
        code(sandcage::docker::not_found),
        help("Install Docker: https://docs.docker.com/get-docker/")
    )]
    DockerNotFound,

    #[error("docker compose plugin not available")]
    #[diagnostic(
        code(sandcage::docker::compose_not_found),
        help("Docker Compose V2 is required. Install it via Docker Desktop or 'docker compose' plugin.")
    )]
    ComposeNotFound,

    #[error("Failed to determine current user UID/GID: {0}")]
    #[diagnostic(code(sandcage::docker::id_failed))]
    IdFailed(String),

    #[error("Cannot determine home directory")]
    #[diagnostic(code(sandcage::docker::no_home_dir))]
    NoHomeDir,

    #[error("Failed to write temporary compose file: {0}")]
    #[diagnostic(code(sandcage::docker::tempfile_failed))]
    TempfileFailed(#[source] std::io::Error),

    #[error("Failed to spawn docker compose: {0}")]
    #[diagnostic(code(sandcage::docker::spawn_failed))]
    SpawnFailed(#[source] std::io::Error),

    #[error("Service '{service}' exited with status {code}")]
    #[diagnostic(code(sandcage::docker::service_failed))]
    ServiceFailed { service: String, code: i32 },

    #[error("Failed to create temporary build directory: {0}")]
    #[diagnostic(code(sandcage::docker::temp_dir_failed))]
    TempDirFailed(#[source] std::io::Error),

    #[error("Failed to write Dockerfile to temp directory: {0}")]
    #[diagnostic(code(sandcage::docker::dockerfile_write_failed))]
    DockerfileWriteFailed(#[source] std::io::Error),

    #[error("Failed to read build-hashes file: {0}")]
    #[diagnostic(code(sandcage::docker::hash_read_failed))]
    HashReadFailed(#[source] std::io::Error),

    #[error("Failed to write build-hashes file: {0}")]
    #[diagnostic(code(sandcage::docker::hash_write_failed))]
    HashWriteFailed(#[source] std::io::Error),

    #[error("docker build for '{image}' exited with status {code}")]
    #[diagnostic(code(sandcage::docker::build_failed))]
    BuildFailed { image: String, code: i32 },

    #[error("Image '{image}' not found locally")]
    #[diagnostic(
        code(sandcage::docker::image_not_found),
        help("Run `sandcage build` to build the required images.")
    )]
    ImageNotFound { image: String },
}

pub type Result<T> = std::result::Result<T, DockerError>;

fn require_docker() -> Result<PathBuf> {
    which::which("docker").map_err(|_| DockerError::DockerNotFound)
}

fn require_compose(docker: &Path) -> Result<()> {
    let output = Command::new(docker)
        .args(["compose", "version"])
        .output()
        .map_err(|_| DockerError::ComposeNotFound)?;

    if !output.status.success() {
        return Err(DockerError::ComposeNotFound);
    }
    Ok(())
}

fn id_flag(flag: &str) -> Result<String> {
    let output = Command::new("id")
        .arg(flag)
        .output()
        .map_err(|e| DockerError::IdFailed(e.to_string()))?;

    if !output.status.success() {
        return Err(DockerError::IdFailed(format!(
            "id {flag} exited with {}",
            output.status
        )));
    }

    Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
}

pub fn build_compose_env(workspace: &Path) -> Result<HashMap<String, String>> {
    let home = dirs::home_dir().ok_or(DockerError::NoHomeDir)?;
    let sandcage_home = home.join(".sandcage");

    let uid = id_flag("-u")?;
    let gid = id_flag("-g")?;

    let mut env = HashMap::new();
    env.insert("SANDCAGE_UID".into(), uid);
    env.insert("SANDCAGE_GID".into(), gid);
    env.insert(
        "SANDCAGE_WORKSPACE".into(),
        workspace.to_string_lossy().into_owned(),
    );
    env.insert(
        "SANDCAGE_HOME".into(),
        sandcage_home.to_string_lossy().into_owned(),
    );
    env.insert(
        "SANDCAGE_GLOBAL_JUSTFILE".into(),
        sandcage_home
            .join("Justfile")
            .to_string_lossy()
            .into_owned(),
    );

    Ok(env)
}

fn write_compose_tempfile() -> Result<tempfile::NamedTempFile> {
    let mut tmp = tempfile::Builder::new()
        .prefix("sandcage-compose-")
        .suffix(".yml")
        .tempfile()
        .map_err(DockerError::TempfileFailed)?;

    tmp.write_all(COMPOSE_YAML.as_bytes())
        .map_err(DockerError::TempfileFailed)?;
    tmp.flush().map_err(DockerError::TempfileFailed)?;

    Ok(tmp)
}

fn image_for_service(service: &str) -> &'static str {
    match service {
        "claude" => "sandcage-claude",
        "codex" => "sandcage-codex",
        _ => "sandcage-base",
    }
}

fn require_image(docker: &Path, image: &str) -> Result<()> {
    let output = Command::new(docker)
        .args(["image", "inspect", image])
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .status();

    match output {
        Ok(status) if status.success() => Ok(()),
        _ => Err(DockerError::ImageNotFound {
            image: image.to_string(),
        }),
    }
}

pub fn run_service(service: &str, workspace: &Path, config: &SandcageConfig) -> Result<()> {
    let docker = require_docker()?;
    require_compose(&docker)?;

    let image = image_for_service(service);
    require_image(&docker, image)?;

    let compose_file = write_compose_tempfile()?;
    let compose_path = compose_file.path().to_string_lossy().into_owned();

    let compose_env = build_compose_env(workspace)?;

    // Build the command
    let mut cmd = Command::new(&docker);
    cmd.args(["compose", "-f", &compose_path, "run", "--rm"]);

    // Add extra -e flags from the merged config env map
    if let Some(ref env_map) = config.env {
        for (key, value) in env_map {
            cmd.args(["-e", &format!("{key}={value}")]);
        }
    }

    cmd.arg(service);

    // Inject compose environment variables
    cmd.envs(&compose_env);

    // Inherit stdio for interactive use
    cmd.stdin(std::process::Stdio::inherit())
        .stdout(std::process::Stdio::inherit())
        .stderr(std::process::Stdio::inherit());

    let status = cmd.status().map_err(DockerError::SpawnFailed)?;

    if !status.success() {
        return Err(DockerError::ServiceFailed {
            service: service.to_string(),
            code: status.code().unwrap_or(-1),
        });
    }

    Ok(())
}

// ---------------------------------------------------------------------------
// Build with cache awareness
// ---------------------------------------------------------------------------

/// SHA-256 hex digest of the given string content.
fn sha256_hex(content: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(content.as_bytes());
    hex::encode(hasher.finalize())
}

/// Path to the build-hashes file: `~/.sandcage/.build-hashes`.
fn build_hashes_path() -> Result<PathBuf> {
    let home = dirs::home_dir().ok_or(DockerError::NoHomeDir)?;
    Ok(home.join(".sandcage").join(".build-hashes"))
}

/// Read the stored hashes from `~/.sandcage/.build-hashes`.
/// Returns an empty map if the file does not exist.
fn read_stored_hashes(path: &Path) -> Result<HashMap<String, String>> {
    if !path.exists() {
        return Ok(HashMap::new());
    }

    let content = std::fs::read_to_string(path).map_err(DockerError::HashReadFailed)?;
    let mut map = HashMap::new();
    for line in content.lines() {
        let line = line.trim();
        if line.is_empty() {
            continue;
        }
        if let Some((name, hash)) = line.split_once(':') {
            map.insert(name.to_string(), hash.to_string());
        }
    }
    Ok(map)
}

/// Persist the full hashes map back to the file.
fn write_stored_hashes(path: &Path, hashes: &HashMap<String, String>) -> Result<()> {
    // Ensure parent directory exists.
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent).map_err(DockerError::HashWriteFailed)?;
    }

    // Build deterministic output (sorted keys).
    let mut lines: Vec<String> = hashes
        .iter()
        .map(|(k, v)| format!("{k}:{v}"))
        .collect();
    lines.sort();
    let content = lines.join("\n") + "\n";

    std::fs::write(path, content).map_err(DockerError::HashWriteFailed)?;
    Ok(())
}

/// Build a single image from inline Dockerfile content (temp build context).
fn build_one_image_from_content(docker: &Path, image: &str, dockerfile_content: &str) -> Result<()> {
    let tmp_dir = tempfile::Builder::new()
        .prefix("sandcage-build-")
        .tempdir()
        .map_err(DockerError::TempDirFailed)?;

    let dockerfile_path = tmp_dir.path().join("Dockerfile");
    std::fs::write(&dockerfile_path, dockerfile_content)
        .map_err(DockerError::DockerfileWriteFailed)?;

    run_docker_build(docker, image, tmp_dir.path(), None)
}

/// Build a single image from a user-provided path.
/// If the path is a directory, it's used as the build context.
/// If the path is a file, it's used as the Dockerfile with a temp context.
fn build_one_image_from_path(docker: &Path, image: &str, override_path: &Path) -> Result<()> {
    if override_path.is_dir() {
        run_docker_build(docker, image, override_path, None)
    } else {
        let tmp_dir = tempfile::Builder::new()
            .prefix("sandcage-build-")
            .tempdir()
            .map_err(DockerError::TempDirFailed)?;

        run_docker_build(docker, image, tmp_dir.path(), Some(override_path))
    }
}

fn run_docker_build(
    docker: &Path,
    image: &str,
    context: &Path,
    dockerfile: Option<&Path>,
) -> Result<()> {
    let mut cmd = Command::new(docker);
    cmd.args(["build", "-t", &format!("{image}:latest")]);

    if let Some(df) = dockerfile {
        cmd.args(["-f", &df.to_string_lossy()]);
    }

    cmd.arg(context)
        .stdin(std::process::Stdio::inherit())
        .stdout(std::process::Stdio::inherit())
        .stderr(std::process::Stdio::inherit());

    let status = cmd.status().map_err(DockerError::SpawnFailed)?;

    if !status.success() {
        return Err(DockerError::BuildFailed {
            image: image.to_string(),
            code: status.code().unwrap_or(-1),
        });
    }

    Ok(())
}

/// Resolve the Dockerfile content for hashing purposes.
/// Returns the content string for override files/dirs, or None if the path is invalid.
fn read_dockerfile_at(path: &Path) -> Result<String> {
    let dockerfile_path = if path.is_dir() {
        path.join("Dockerfile")
    } else {
        path.to_path_buf()
    };
    std::fs::read_to_string(&dockerfile_path).map_err(DockerError::HashReadFailed)
}

/// Build all three sandcage images with cache awareness.
///
/// * Resolves layered config to check for Dockerfile overrides.
/// * When `force` is `true`, every image is rebuilt regardless of the stored hash.
/// * Otherwise, an image is skipped when its Dockerfile hash matches the stored value.
/// * After a successful build the hash is updated in `~/.sandcage/.build-hashes`.
pub fn build_images(force: bool) -> Result<()> {
    let docker = require_docker()?;

    let home = dirs::home_dir().ok_or(DockerError::NoHomeDir)?;
    let global_config = home.join(".sandcage").join("config.toml");
    let project_dir = std::env::current_dir().unwrap_or_default();
    let project_config = project_dir.join(".sandcage.yml");

    // Global overrides are always trusted.
    let global_cfg = crate::config::resolve_config(Some(&global_config), None)
        .unwrap_or_default();
    let mut overrides = global_cfg.dockerfiles.unwrap_or_default();

    // Project overrides require the project to be in trusted_projects.
    let project_cfg = if project_config.exists() {
        crate::config::load(&project_config).ok()
    } else {
        None
    };

    if let Some(ref pcfg) = project_cfg
        && pcfg.dockerfiles.is_some()
    {
        let trusted = crate::config::load_trusted_projects(&global_config);
        let is_trusted = trusted.iter().any(|p| {
            let canonical_trusted = std::fs::canonicalize(p).unwrap_or_else(|_| p.clone());
            let canonical_project = std::fs::canonicalize(&project_dir)
                .unwrap_or_else(|_| project_dir.clone());
            canonical_project.starts_with(canonical_trusted)
        });

        if is_trusted {
            overrides.extend(pcfg.dockerfiles.clone().unwrap_or_default());
        } else {
            eprintln!(
                "sandcage: warning: project Dockerfile overrides ignored (not in trusted_projects)"
            );
            eprintln!(
                "sandcage: add this project to trusted_projects in ~/.sandcage/config.toml"
            );
        }
    }


    let images: &[(&str, &str)] = &[
        ("sandcage-base", DOCKERFILE_BASE),
        ("sandcage-claude", DOCKERFILE_CLAUDE),
        ("sandcage-codex", DOCKERFILE_CODEX),
    ];

    let hashes_path = build_hashes_path()?;
    let mut stored = read_stored_hashes(&hashes_path)?;

    for (image, bundled_dockerfile) in images {
        let short_name = image.strip_prefix("sandcage-").unwrap_or(image);
        let override_path = overrides.get(short_name);

        let current_hash = if let Some(path) = override_path {
            sha256_hex(&read_dockerfile_at(path)?)
        } else {
            sha256_hex(bundled_dockerfile)
        };

        let needs_build = force || stored.get(*image) != Some(&current_hash);

        if needs_build {
            if force {
                eprintln!("{image}: rebuilding (--force)");
            } else if override_path.is_some() {
                eprintln!("{image}: rebuilding (custom Dockerfile changed)");
            } else {
                eprintln!("{image}: rebuilding (Dockerfile changed)");
            }

            if let Some(path) = override_path {
                build_one_image_from_path(&docker, image, path)?;
            } else {
                build_one_image_from_content(&docker, image, bundled_dockerfile)?;
            }

            stored.insert(image.to_string(), current_hash);
            write_stored_hashes(&hashes_path, &stored)?;
        } else {
            eprintln!("{image}: up to date");
        }
    }

    Ok(())
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[test]
    fn compose_yaml_is_bundled() {
        assert!(
            COMPOSE_YAML.contains("services:"),
            "bundled compose YAML should contain 'services:'"
        );
        assert!(
            COMPOSE_YAML.contains("claude:"),
            "bundled compose YAML should define 'claude' service"
        );
        assert!(
            COMPOSE_YAML.contains("codex:"),
            "bundled compose YAML should define 'codex' service"
        );
        assert!(
            COMPOSE_YAML.contains("shell:"),
            "bundled compose YAML should define 'shell' service"
        );
    }

    #[test]
    fn build_compose_env_contains_required_keys() {
        let workspace = PathBuf::from("/tmp/test-workspace");
        let env = build_compose_env(&workspace).expect("build_compose_env");

        assert!(env.contains_key("SANDCAGE_UID"), "missing SANDCAGE_UID");
        assert!(env.contains_key("SANDCAGE_GID"), "missing SANDCAGE_GID");
        assert!(env.contains_key("SANDCAGE_WORKSPACE"), "missing SANDCAGE_WORKSPACE");
        assert!(env.contains_key("SANDCAGE_HOME"), "missing SANDCAGE_HOME");
        assert!(
            env.contains_key("SANDCAGE_GLOBAL_JUSTFILE"),
            "missing SANDCAGE_GLOBAL_JUSTFILE"
        );
    }

    #[test]
    fn build_compose_env_workspace_matches() {
        let workspace = PathBuf::from("/my/project");
        let env = build_compose_env(&workspace).expect("build_compose_env");
        assert_eq!(env["SANDCAGE_WORKSPACE"], "/my/project");
    }

    #[test]
    fn build_compose_env_home_ends_with_sandcage() {
        let workspace = PathBuf::from("/tmp");
        let env = build_compose_env(&workspace).expect("build_compose_env");
        assert!(
            env["SANDCAGE_HOME"].ends_with(".sandcage"),
            "SANDCAGE_HOME should end with .sandcage, got: {}",
            env["SANDCAGE_HOME"]
        );
    }

    #[test]
    fn build_compose_env_justfile_is_under_home() {
        let workspace = PathBuf::from("/tmp");
        let env = build_compose_env(&workspace).expect("build_compose_env");
        assert!(
            env["SANDCAGE_GLOBAL_JUSTFILE"].starts_with(&env["SANDCAGE_HOME"]),
            "SANDCAGE_GLOBAL_JUSTFILE should be under SANDCAGE_HOME"
        );
        assert!(
            env["SANDCAGE_GLOBAL_JUSTFILE"].ends_with("Justfile"),
            "SANDCAGE_GLOBAL_JUSTFILE should end with Justfile"
        );
    }

    #[test]
    fn uid_gid_are_numeric() {
        let workspace = PathBuf::from("/tmp");
        let env = build_compose_env(&workspace).expect("build_compose_env");

        let uid: u32 = env["SANDCAGE_UID"]
            .parse()
            .expect("UID should be numeric");
        let gid: u32 = env["SANDCAGE_GID"]
            .parse()
            .expect("GID should be numeric");

        // Basic sanity — UIDs/GIDs are non-negative (they're u32)
        assert!(uid < 1_000_000, "UID seems unreasonably large: {uid}");
        assert!(gid < 1_000_000, "GID seems unreasonably large: {gid}");
    }

    #[test]
    fn write_compose_tempfile_creates_valid_yaml() {
        let tmp = write_compose_tempfile().expect("write_compose_tempfile");
        let content = std::fs::read_to_string(tmp.path()).expect("read tempfile");
        assert_eq!(content, COMPOSE_YAML, "tempfile content should match bundled YAML");
    }

    #[test]
    fn require_docker_finds_docker_or_fails_cleanly() {
        // This test documents the behavior — it either finds docker or returns
        // DockerNotFound. We don't assert which, since CI may or may not have docker.
        let result = require_docker();
        match result {
            Ok(path) => assert!(path.to_string_lossy().contains("docker")),
            Err(DockerError::DockerNotFound) => { /* expected on machines without docker */ }
            Err(e) => panic!("unexpected error: {e}"),
        }
    }

    // -----------------------------------------------------------------------
    // build_images / cache-awareness tests
    // -----------------------------------------------------------------------

    #[test]
    fn dockerfiles_are_bundled() {
        assert!(!DOCKERFILE_BASE.is_empty(), "base Dockerfile should not be empty");
        assert!(!DOCKERFILE_CLAUDE.is_empty(), "claude Dockerfile should not be empty");
        assert!(!DOCKERFILE_CODEX.is_empty(), "codex Dockerfile should not be empty");
    }

    #[test]
    fn sha256_hex_is_deterministic() {
        let content = "FROM ubuntu:22.04\nRUN echo hello\n";
        let h1 = sha256_hex(content);
        let h2 = sha256_hex(content);
        assert_eq!(h1, h2, "sha256_hex must be deterministic");
        // SHA-256 produces 64 hex chars
        assert_eq!(h1.len(), 64, "sha256_hex should produce 64-char hex string");
    }

    #[test]
    fn sha256_hex_differs_for_different_input() {
        let h1 = sha256_hex("FROM ubuntu:22.04\n");
        let h2 = sha256_hex("FROM debian:bookworm\n");
        assert_ne!(h1, h2, "different content must yield different hash");
    }

    #[test]
    fn hash_file_roundtrip() {
        let tmp = tempfile::tempdir().expect("create temp dir");
        let path = tmp.path().join(".build-hashes");

        let mut hashes = HashMap::new();
        hashes.insert("sandcage-base".to_string(), "aabbcc".to_string());
        hashes.insert("sandcage-claude".to_string(), "ddeeff".to_string());

        write_stored_hashes(&path, &hashes).expect("write hashes");
        let loaded = read_stored_hashes(&path).expect("read hashes");

        assert_eq!(loaded.get("sandcage-base").map(String::as_str), Some("aabbcc"));
        assert_eq!(loaded.get("sandcage-claude").map(String::as_str), Some("ddeeff"));
    }

    #[test]
    fn read_stored_hashes_returns_empty_when_file_missing() {
        let tmp = tempfile::tempdir().expect("create temp dir");
        let path = tmp.path().join("nonexistent");

        let result = read_stored_hashes(&path).expect("read hashes on missing file");
        assert!(result.is_empty(), "should return empty map when file does not exist");
    }

    #[test]
    fn cache_hit_detection_same_hash() {
        let dockerfile = "FROM ubuntu:22.04\n";
        let hash = sha256_hex(dockerfile);

        let mut stored: HashMap<String, String> = HashMap::new();
        stored.insert("sandcage-base".to_string(), hash.clone());

        let needs_build = stored.get("sandcage-base").map_or(true, |h| h != &hash);
        assert!(!needs_build, "same hash should be a cache hit (no rebuild needed)");
    }

    #[test]
    fn cache_miss_detection_different_hash() {
        let dockerfile = "FROM ubuntu:22.04\n";
        let current_hash = sha256_hex(dockerfile);

        let mut stored: HashMap<String, String> = HashMap::new();
        stored.insert("sandcage-base".to_string(), "old_stale_hash".to_string());

        let needs_build = stored.get("sandcage-base").map_or(true, |h| h != &current_hash);
        assert!(needs_build, "different hash should be a cache miss (rebuild needed)");
    }

    #[test]
    fn cache_miss_detection_missing_entry() {
        let dockerfile = "FROM ubuntu:22.04\n";
        let current_hash = sha256_hex(dockerfile);
        let stored: HashMap<String, String> = HashMap::new();

        let needs_build = stored.get("sandcage-base").map_or(true, |h| h != &current_hash);
        assert!(needs_build, "missing entry should be treated as a cache miss");
    }

    #[test]
    fn force_flag_bypasses_cache() {
        // Even when hash matches, force=true means needs_build is true
        let dockerfile = "FROM ubuntu:22.04\n";
        let hash = sha256_hex(dockerfile);

        let mut stored: HashMap<String, String> = HashMap::new();
        stored.insert("sandcage-base".to_string(), hash.clone());

        let force = true;
        let needs_build = force || stored.get("sandcage-base").map_or(true, |h| h != &hash);
        assert!(needs_build, "force flag should always trigger a rebuild");
    }

    #[test]
    fn write_stored_hashes_creates_parent_dirs() {
        let tmp = tempfile::tempdir().expect("create temp dir");
        // Nested path that doesn't exist yet
        let path = tmp.path().join("a").join("b").join(".build-hashes");

        let mut hashes = HashMap::new();
        hashes.insert("sandcage-base".to_string(), "abc123".to_string());

        write_stored_hashes(&path, &hashes).expect("should create parent dirs and write");
        assert!(path.exists(), "hash file should have been created");
    }

    #[test]
    fn image_for_service_maps_correctly() {
        assert_eq!(image_for_service("claude"), "sandcage-claude");
        assert_eq!(image_for_service("codex"), "sandcage-codex");
        assert_eq!(image_for_service("shell"), "sandcage-base");
        assert_eq!(image_for_service("anything-else"), "sandcage-base");
    }

    #[test]
    fn require_image_fails_for_nonexistent_image() {
        let docker = match require_docker() {
            Ok(path) => path,
            Err(_) => return, // skip if docker not installed
        };
        let result = require_image(&docker, "sandcage-nonexistent-image-abc123");
        assert!(
            matches!(result, Err(DockerError::ImageNotFound { .. })),
            "should fail for nonexistent image: {result:?}"
        );
    }
}