//! GGUF models on this machine, and the downloads that produce them.
//!
//! The registry pattern again: one owner, one lock, so what is on disk and what
//! this server believes cannot come apart. Three things shape the design, all
//! consequences of a model file being gigabytes rather than kilobytes:
//!
//! **A download belongs to the model, not to whoever asked for it.** It is
//! keyed by the model it produces and lives here, so any device can watch it --
//! including one that did not start it. State in a per-connection channel would
//! not survive the phone locking its screen, which for an hour-long download is
//! the normal case.
//!
//! **Every run has an id, and its outcome outlives it.** Without those, "not
//! downloading" is three answers at once -- it finished, it never started, or a
//! different run finished while you were away.
//!
//! **Progress is measured, never estimated.** `total` is whatever
//! `Content-Length` said and nothing else; when the server does not send one it
//! stays `None` and the phone shows that it does not know.

use std::collections::HashMap;
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use anyhow::{Context, Result, bail};
use serde::Serialize;

use wg_app_link::private;

/// Identifies this client to HuggingFace. They ask for one, and a request
/// without it is more likely to be rate-limited.
const USER_AGENT: &str = concat!("ai-server/", env!("CARGO_PKG_VERSION"));

/// Read size per loop iteration. Big enough that the syscall overhead is
/// nothing against a multi-gigabyte file, small enough that a cancel is noticed
/// promptly -- the flag is only checked between chunks.
const CHUNK: usize = 256 * 1024;

/// A model file sitting on this machine, ready to run.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LocalModel {
    /// `owner/repo/file.gguf` -- the HuggingFace coordinates, which are already
    /// unique, so nothing has to invent an id.
    pub key: String,
    pub repo: String,
    pub file: String,
    pub bytes: u64,
}

/// What a run is doing, or did. Flat rather than a tagged enum carrying its
/// message, because the phone switches on this and a string it can compare is
/// easier to render than a variant it has to destructure.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum DownloadState {
    Running,
    /// Reading the finished file back to check it against the hash HuggingFace
    /// publishes. Its own state because it takes real time on a multi-gigabyte
    /// file and "still working" is the honest thing to show, rather than a bar
    /// sitting at 100% for half a minute.
    Verifying,
    Finished,
    Failed,
    Cancelled,
}

/// One download run, as the phone sees it.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DownloadStatus {
    pub key: String,
    /// Distinguishes this run from any earlier one for the same model, so a
    /// device that was watching run 3 can tell it is now looking at run 4.
    pub run: u64,
    pub repo: String,
    pub file: String,
    pub state: DownloadState,
    pub done: u64,
    /// What `Content-Length` said, or absent when the server did not say.
    /// Absent means "unknown", never "zero".
    #[serde(skip_serializing_if = "Option::is_none")]
    pub total: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
    pub started: f64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub finished: Option<f64>,
}

/// The mutable half of a run, behind one lock.
#[derive(Debug)]
struct Progress {
    state: DownloadState,
    done: u64,
    total: Option<u64>,
    error: Option<String>,
    started: f64,
    finished: Option<f64>,
}

/// A run, shared between the thread doing the work and everyone watching.
struct Run {
    id: u64,
    key: String,
    repo: String,
    file: String,
    progress: Mutex<Progress>,
    /// Set by [`ModelStore::cancel`]; the download loop checks it between
    /// chunks and stops, leaving the partial file for a later resume.
    cancel: AtomicBool,
}

impl Run {
    fn status(&self) -> DownloadStatus {
        let p = self.progress.lock().unwrap();
        DownloadStatus {
            key: self.key.clone(),
            run: self.id,
            repo: self.repo.clone(),
            file: self.file.clone(),
            state: p.state,
            done: p.done,
            total: p.total,
            error: p.error.clone(),
            started: p.started,
            finished: p.finished,
        }
    }

    fn finish(&self, state: DownloadState, error: Option<String>) {
        let mut p = self.progress.lock().unwrap();
        p.state = state;
        p.error = error;
        p.finished = Some(crate::session::now());
    }
}

/// Every model this machine has, and every download in flight or finished.
pub struct ModelStore {
    dir: PathBuf,
    /// Keyed by model key: one run per model at a time, and the last run for a
    /// model stays here after it ends so its outcome can still be read.
    runs: Mutex<HashMap<String, Arc<Run>>>,
    next_run: AtomicU64,
}

impl ModelStore {
    pub fn new(dir: PathBuf) -> Self {
        Self {
            dir,
            runs: Mutex::new(HashMap::new()),
            next_run: AtomicU64::new(1),
        }
    }

    /// Where a model's file lives, refusing anything that would escape the
    /// models directory.
    ///
    /// The repo and file come from a phone, so they are treated as hostile:
    /// every component must be an ordinary name. Rejecting rather than
    /// sanitising, since a silently rewritten path would download the right
    /// bytes to the wrong place.
    fn path_for(&self, repo: &str, file: &str) -> Result<PathBuf> {
        let mut path = self.dir.clone();
        for part in repo.split('/').chain(file.split('/')) {
            if part.is_empty() || part == "." || part == ".." || part.contains('\\') {
                bail!("\"{repo}/{file}\" is not a name this can store: \"{part}\"");
            }
            path.push(part);
        }
        Ok(path)
    }

    pub fn key_for(repo: &str, file: &str) -> String {
        format!("{repo}/{file}")
    }

    /// Every `.gguf` found under the models directory, newest first. Read from
    /// disk on each call rather than cached: a file deleted by hand should stop
    /// being offered.
    pub fn list(&self) -> Vec<LocalModel> {
        let mut found = Vec::new();
        collect(&self.dir, &self.dir, &mut found);
        found.sort_by(|a, b| a.key.cmp(&b.key));
        found
    }

    /// The status of every run this server remembers.
    pub fn downloads(&self) -> Vec<DownloadStatus> {
        let runs = self.runs.lock().unwrap();
        let mut all: Vec<_> = runs.values().map(|run| run.status()).collect();
        all.sort_by_key(|status| std::cmp::Reverse(status.run));
        all
    }

    /// Starts fetching `file` from `repo`, or returns the run already doing so.
    /// Idempotent on purpose: a phone that lost its connection will press the
    /// button again, and that must join the existing run rather than start a
    /// second one writing the same file.
    pub fn start(self: &Arc<Self>, repo: &str, file: &str) -> Result<DownloadStatus> {
        let key = Self::key_for(repo, file);
        let target = self.path_for(repo, file)?;
        if target.is_file() {
            bail!("{key} is already downloaded");
        }

        let mut runs = self.runs.lock().unwrap();
        if let Some(existing) = runs.get(&key)
            && existing.progress.lock().unwrap().state == DownloadState::Running
        {
            return Ok(existing.status());
        }

        let run = Arc::new(Run {
            id: self.next_run.fetch_add(1, Ordering::Relaxed),
            key: key.clone(),
            repo: repo.to_string(),
            file: file.to_string(),
            progress: Mutex::new(Progress {
                state: DownloadState::Running,
                done: 0,
                total: None,
                error: None,
                started: crate::session::now(),
                finished: None,
            }),
            cancel: AtomicBool::new(false),
        });
        runs.insert(key, Arc::clone(&run));
        let status = run.status();
        drop(runs);

        // A dedicated thread rather than the blocking pool: this holds its
        // thread for as long as the download takes, which is minutes to hours,
        // and the pool exists for short work.
        let store = Arc::clone(self);
        std::thread::spawn(move || {
            let outcome = store.fetch(&run, &target);
            match outcome {
                Ok(()) if run.cancel.load(Ordering::Relaxed) => {
                    run.finish(DownloadState::Cancelled, None);
                    tracing::info!("download {} cancelled", run.key);
                }
                Ok(()) => {
                    run.finish(DownloadState::Finished, None);
                    tracing::info!("download {} finished", run.key);
                }
                Err(err) => {
                    let message = format!("{err:#}");
                    tracing::warn!("download {} failed: {message}", run.key);
                    run.finish(DownloadState::Failed, Some(message));
                }
            }
        });
        Ok(status)
    }

    /// Asks a running download to stop. The partial file stays, so starting
    /// again resumes rather than refetching.
    pub fn cancel(&self, key: &str) -> Result<DownloadStatus> {
        let runs = self.runs.lock().unwrap();
        let Some(run) = runs.get(key) else {
            bail!("no download for {key}");
        };
        run.cancel.store(true, Ordering::Relaxed);
        Ok(run.status())
    }

    /// Removes a downloaded model, and any partial file for it.
    pub fn delete(&self, key: &str) -> Result<()> {
        let (repo, file) = key.rsplit_once('/').context("a key is repo/file")?;
        let target = self.path_for(repo, file)?;
        let partial = partial_of(&target);
        if !target.is_file() && !partial.is_file() {
            bail!("{key} is not downloaded");
        }
        for path in [&target, &partial] {
            if path.is_file() {
                std::fs::remove_file(path).with_context(|| format!("remove {}", path.display()))?;
            }
        }
        self.runs.lock().unwrap().remove(key);
        Ok(())
    }

    fn fetch(&self, run: &Run, target: &Path) -> Result<()> {
        let partial = partial_of(target);
        let identity = identity_of(target);
        if let Some(parent) = target.parent() {
            private::create_dir(parent)?;
        }

        // What we have, and what it was part of. A partial with no recorded
        // identity is not resumable -- it could be a fragment of any revision.
        let known = std::fs::read_to_string(&identity)
            .ok()
            .map(|s| s.trim().to_string());
        let have = match known {
            Some(_) => partial.metadata().map(|m| m.len()).unwrap_or(0),
            None => 0,
        };

        let url = format!(
            "https://huggingface.co/{}/resolve/main/{}",
            run.repo,
            run.file.replace(' ', "%20")
        );
        let (mut response, mut resumed) = request(&url, have)?;
        let mut etag = etag_of(&response);

        // HuggingFace's CDN ignores `If-Range` -- probed 2026-08-28: a
        // deliberately stale validator still answers 206 with the ranged bytes.
        // So the header cannot be relied on to restart us, and the check is done
        // here instead: if what arrived is not the revision our partial belongs
        // to, resuming would splice two files into something of exactly the
        // right length and the wrong contents.
        if resumed && etag.is_some() && etag != known {
            tracing::info!(
                "{} changed upstream since the partial was written -- starting again",
                run.key,
            );
            let (fresh, fresh_resumed) = request(&url, 0)?;
            response = fresh;
            resumed = fresh_resumed;
            etag = etag_of(&response);
        }

        // On a 206, Content-Length is the length of the *range*, not of the file
        // -- taken at face value it would fill the bar at 72 MB of a 234 MB
        // model. The whole size is the last field of Content-Range, which has
        // the further merit of not depending on where the range began.
        let total: Option<u64> = if resumed {
            response
                .headers()
                .get("content-range")
                .and_then(|v| v.to_str().ok())
                .and_then(|v| {
                    v.rsplit_once('/')
                        .map(|(_, whole)| whole.trim().to_string())
                })
                .and_then(|whole| whole.parse().ok())
        } else {
            response
                .headers()
                .get("content-length")
                .and_then(|v| v.to_str().ok()?.parse().ok())
        };
        let mut done = if resumed { have } else { 0 };
        {
            let mut p = run.progress.lock().unwrap();
            p.done = done;
            p.total = total;
        }

        // `truncate(false)` is the whole resume story: the file is opened to be
        // seeked into and appended to, and truncating would throw away exactly
        // the bytes the Range request just asked the server not to send again.
        // Stated rather than left to the default.
        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(false)
            .open(&partial)
            .with_context(|| format!("open {}", partial.display()))?;
        if resumed {
            file.seek(SeekFrom::Start(have))
                .context("seek to resume point")?;
        } else {
            file.set_len(0)
                .context("truncate a partial we cannot resume onto")?;
        }
        // Written before the body, so an interrupted download leaves a partial
        // that can still say which revision it belongs to.
        if let Some(etag) = &etag {
            std::fs::write(&identity, etag).ok();
        }

        let mut reader = response.body_mut().as_reader();
        let mut buffer = vec![0u8; CHUNK];
        loop {
            if run.cancel.load(Ordering::Relaxed) {
                file.flush().ok();
                return Ok(());
            }
            let read = reader
                .read(&mut buffer)
                .context("reading from HuggingFace")?;
            if read == 0 {
                break;
            }
            file.write_all(&buffer[..read])
                .context("writing the model file")?;
            done += read as u64;
            run.progress.lock().unwrap().done = done;
        }
        file.flush().context("flushing the model file")?;
        drop(file);

        // Checked before the rename, so a file that fails never gets the real
        // name and `list` never offers it. With the identity check above this
        // should not fire; it is here because a wrong model is the kind of
        // failure that surfaces as bad output rather than as an error.
        if let Some(expected) = published_sha256(&run.repo, &run.file) {
            run.progress.lock().unwrap().state = DownloadState::Verifying;
            let actual = sha256_of(&partial)?;
            if actual != expected {
                std::fs::remove_file(&partial).ok();
                std::fs::remove_file(&identity).ok();
                bail!(
                    "{} arrived corrupted -- HuggingFace publishes sha256 {expected}, what \
                     arrived hashes to {actual}. It has been deleted; downloading again \
                     starts clean.",
                    run.key,
                );
            }
        }

        // Renamed only once complete, so a file at its real name is always a
        // whole model -- `list` needs no other way to tell.
        std::fs::rename(&partial, target)
            .with_context(|| format!("finish {}", target.display()))?;
        std::fs::remove_file(&identity).ok();
        Ok(())
    }
}

/// The sha256 of a file, read in chunks -- these are gigabytes, and reading one
/// into memory to hash it would be the largest allocation this server makes.
fn sha256_of(path: &Path) -> Result<String> {
    use sha2::{Digest, Sha256};
    let mut file =
        std::fs::File::open(path).with_context(|| format!("reopen {}", path.display()))?;
    let mut hasher = Sha256::new();
    let mut buffer = vec![0u8; CHUNK];
    loop {
        let read = file.read(&mut buffer).context("reading back to verify")?;
        if read == 0 {
            break;
        }
        hasher.update(&buffer[..read]);
    }
    // Hex by hand, as `wg_app_link::enroll::token_hash_hex` also has to, since
    // this sha2 version's output type does not implement LowerHex.
    Ok(hasher
        .finalize()
        .iter()
        .map(|byte| format!("{byte:02x}"))
        .collect())
}

/// One GET, ranged when there is something to resume onto.
fn request(url: &str, from: u64) -> Result<(ureq::http::Response<ureq::Body>, bool)> {
    let mut get = ureq::get(url).header("User-Agent", USER_AGENT);
    if from > 0 {
        get = get.header("Range", &format!("bytes={from}-"));
    }
    let response = get.call().with_context(|| format!("GET {url}"))?;
    // Trust the status, not the request: a server that ignores Range answers
    // 200 with the whole file, and appending to that would corrupt it.
    let resumed = response.status() == 206;
    Ok((response, resumed))
}

fn etag_of(response: &ureq::http::Response<ureq::Body>) -> Option<String> {
    Some(
        response
            .headers()
            .get("etag")?
            .to_str()
            .ok()?
            .trim()
            .to_string(),
    )
}

/// `x.gguf` -> `x.gguf.part.etag`, holding which revision the partial beside it
/// is a piece of.
fn identity_of(target: &Path) -> PathBuf {
    let mut name = target.as_os_str().to_os_string();
    name.push(".part.etag");
    PathBuf::from(name)
}

/// `x.gguf` -> `x.gguf.part`, the in-progress name.
fn partial_of(target: &Path) -> PathBuf {
    let mut name = target.as_os_str().to_os_string();
    name.push(".part");
    PathBuf::from(name)
}

/// Walks `dir` collecting `.gguf` files, keyed by their path under `root`.
fn collect(root: &Path, dir: &Path, found: &mut Vec<LocalModel>) {
    let Ok(entries) = std::fs::read_dir(dir) else {
        return;
    };
    for entry in entries.flatten() {
        let path = entry.path();
        if path.is_dir() {
            collect(root, &path, found);
            continue;
        }
        if path.extension().is_none_or(|e| e != "gguf") {
            continue;
        }
        let Ok(relative) = path.strip_prefix(root) else {
            continue;
        };
        let key = relative.to_string_lossy().replace('\\', "/");
        let Some((repo, file)) = key.rsplit_once('/') else {
            continue;
        };
        found.push(LocalModel {
            key: key.clone(),
            repo: repo.to_string(),
            file: file.to_string(),
            bytes: entry.metadata().map(|m| m.len()).unwrap_or(0),
        });
    }
}

/// A model repository on HuggingFace, as the browse screen shows it.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RemoteRepo {
    /// `owner/name`, which is what everything else here is keyed by.
    pub id: String,
    pub downloads: u64,
    pub likes: u64,
}

/// One downloadable file inside a repository.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RemoteFile {
    pub path: String,
    pub bytes: u64,
    /// Already on this machine, so the phone can say so rather than offering to
    /// fetch it again.
    pub have: bool,
}

/// Searches HuggingFace for GGUF repositories matching `query`.
///
/// Proxied through this server rather than called from the phone, for two
/// reasons that both matter: the app trusts exactly one certificate -- this
/// server's -- and has no general internet trust to spend on huggingface.co,
/// and the machine that has to do the downloading is this one.
pub fn search(query: &str) -> Result<Vec<RemoteRepo>> {
    let url = format!(
        "https://huggingface.co/api/models?search={}&filter=gguf&limit=25&sort=downloads&direction=-1",
        urlencode(query)
    );
    let body = get_json(&url)?;
    let list = body
        .as_array()
        .context("HuggingFace returned something that is not a list")?;
    Ok(list
        .iter()
        .filter_map(|m| {
            Some(RemoteRepo {
                id: m.get("id")?.as_str()?.to_string(),
                downloads: m
                    .get("downloads")
                    .and_then(serde_json::Value::as_u64)
                    .unwrap_or(0),
                likes: m
                    .get("likes")
                    .and_then(serde_json::Value::as_u64)
                    .unwrap_or(0),
            })
        })
        .collect())
}

/// The sha256 HuggingFace publishes for one file, if it publishes one. It is
/// the LFS object id, which for these repositories is the sha256 of the content
/// -- so it is a free integrity check rather than a second source of truth.
fn published_sha256(repo: &str, file: &str) -> Option<String> {
    let url = format!("https://huggingface.co/api/models/{repo}/tree/main?expand=true");
    let body = get_json(&url).ok()?;
    body.as_array()?.iter().find_map(|f| {
        (f.get("path")?.as_str()? == file)
            .then(|| f.get("lfs")?.get("oid")?.as_str().map(str::to_string))?
    })
}

/// The GGUF files in one repository, largest last, with the ones already
/// downloaded marked.
pub fn files(repo: &str, store: &ModelStore) -> Result<Vec<RemoteFile>> {
    let url = format!("https://huggingface.co/api/models/{repo}/tree/main");
    let body = get_json(&url)?;
    let list = body
        .as_array()
        .context("HuggingFace returned something that is not a list")?;
    let have: std::collections::HashSet<String> = store.list().into_iter().map(|m| m.key).collect();
    let mut files: Vec<RemoteFile> = list
        .iter()
        .filter_map(|f| {
            let path = f.get("path")?.as_str()?.to_string();
            if !path.ends_with(".gguf") {
                return None;
            }
            Some(RemoteFile {
                bytes: f
                    .get("size")
                    .and_then(serde_json::Value::as_u64)
                    .unwrap_or(0),
                have: have.contains(&ModelStore::key_for(repo, &path)),
                path,
            })
        })
        .collect();
    files.sort_by_key(|f| f.bytes);
    Ok(files)
}

fn get_json(url: &str) -> Result<serde_json::Value> {
    let text = ureq::get(url)
        .header("User-Agent", USER_AGENT)
        .call()
        .and_then(|mut r| r.body_mut().read_to_string())
        .with_context(|| format!("GET {url}"))?;
    serde_json::from_str(&text).with_context(|| format!("{url} did not return JSON"))
}

/// Percent-encodes a query string. Deliberately minimal -- this escapes what a
/// model search actually contains rather than implementing the whole rule set,
/// and anything unexpected becomes `%XX` rather than being passed through.
fn urlencode(value: &str) -> String {
    value
        .bytes()
        .map(|b| match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
                (b as char).to_string()
            }
            b' ' => "+".to_string(),
            other => format!("%{other:02X}"),
        })
        .collect()
}