Lumotia/crates/transcription/src/model_manager.rs

use std::path::{Path, PathBuf};

use kon_core::error::{KonError, Result};
use kon_core::model_registry::{find_model, ModelFile};
use kon_core::types::{DownloadProgress, ModelId};

/// Resolve the models storage directory.
/// Windows: %LOCALAPPDATA%/kon/models
/// Unix: ~/.kon/models
pub fn models_dir() -> PathBuf {
    if cfg!(target_os = "windows") {
        let local_app_data = std::env::var("LOCALAPPDATA").unwrap_or_else(|_| ".".to_string());
        PathBuf::from(local_app_data).join("kon").join("models")
    } else {
        dirs_path().join("models")
    }
}

fn dirs_path() -> PathBuf {
    if cfg!(target_os = "windows") {
        let local_app_data = std::env::var("LOCALAPPDATA").unwrap_or_else(|_| ".".to_string());
        PathBuf::from(local_app_data).join("kon")
    } else {
        let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
        PathBuf::from(home).join(".kon")
    }
}

/// Get the directory path where a specific model's files are stored.
pub fn model_dir(id: &ModelId) -> PathBuf {
    models_dir().join(id.as_str())
}

/// Check whether all files for a model have been downloaded.
pub fn is_downloaded(id: &ModelId) -> bool {
    let entry = match find_model(id) {
        Some(e) => e,
        None => return false,
    };
    let dir = model_dir(id);
    entry.files.iter().all(|f| dir.join(f.filename).exists())
}

/// List all downloaded model IDs.
pub fn list_downloaded() -> Vec<ModelId> {
    kon_core::model_registry::all_models()
        .iter()
        .filter(|m| is_downloaded(&m.id))
        .map(|m| m.id.clone())
        .collect()
}

/// Download all files for a model, calling the progress callback per chunk.
/// Files are downloaded to a .part suffix and atomically renamed on completion.
pub async fn download(
    id: &ModelId,
    progress: impl Fn(DownloadProgress) + Send + 'static,
) -> Result<()> {
    let entry = find_model(id).ok_or_else(|| KonError::ModelNotFound(id.clone()))?;

    let dir = model_dir(id);
    std::fs::create_dir_all(&dir)?;

    for file in &entry.files {
        let dest = dir.join(file.filename);
        if dest.exists() {
            continue;
        }
        download_file(file, &dest, id, &progress).await?;
    }

    Ok(())
}

/// Download a single file with HTTP Range resume and optional SHA256 verification.
///
/// Resume pattern from Buzz (chidiwilliams/buzz): if a .part file exists,
/// send a Range header to resume from where we left off. SHA256 is checked
/// incrementally during download — no second pass over the file.
async fn download_file(
    file: &ModelFile,
    dest: &Path,
    model_id: &ModelId,
    progress: &(impl Fn(DownloadProgress) + Send),
) -> Result<()> {
    use futures_util::StreamExt;
    use sha2::{Digest, Sha256};

    let part_path = dest.with_extension(
        dest.extension()
            .map(|e| format!("{}.part", e.to_string_lossy()))
            .unwrap_or_else(|| "part".to_string()),
    );

    let client = reqwest::Client::builder()
        .connect_timeout(std::time::Duration::from_secs(30))
        .build()
        .map_err(|e| KonError::DownloadFailed(e.to_string()))?;

    // Check for existing partial download (resume support)
    let existing_bytes = if part_path.exists() {
        std::fs::metadata(&part_path).map(|m| m.len()).unwrap_or(0)
    } else {
        0
    };

    let mut request = client.get(file.url);

    // If we have a partial file and no SHA256 to verify (can't verify partial),
    // request a range resume. If SHA256 is set, we restart to ensure integrity.
    let resuming = existing_bytes > 0 && file.sha256.is_none();
    if resuming {
        request = request.header("Range", format!("bytes={existing_bytes}-"));
    }

    let response = request
        .send()
        .await
        .map_err(|e| KonError::DownloadFailed(e.to_string()))?;

    // Check if server supports range (206 Partial Content) or gave full file (200)
    let actually_resuming = resuming && response.status().as_u16() == 206;

    let total_bytes = if actually_resuming {
        // Content-Range: bytes START-END/TOTAL — extract TOTAL
        response
            .headers()
            .get("content-range")
            .and_then(|v| v.to_str().ok())
            .and_then(|s| s.rsplit('/').next())
            .and_then(|s| s.parse::<u64>().ok())
            .unwrap_or(0)
    } else {
        response.content_length().unwrap_or(0)
    };

    let mut stream = response.bytes_stream();
    let mut downloaded: u64 = if actually_resuming { existing_bytes } else { 0 };
    let mut last_percent: u8 = 0;

    // Open file for append (resume) or create (fresh start)
    let mut out = if actually_resuming {
        std::fs::OpenOptions::new().append(true).open(&part_path)?
    } else {
        std::fs::File::create(&part_path)?
    };

    // Incremental SHA256 — only when a checksum is provided
    let mut hasher = file.sha256.map(|_| Sha256::new());

    // If resuming without SHA256, we can't hash the already-downloaded portion,
    // but we also don't need to — we only hash when sha256 is set, and we
    // restart from scratch in that case.

    while let Some(chunk) = stream.next().await {
        let chunk = chunk.map_err(|e| KonError::DownloadFailed(e.to_string()))?;
        std::io::Write::write_all(&mut out, &chunk)?;
        if let Some(ref mut h) = hasher {
            h.update(&chunk);
        }
        downloaded += chunk.len() as u64;

        let percent = if total_bytes > 0 {
            ((downloaded as f64 / total_bytes as f64) * 100.0) as u8
        } else {
            0
        };

        if percent != last_percent {
            last_percent = percent;
            progress(DownloadProgress {
                model_id: model_id.clone(),
                file_name: file.filename.to_string(),
                bytes_downloaded: downloaded,
                total_bytes,
                percent,
            });
        }
    }

    drop(out);

    // Verify SHA256 if provided
    if let (Some(expected), Some(hasher)) = (file.sha256, hasher) {
        let actual = format!("{:x}", hasher.finalize());
        if actual != expected {
            // Delete corrupt file so next attempt starts fresh
            let _ = std::fs::remove_file(&part_path);
            return Err(KonError::DownloadFailed(format!(
                "SHA256 mismatch for {}: expected {}, got {}",
                file.filename, expected, actual
            )));
        }
    }

    // Atomic rename — file is complete and verified
    std::fs::rename(&part_path, dest)?;

    Ok(())
}

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

    #[test]
    fn model_dir_returns_correct_path() {
        let id = ModelId::new("whisper-tiny-en");
        let path = model_dir(&id);
        assert!(path.to_string_lossy().contains("whisper-tiny-en"));
    }

    #[test]
    fn is_downloaded_returns_false_for_missing_model() {
        let id = ModelId::new("nonexistent-model");
        assert!(!is_downloaded(&id));
    }

    #[test]
    fn list_downloaded_returns_empty_when_no_models() {
        let list = list_downloaded();
        // In test environment, no models are downloaded
        // This just verifies the function doesn't panic
        assert!(list.len() <= kon_core::model_registry::all_models().len());
    }
}