//! Explicit model setup shared by the GUI and headless CLI. These functions do
//! not configure an inference provider or download executable workers. Search
//! callers must never call provisioning implicitly.
use anyhow::{bail, Context, Result};
use directories::ProjectDirs;
use reqwest::blocking::Client;
use serde::Deserialize;
use sha2::{Digest, Sha256};
use std::{
    fs::{self, File},
    io::{BufReader, Read, Write},
    path::{Component, Path, PathBuf},
    sync::atomic::{AtomicU64, Ordering},
    time::Duration,
};

/// Exact bytes of the pinned text model plus its multimedia projection.
pub const TOTAL_BYTES: u64 = 864_676_480;
const ATTEMPTS: usize = 3;

#[derive(Deserialize)]
struct Manifest {
    model_id: String,
    export_repository: String,
    revision: String,
    files: Vec<Artifact>,
}

#[derive(Deserialize)]
struct Artifact {
    path: String,
    bytes: u64,
    sha256: String,
}

fn manifest() -> Result<Manifest> {
    let manifest: Manifest = serde_json::from_str(include_str!(
        "../../../tools/embeddinggemma2/model-manifest.json"
    ))?;
    if manifest.model_id != "google/embeddinggemma-2"
        || manifest.export_repository != "ggml-org/embeddinggemma-2-GGUF"
        || manifest.revision.len() != 40
        || !manifest
            .revision
            .bytes()
            .all(|byte| byte.is_ascii_hexdigit())
        || manifest.files.len() != 2
        || manifest
            .files
            .iter()
            .map(|artifact| artifact.bytes)
            .sum::<u64>()
            != TOTAL_BYTES
    {
        bail!("invalid bundled model manifest");
    }
    for artifact in &manifest.files {
        validate_artifact(artifact)?;
        if !matches!(
            artifact.path.as_str(),
            "embeddinggemma-2-Q8_0.gguf" | "mmproj-embeddinggemma-2-Q8_0.gguf"
        ) {
            bail!("bundled model manifest contains an unexpected artifact");
        }
    }
    if manifest.files[0].path == manifest.files[1].path {
        bail!("duplicate model manifest artifact");
    }
    Ok(manifest)
}

fn validate_artifact(artifact: &Artifact) -> Result<()> {
    let mut components = Path::new(&artifact.path).components();
    if !matches!(components.next(), Some(Component::Normal(_)))
        || components.next().is_some()
        || artifact.path.contains(['/', '\\'])
        || artifact.bytes == 0
        || artifact.sha256.len() != 64
        || !artifact.sha256.bytes().all(|byte| byte.is_ascii_hexdigit())
    {
        bail!("invalid artifact path, size, or digest");
    }
    Ok(())
}

/// Platform-local application data, consistently used by CLI and desktop.
/// This computes the path without creating files or accessing the network.
pub fn default_model_directory() -> Result<PathBuf> {
    let project = ProjectDirs::from("dev", "Searchi", "Searchi")
        .context("platform application data directory is unavailable")?;
    Ok(project
        .data_local_dir()
        .join("models")
        .join("embeddinggemma2"))
}

/// Cheap location selection for optional bundled assets. This checks regular
/// files and expected lengths only; `verify_cached` still verifies SHA-256
/// before setup considers the model usable. Partial bundles use managed cache.
pub fn has_model_files(directory: &Path) -> Result<bool> {
    model_files_present(directory, &manifest()?.files)
}

fn model_files_present(directory: &Path, artifacts: &[Artifact]) -> Result<bool> {
    for artifact in artifacts {
        let metadata = match fs::metadata(directory.join(&artifact.path)) {
            Ok(metadata) => metadata,
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
            Err(error) => return Err(error).context("inspect bundled model artifact"),
        };
        if !metadata.is_file() || metadata.len() != artifact.bytes {
            return Ok(false);
        }
    }
    Ok(true)
}

/// Offline validation includes file length and full SHA-256, not mere existence.
pub fn verify_cached(directory: &Path) -> Result<bool> {
    for artifact in manifest()?.files {
        if !verified(&directory.join(&artifact.path), &artifact)? {
            return Ok(false);
        }
    }
    Ok(true)
}

fn verified(path: &Path, artifact: &Artifact) -> Result<bool> {
    let file = match File::open(path) {
        Ok(file) => file,
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
        Err(error) => return Err(error).context("open cached model artifact"),
    };
    if file.metadata()?.len() != artifact.bytes {
        return Ok(false);
    }
    let mut reader = BufReader::new(file).take(artifact.bytes + 1);
    let mut hash = Sha256::new();
    let mut size = 0u64;
    let mut buffer = [0u8; 65536];
    loop {
        let read = reader.read(&mut buffer)?;
        if read == 0 {
            break;
        }
        size += read as u64;
        if size > artifact.bytes {
            return Ok(false);
        }
        hash.update(&buffer[..read]);
    }
    Ok(size == artifact.bytes
        && format!("{:x}", hash.finalize()).eq_ignore_ascii_case(&artifact.sha256))
}

/// Explicit blocking setup: pinned HTTPS artifacts, three attempts per missing
/// or corrupt artifact, streaming size/SHA checks, atomic replacement. Progress
/// is completed verified bytes plus bytes of the current attempt. On failure it
/// resets to verified completed bytes; cached files never trigger network access.
pub fn provision(directory: &Path, progress: &AtomicU64) -> Result<()> {
    let manifest = manifest()?;
    let client = Client::builder()
        .connect_timeout(Duration::from_secs(20))
        .timeout(Duration::from_secs(900))
        .user_agent(concat!("Searchi/", env!("CARGO_PKG_VERSION")))
        .build()?;
    provision_with(directory, progress, &manifest, &client, |artifact| {
        format!(
            "https://huggingface.co/{}/resolve/{}/{}",
            manifest.export_repository, manifest.revision, artifact.path
        )
    })
}

// URL injection is private so tests can use a local HTTP fixture. Production
// exposes no arbitrary download source or unverified executable provisioner.
fn provision_with(
    directory: &Path,
    progress: &AtomicU64,
    manifest: &Manifest,
    client: &Client,
    url: impl Fn(&Artifact) -> String,
) -> Result<()> {
    fs::create_dir_all(directory).context("create model cache directory")?;
    progress.store(0, Ordering::Release);
    let mut completed = 0;
    for artifact in &manifest.files {
        validate_artifact(artifact)?;
        let destination = directory.join(&artifact.path);
        if verified(&destination, artifact)? {
            completed += artifact.bytes;
            progress.store(completed, Ordering::Release);
            continue;
        }
        let mut last_error = None;
        for attempt in 0..ATTEMPTS {
            match download_one(
                &destination,
                artifact,
                progress,
                completed,
                client,
                &url(artifact),
            ) {
                Ok(()) => {
                    last_error = None;
                    break;
                }
                Err(error) => {
                    last_error = Some(error);
                    progress.store(completed, Ordering::Release);
                    if attempt + 1 < ATTEMPTS {
                        std::thread::sleep(Duration::from_millis(250 << attempt));
                    }
                }
            }
        }
        if let Some(error) = last_error {
            return Err(error).with_context(|| {
                format!(
                    "model artifact {} failed after {ATTEMPTS} attempts",
                    artifact.path
                )
            });
        }
        completed += artifact.bytes;
        progress.store(completed, Ordering::Release);
    }
    Ok(())
}

fn download_one(
    destination: &Path,
    artifact: &Artifact,
    progress: &AtomicU64,
    completed: u64,
    client: &Client,
    url: &str,
) -> Result<()> {
    let mut response = client
        .get(url)
        .send()
        .and_then(reqwest::blocking::Response::error_for_status)
        .map_err(reqwest::Error::without_url)
        .context("request model artifact")?;
    if response
        .content_length()
        .is_some_and(|size| size != artifact.bytes)
    {
        bail!("model artifact response length differs from pinned size");
    }
    let directory = destination
        .parent()
        .context("model artifact has no parent directory")?;
    let mut output = tempfile::Builder::new()
        .prefix(".searchi-model-")
        .suffix(".part")
        .tempfile_in(directory)?;
    let mut hash = Sha256::new();
    let mut downloaded = 0;
    let mut buffer = [0u8; 65536];
    loop {
        let read = response
            .read(&mut buffer)
            .context("read model artifact stream")?;
        if read == 0 {
            break;
        }
        downloaded += read as u64;
        if downloaded > artifact.bytes {
            bail!("model artifact exceeds pinned size");
        }
        output.write_all(&buffer[..read])?;
        hash.update(&buffer[..read]);
        progress.store(completed + downloaded, Ordering::Release);
    }
    if downloaded != artifact.bytes
        || !format!("{:x}", hash.finalize()).eq_ignore_ascii_case(&artifact.sha256)
    {
        bail!("model artifact size or SHA-256 mismatch");
    }
    output.as_file().sync_all()?;
    output
        .persist(destination)
        .map_err(|error| error.error)
        .context("atomically replace model artifact")?;
    #[cfg(unix)]
    File::open(directory)?.sync_all()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::{net::TcpListener, sync::atomic::AtomicUsize, thread, time::Instant};

    fn artifact(path: &str, bytes: &[u8]) -> Artifact {
        Artifact {
            path: path.into(),
            bytes: bytes.len() as u64,
            sha256: format!("{:x}", Sha256::digest(bytes)),
        }
    }

    fn test_manifest(files: Vec<Artifact>) -> Manifest {
        Manifest {
            model_id: "test".into(),
            export_repository: "test".into(),
            revision: "test".into(),
            files,
        }
    }

    fn client() -> Client {
        Client::builder()
            .no_proxy()
            .connect_timeout(Duration::from_secs(1))
            .timeout(Duration::from_secs(2))
            .build()
            .unwrap()
    }

    // Local HTTP only. The production API does not expose this override.
    fn server(responses: Vec<Vec<u8>>) -> (String, thread::JoinHandle<usize>) {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let url = format!("http://{}/fixture", listener.local_addr().unwrap());
        listener.set_nonblocking(true).unwrap();
        let handle = thread::spawn(move || {
            let mut count = 0;
            for response in responses {
                let deadline = Instant::now() + Duration::from_secs(5);
                let mut stream = loop {
                    match listener.accept() {
                        Ok((stream, _)) => break stream,
                        Err(error)
                            if error.kind() == std::io::ErrorKind::WouldBlock
                                && Instant::now() < deadline =>
                        {
                            thread::sleep(Duration::from_millis(2))
                        }
                        Err(error) => panic!("local fixture accept failed: {error}"),
                    }
                };
                stream.set_nonblocking(false).unwrap();
                stream
                    .set_read_timeout(Some(Duration::from_secs(1)))
                    .unwrap();
                let mut request = Vec::new();
                let mut buffer = [0u8; 1024];
                while !request.windows(4).any(|window| window == b"\r\n\r\n") {
                    let read = stream.read(&mut buffer).unwrap();
                    assert!(read > 0);
                    request.extend_from_slice(&buffer[..read]);
                    assert!(request.len() <= 8192);
                }
                stream.write_all(&response).unwrap();
                count += 1;
            }
            count
        });
        (url, handle)
    }

    fn response(body: &[u8], declared: Option<usize>) -> Vec<u8> {
        let mut response = b"HTTP/1.1 200 OK\r\nConnection: close\r\n".to_vec();
        if let Some(length) = declared {
            response.extend_from_slice(format!("Content-Length: {length}\r\n").as_bytes());
        }
        response.extend_from_slice(b"\r\n");
        response.extend_from_slice(body);
        response
    }

    fn assert_no_part_files(directory: &Path) {
        assert!(!fs::read_dir(directory).unwrap().any(|entry| entry
            .unwrap()
            .file_name()
            .to_string_lossy()
            .ends_with(".part")));
    }

    #[test]
    fn manifest_is_exact_pinned_google_text_and_media_contract() {
        let manifest = manifest().unwrap();
        assert_eq!(
            manifest.revision,
            "bfcd298762cc34d0357ece5ebdd31791a3a374d8"
        );
        assert_eq!(
            manifest
                .files
                .iter()
                .map(|artifact| artifact.bytes)
                .sum::<u64>(),
            TOTAL_BYTES
        );
        assert_eq!(
            manifest.files[0].sha256,
            "2188ac1deca4b77dffefd603c2776a9d76d9d74ec01841392982ebb840b09135"
        );
        assert_eq!(
            manifest.files[1].sha256,
            "c4a8a52691ecef40618438928bdf9e68379b854e24166f292592353db0aab64f"
        );
    }

    #[test]
    fn optional_bundle_requires_both_regular_files_at_expected_sizes() {
        let directory = tempfile::tempdir().unwrap();
        let artifacts = [artifact("one.gguf", b"abc"), artifact("two.gguf", b"xyz")];
        assert!(!model_files_present(directory.path(), &artifacts).unwrap());
        fs::write(directory.path().join("one.gguf"), b"abc").unwrap();
        assert!(!model_files_present(directory.path(), &artifacts).unwrap());
        fs::create_dir(directory.path().join("two.gguf")).unwrap();
        assert!(!model_files_present(directory.path(), &artifacts).unwrap());
        fs::remove_dir(directory.path().join("two.gguf")).unwrap();
        fs::write(directory.path().join("two.gguf"), b"xyz").unwrap();
        assert!(model_files_present(directory.path(), &artifacts).unwrap());
        fs::write(directory.path().join("one.gguf"), b"ab").unwrap();
        assert!(!model_files_present(directory.path(), &artifacts).unwrap());
    }

    #[test]
    fn bundled_file_presence_never_substitutes_for_sha_verification() {
        let directory = tempfile::tempdir().unwrap();
        let artifact = artifact("one.gguf", b"abc");
        let path = directory.path().join(&artifact.path);
        fs::write(&path, b"xyz").unwrap();
        assert!(model_files_present(directory.path(), std::slice::from_ref(&artifact)).unwrap());
        assert!(!verified(&path, &artifact).unwrap());
    }

    #[test]
    fn same_size_corruption_is_detected_and_missing_cache_is_offline() {
        let directory = tempfile::tempdir().unwrap();
        assert!(!verify_cached(directory.path()).unwrap());
        let artifact = artifact("model.gguf", b"abc");
        let path = directory.path().join(&artifact.path);
        fs::write(&path, b"abc").unwrap();
        assert!(verified(&path, &artifact).unwrap());
        fs::write(&path, b"xyz").unwrap();
        assert!(!verified(&path, &artifact).unwrap());
    }

    #[test]
    fn cached_valid_artifacts_do_not_request_a_url() {
        let directory = tempfile::tempdir().unwrap();
        fs::write(directory.path().join("one.gguf"), b"abc").unwrap();
        let manifest = test_manifest(vec![artifact("one.gguf", b"abc")]);
        let calls = AtomicUsize::new(0);
        let progress = AtomicU64::new(999);
        provision_with(directory.path(), &progress, &manifest, &client(), |_| {
            calls.fetch_add(1, Ordering::Relaxed);
            "http://127.0.0.1:1/unreachable".into()
        })
        .unwrap();
        assert_eq!(calls.load(Ordering::Relaxed), 0);
        assert_eq!(progress.load(Ordering::Acquire), 3);
    }

    #[test]
    fn interrupted_attempt_retries_and_atomically_replaces_corrupt_cache() {
        let directory = tempfile::tempdir().unwrap();
        fs::write(directory.path().join("one.gguf"), b"badbad").unwrap();
        fs::write(directory.path().join("two.gguf"), b"abcdef").unwrap();
        let manifest = test_manifest(vec![
            artifact("one.gguf", b"abcdef"),
            artifact("two.gguf", b"abcdef"),
        ]);
        let (url, server) = server(vec![
            response(b"abc", Some(6)),
            response(b"abcdef", Some(6)),
        ]);
        let progress = AtomicU64::new(0);
        provision_with(directory.path(), &progress, &manifest, &client(), |_| {
            url.clone()
        })
        .unwrap();
        assert_eq!(server.join().unwrap(), 2);
        assert_eq!(progress.load(Ordering::Acquire), 12);
        assert_eq!(
            fs::read(directory.path().join("one.gguf")).unwrap(),
            b"abcdef"
        );
        assert_no_part_files(directory.path());
    }

    #[test]
    fn oversized_stream_is_bounded_and_failure_cleans_partial_files() {
        let directory = tempfile::tempdir().unwrap();
        fs::write(directory.path().join("one.gguf"), b"old").unwrap();
        let manifest = test_manifest(vec![artifact("one.gguf", b"abc")]);
        let (url, server) = server(vec![response(b"abcde", None); ATTEMPTS]);
        let progress = AtomicU64::new(0);
        let error = provision_with(directory.path(), &progress, &manifest, &client(), |_| {
            url.clone()
        })
        .unwrap_err();
        assert!(format!("{error:#}").contains("exceeds pinned size"));
        assert_eq!(server.join().unwrap(), ATTEMPTS);
        assert_eq!(progress.load(Ordering::Acquire), 0);
        assert_eq!(fs::read(directory.path().join("one.gguf")).unwrap(), b"old");
        assert_no_part_files(directory.path());
    }

    #[test]
    fn wrong_digest_cannot_be_installed_even_with_correct_size() {
        let directory = tempfile::tempdir().unwrap();
        let manifest = test_manifest(vec![artifact("one.gguf", b"abc")]);
        let (url, server) = server(vec![response(b"xyz", Some(3)); ATTEMPTS]);
        let progress = AtomicU64::new(0);
        let error = provision_with(directory.path(), &progress, &manifest, &client(), |_| {
            url.clone()
        })
        .unwrap_err();
        assert!(format!("{error:#}").contains("SHA-256 mismatch"));
        assert_eq!(server.join().unwrap(), ATTEMPTS);
        assert!(!directory.path().join("one.gguf").exists());
        assert_no_part_files(directory.path());
    }

    #[test]
    fn completed_artifact_survives_later_failure_and_is_reused_on_retry() {
        let directory = tempfile::tempdir().unwrap();
        let manifest = test_manifest(vec![
            artifact("one.gguf", b"abc"),
            artifact("two.gguf", b"def"),
        ]);
        let mut responses = vec![response(b"abc", Some(3))];
        responses.extend(vec![response(b"bad", Some(3)); ATTEMPTS]);
        let (url, first_server) = server(responses);
        let progress = AtomicU64::new(0);
        assert!(
            provision_with(directory.path(), &progress, &manifest, &client(), |_| url
                .clone())
            .is_err()
        );
        assert_eq!(first_server.join().unwrap(), ATTEMPTS + 1);
        assert_eq!(progress.load(Ordering::Acquire), 3);
        assert_eq!(fs::read(directory.path().join("one.gguf")).unwrap(), b"abc");
        assert!(!directory.path().join("two.gguf").exists());
        let (url, second_server) = server(vec![response(b"def", Some(3))]);
        provision_with(
            directory.path(),
            &progress,
            &manifest,
            &client(),
            |artifact| {
                assert_eq!(artifact.path, "two.gguf");
                url.clone()
            },
        )
        .unwrap();
        assert_eq!(second_server.join().unwrap(), 1);
        assert_eq!(progress.load(Ordering::Acquire), 6);
        assert_no_part_files(directory.path());
    }

    #[test]
    fn unsafe_artifact_names_are_rejected_and_default_path_is_platform_local() {
        assert!(validate_artifact(&artifact("../outside", b"abc")).is_err());
        assert!(validate_artifact(&artifact("..\\outside", b"abc")).is_err());
        assert!(default_model_directory()
            .unwrap()
            .ends_with("models/embeddinggemma2"));
    }
}