Lumotia/crates/llm/src/model_manager.rs

use std::fmt;
use std::io;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::sync::{LazyLock, Mutex};

use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

#[allow(non_camel_case_types)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum LlmModelId {
    #[serde(rename = "qwen3_5_2b")]
    Qwen3_5_2B_Q4,
    #[serde(rename = "qwen3_5_4b")]
    Qwen3_5_4B_Q4,
    #[serde(rename = "qwen3_5_9b")]
    Qwen3_5_9B_Q4,
    #[serde(rename = "qwen3_6_27b")]
    Qwen3_6_27B_Q4,
}

impl LlmModelId {
    pub fn default_tier() -> Self {
        Self::Qwen3_5_4B_Q4
    }

    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Qwen3_5_2B_Q4 => "qwen3_5_2b",
            Self::Qwen3_5_4B_Q4 => "qwen3_5_4b",
            Self::Qwen3_5_9B_Q4 => "qwen3_5_9b",
            Self::Qwen3_6_27B_Q4 => "qwen3_6_27b",
        }
    }

    pub fn display_name(&self) -> &'static str {
        match self {
            Self::Qwen3_5_2B_Q4 => "Qwen3.5 2B",
            Self::Qwen3_5_4B_Q4 => "Qwen3.5 4B",
            Self::Qwen3_5_9B_Q4 => "Qwen3.5 9B",
            Self::Qwen3_6_27B_Q4 => "Qwen3.6 27B",
        }
    }

    pub fn file_name(&self) -> &'static str {
        match self {
            Self::Qwen3_5_2B_Q4 => "Qwen3.5-2B-Q4_K_M.gguf",
            Self::Qwen3_5_4B_Q4 => "Qwen3.5-4B-Q4_K_M.gguf",
            Self::Qwen3_5_9B_Q4 => "Qwen3.5-9B-Q4_K_M.gguf",
            Self::Qwen3_6_27B_Q4 => "Qwen3.6-27B-Q4_K_M.gguf",
        }
    }

    pub fn size_bytes(&self) -> u64 {
        match self {
            Self::Qwen3_5_2B_Q4 => 1_280_835_840,
            Self::Qwen3_5_4B_Q4 => 2_740_937_888,
            Self::Qwen3_5_9B_Q4 => 5_680_522_464,
            Self::Qwen3_6_27B_Q4 => 16_817_244_384,
        }
    }

    pub fn minimum_ram_bytes(&self) -> u64 {
        match self {
            Self::Qwen3_5_2B_Q4 => 8 * 1024_u64.pow(3),
            Self::Qwen3_5_4B_Q4 => 16 * 1024_u64.pow(3),
            Self::Qwen3_5_9B_Q4 => 32 * 1024_u64.pow(3),
            Self::Qwen3_6_27B_Q4 => 64 * 1024_u64.pow(3),
        }
    }

    pub fn recommended_vram_bytes(&self) -> Option<u64> {
        match self {
            Self::Qwen3_5_2B_Q4 => None,
            Self::Qwen3_5_4B_Q4 => Some(6 * 1024_u64.pow(3)),
            Self::Qwen3_5_9B_Q4 => Some(12 * 1024_u64.pow(3)),
            Self::Qwen3_6_27B_Q4 => Some(24 * 1024_u64.pow(3)),
        }
    }

    pub fn description(&self) -> &'static str {
        match self {
            Self::Qwen3_5_2B_Q4 => "Minimal tier for 8 GB RAM and CPU-heavy machines.",
            Self::Qwen3_5_4B_Q4 => {
                "Standard tier for cleanup and task extraction on 16 GB systems."
            }
            Self::Qwen3_5_9B_Q4 => "High tier for 32 GB RAM with a 12 GB+ GPU.",
            Self::Qwen3_6_27B_Q4 => {
                "Maximum tier for 64 GB RAM with a 24 GB GPU; partial CPU offload below that."
            }
        }
    }

    pub fn hf_url(&self) -> &'static str {
        match self {
            Self::Qwen3_5_2B_Q4 => {
                "https://huggingface.co/unsloth/Qwen3.5-2B-GGUF/resolve/f6d5376be1edb4d416d56da11e5397a961aca8ae/Qwen3.5-2B-Q4_K_M.gguf"
            }
            Self::Qwen3_5_4B_Q4 => {
                "https://huggingface.co/unsloth/Qwen3.5-4B-GGUF/resolve/e87f176479d0855a907a41277aca2f8ee7a09523/Qwen3.5-4B-Q4_K_M.gguf"
            }
            Self::Qwen3_5_9B_Q4 => {
                "https://huggingface.co/unsloth/Qwen3.5-9B-GGUF/resolve/3885219b6810b007914f3a7950a8d1b469d598a5/Qwen3.5-9B-Q4_K_M.gguf"
            }
            Self::Qwen3_6_27B_Q4 => {
                "https://huggingface.co/unsloth/Qwen3.6-27B-GGUF/resolve/82d411acf4a06cfb8d9b073a5211bf410bfc29bf/Qwen3.6-27B-Q4_K_M.gguf"
            }
        }
    }

    pub fn sha256(&self) -> &'static str {
        match self {
            Self::Qwen3_5_2B_Q4 => {
                "aaf42c8b7c3cab2bf3d69c355048d4a0ee9973d48f16c731c0520ee914699223"
            }
            Self::Qwen3_5_4B_Q4 => {
                "00fe7986ff5f6b463e62455821146049db6f9313603938a70800d1fb69ef11a4"
            }
            Self::Qwen3_5_9B_Q4 => {
                "03b74727a860a56338e042c4420bb3f04b2fec5734175f4cb9fa853daf52b7e8"
            }
            Self::Qwen3_6_27B_Q4 => {
                "5ed60d0af4650a854b1755bd392f9aef4872643dc25a254bc68043fa638392a0"
            }
        }
    }
}

impl fmt::Display for LlmModelId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

impl FromStr for LlmModelId {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        match value {
            "qwen3_5_2b" => Ok(Self::Qwen3_5_2B_Q4),
            "qwen3_5_4b" => Ok(Self::Qwen3_5_4B_Q4),
            "qwen3_5_9b" => Ok(Self::Qwen3_5_9B_Q4),
            "qwen3_6_27b" => Ok(Self::Qwen3_6_27B_Q4),
            other => Err(format!("Unknown LLM model id: {other}")),
        }
    }
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LlmModelInfo {
    pub id: String,
    pub display_name: &'static str,
    pub file_name: &'static str,
    pub size_bytes: u64,
    pub description: &'static str,
    pub minimum_ram_bytes: u64,
    pub recommended_vram_bytes: Option<u64>,
}

#[derive(Debug, thiserror::Error)]
pub enum DownloadError {
    #[error("http error: {0}")]
    Http(String),
    #[error("io error: {0}")]
    Io(#[from] io::Error),
    #[error("sha256 mismatch: expected {expected}, got {actual}")]
    ShaMismatch { expected: String, actual: String },
    #[error("resume failed: server does not support range requests")]
    ResumeUnsupported,
}

const ALL_MODELS: &[LlmModelId] = &[
    LlmModelId::Qwen3_5_2B_Q4,
    LlmModelId::Qwen3_5_4B_Q4,
    LlmModelId::Qwen3_5_9B_Q4,
    LlmModelId::Qwen3_6_27B_Q4,
];

static ACTIVE_DOWNLOADS: LazyLock<Mutex<std::collections::HashSet<LlmModelId>>> =
    LazyLock::new(|| Mutex::new(std::collections::HashSet::new()));

struct DownloadReservation {
    id: LlmModelId,
}

impl DownloadReservation {
    fn acquire(id: LlmModelId) -> Result<Self, DownloadError> {
        let mut active = ACTIVE_DOWNLOADS
            .lock()
            .map_err(|_| DownloadError::Http("download lock poisoned".into()))?;
        if !active.insert(id) {
            return Err(DownloadError::Http(format!(
                "download already in progress for {}",
                id.as_str()
            )));
        }
        Ok(Self { id })
    }
}

impl Drop for DownloadReservation {
    fn drop(&mut self) {
        if let Ok(mut active) = ACTIVE_DOWNLOADS.lock() {
            active.remove(&self.id);
        }
    }
}

pub fn all_models() -> &'static [LlmModelId] {
    ALL_MODELS
}

pub fn model_info(id: LlmModelId) -> LlmModelInfo {
    LlmModelInfo {
        id: id.as_str().to_string(),
        display_name: id.display_name(),
        file_name: id.file_name(),
        size_bytes: id.size_bytes(),
        description: id.description(),
        minimum_ram_bytes: id.minimum_ram_bytes(),
        recommended_vram_bytes: id.recommended_vram_bytes(),
    }
}

pub fn recommend_tier(total_ram_bytes: u64, total_vram_bytes: Option<u64>) -> LlmModelId {
    let vram = total_vram_bytes.unwrap_or(0);
    if vram >= 24 * 1024_u64.pow(3) && total_ram_bytes >= 64 * 1024_u64.pow(3) {
        LlmModelId::Qwen3_6_27B_Q4
    } else if vram >= 12 * 1024_u64.pow(3) && total_ram_bytes >= 32 * 1024_u64.pow(3) {
        LlmModelId::Qwen3_5_9B_Q4
    } else if vram >= 6 * 1024_u64.pow(3) || total_ram_bytes >= 16 * 1024_u64.pow(3) {
        LlmModelId::Qwen3_5_4B_Q4
    } else {
        LlmModelId::Qwen3_5_2B_Q4
    }
}

pub fn model_dir() -> PathBuf {
    magnotia_core::paths::app_paths().llm_models_dir()
}

pub fn model_path(id: LlmModelId) -> PathBuf {
    model_dir().join(id.file_name())
}

pub fn partial_download_path(id: LlmModelId) -> PathBuf {
    model_path(id).with_extension("gguf.part")
}

pub fn is_downloaded(id: LlmModelId) -> bool {
    model_path(id).exists()
}

pub fn delete_model(id: LlmModelId) -> io::Result<()> {
    let final_path = model_path(id);
    let partial_path = partial_download_path(id);

    if final_path.exists() {
        std::fs::remove_file(final_path)?;
    }
    if partial_path.exists() {
        std::fs::remove_file(partial_path)?;
    }

    Ok(())
}

pub async fn download_model<F>(id: LlmModelId, on_progress: F) -> Result<(), DownloadError>
where
    F: FnMut(u64, u64) + Send + 'static,
{
    let _reservation = DownloadReservation::acquire(id)?;
    let dest = model_path(id);
    tokio::fs::create_dir_all(model_dir()).await?;

    if dest.exists() {
        let actual = sha256_file(&dest).await?;
        if actual == id.sha256() {
            return Ok(());
        }
        tokio::fs::remove_file(&dest).await?;
    }

    download_impl(id.hf_url(), id.sha256(), &dest, on_progress).await
}

async fn sha256_file(path: &Path) -> Result<String, io::Error> {
    let mut hasher = Sha256::new();
    let mut file = tokio::fs::File::open(path).await?;
    let mut buffer = [0u8; 8192];

    loop {
        let count = file.read(&mut buffer).await?;
        if count == 0 {
            break;
        }
        hasher.update(&buffer[..count]);
    }

    Ok(format!("{:x}", hasher.finalize()))
}

async fn download_impl<F>(
    url: &str,
    expected_sha: &str,
    dest: &Path,
    mut on_progress: F,
) -> Result<(), DownloadError>
where
    F: FnMut(u64, u64) + Send + 'static,
{
    let tmp = dest.with_extension("gguf.part");
    let resume_from = tokio::fs::metadata(&tmp)
        .await
        .ok()
        .map(|m| m.len())
        .unwrap_or(0);

    let client = reqwest::Client::builder()
        .user_agent("magnotia/0.1.0")
        .connect_timeout(std::time::Duration::from_secs(30))
        .build()
        .map_err(|e| DownloadError::Http(e.to_string()))?;

    let mut request = client.get(url);
    if resume_from > 0 {
        request = request.header(reqwest::header::RANGE, format!("bytes={resume_from}-"));
    }

    let response = request
        .send()
        .await
        .map_err(|e| DownloadError::Http(e.to_string()))?;
    if resume_from > 0 && response.status() != reqwest::StatusCode::PARTIAL_CONTENT {
        return Err(DownloadError::ResumeUnsupported);
    }
    if !response.status().is_success() && response.status() != reqwest::StatusCode::PARTIAL_CONTENT
    {
        return Err(DownloadError::Http(format!("status {}", response.status())));
    }

    let total = if resume_from > 0 {
        response
            .headers()
            .get(reqwest::header::CONTENT_RANGE)
            .and_then(|value| value.to_str().ok())
            .and_then(|value| value.rsplit('/').next())
            .and_then(|value| value.parse::<u64>().ok())
            .unwrap_or_else(|| response.content_length().unwrap_or(0) + resume_from)
    } else {
        response.content_length().unwrap_or(0)
    };

    let mut hasher = Sha256::new();
    if resume_from > 0 {
        let mut partial = tokio::fs::File::open(&tmp).await?;
        let mut buffer = [0u8; 8192];
        loop {
            let count = partial.read(&mut buffer).await?;
            if count == 0 {
                break;
            }
            hasher.update(&buffer[..count]);
        }
    }

    let mut output = tokio::fs::OpenOptions::new()
        .create(true)
        .append(true)
        .open(&tmp)
        .await?;

    let mut downloaded = resume_from;
    let mut stream = response.bytes_stream();
    while let Some(chunk) = stream.next().await {
        let chunk = chunk.map_err(|e| DownloadError::Http(e.to_string()))?;
        output.write_all(&chunk).await?;
        hasher.update(&chunk);
        downloaded += chunk.len() as u64;
        on_progress(downloaded, total);
    }
    output.flush().await?;
    drop(output);

    let actual = format!("{:x}", hasher.finalize());
    if actual != expected_sha {
        tokio::fs::remove_file(&tmp).await.ok();
        return Err(DownloadError::ShaMismatch {
            expected: expected_sha.to_string(),
            actual,
        });
    }

    tokio::fs::rename(&tmp, dest).await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{Arc, Mutex};
    use tempfile::tempdir;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::net::TcpListener;

    #[test]
    fn model_path_contains_model_dir_and_filename() {
        let path = model_path(LlmModelId::Qwen3_5_2B_Q4);
        assert!(path.to_string_lossy().ends_with("Qwen3.5-2B-Q4_K_M.gguf"));
        assert!(path.starts_with(model_dir()));
    }

    #[test]
    fn recommend_tier_prefers_mid_by_default() {
        let tier = recommend_tier(16 * 1024_u64.pow(3), None);
        assert_eq!(tier, LlmModelId::Qwen3_5_4B_Q4);
    }

    #[tokio::test]
    async fn download_impl_supports_resume_and_sha_verification() {
        let fixture = b"hello resumed download".to_vec();
        let expected_sha = format!("{:x}", Sha256::digest(&fixture));
        let server = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = server.local_addr().unwrap();
        let content = fixture.clone();

        let server_task = tokio::spawn(async move {
            let (mut socket, _) = server.accept().await.unwrap();
            let mut request = vec![0u8; 2048];
            let size = socket.read(&mut request).await.unwrap();
            let request = String::from_utf8_lossy(&request[..size]).to_lowercase();
            let range_start = request
                .lines()
                .find_map(|line| line.strip_prefix("range: bytes="))
                .and_then(|line| line.strip_suffix('-'))
                .and_then(|line| line.trim().parse::<usize>().ok());

            if let Some(start) = range_start {
                let body = &content[start..];
                let response = format!(
                    "HTTP/1.1 206 Partial Content\r\nContent-Length: {}\r\nContent-Range: bytes {}-{}/{}\r\nAccept-Ranges: bytes\r\n\r\n",
                    body.len(),
                    start,
                    content.len() - 1,
                    content.len()
                );
                socket.write_all(response.as_bytes()).await.unwrap();
                socket.write_all(body).await.unwrap();
            } else {
                let response = format!(
                    "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nAccept-Ranges: bytes\r\n\r\n",
                    content.len()
                );
                socket.write_all(response.as_bytes()).await.unwrap();
                socket.write_all(&content).await.unwrap();
            }
        });

        let dir = tempdir().unwrap();
        let dest = dir.path().join("fixture.gguf");
        let part = dest.with_extension("gguf.part");
        tokio::fs::write(&part, &fixture[..10]).await.unwrap();

        let progress = Arc::new(Mutex::new(Vec::new()));
        let progress_clone = progress.clone();
        download_impl(
            &format!("http://{addr}/fixture.gguf"),
            &expected_sha,
            &dest,
            move |done, total| progress_clone.lock().unwrap().push((done, total)),
        )
        .await
        .unwrap();

        let saved = tokio::fs::read(&dest).await.unwrap();
        assert_eq!(saved, fixture);
        assert!(!part.exists());
        assert!(!progress.lock().unwrap().is_empty());

        server_task.await.unwrap();
    }
}