minni 0.1.1

use anyhow::{Context, Result};
use indicatif::{ProgressBar, ProgressStyle};
use std::fs::{self, File};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};

/// Hugging Face model ID for reranking.
pub const MODEL_ID: &str = "cross-encoder/ms-marco-MiniLM-L6-v2";
/// Hugging Face model ID for dense retrieval.
pub const DENSE_MODEL_ID: &str = "sentence-transformers/all-MiniLM-L6-v2";

// Files needed for the reranker.
const REQUIRED_FILES: &[&str] = &["model.onnx", "tokenizer.json"];
const DENSE_REQUIRED_FILES: &[&str] = &["model.onnx", "tokenizer.json"];

// Hugging Face API endpoint.
const HF_API_BASE: &str = "https://huggingface.co";

pub fn get_model_dir() -> Result<PathBuf> {
    let home =
        dirs::home_dir().ok_or_else(|| anyhow::anyhow!("Could not determine home directory"))?;
    let model_dir = home
        .join(".minni")
        .join("models")
        .join("ms-marco-MiniLM-L6-v2");
    Ok(model_dir)
}

pub fn is_model_installed() -> bool {
    if let Ok(model_dir) = get_model_dir() {
        REQUIRED_FILES.iter().all(|f| model_dir.join(f).exists())
    } else {
        false
    }
}

pub fn ensure_model_available() -> Result<PathBuf> {
    let model_dir = get_model_dir()?;

    if is_model_installed() {
        return Ok(model_dir);
    }

    println!("Downloading reranker model (first run only)...");
    println!("Model: {}", MODEL_ID);
    println!();

    fs::create_dir_all(&model_dir)?;

    // Download each required file
    for file_name in REQUIRED_FILES {
        download_file_from_hf(MODEL_ID, file_name, &model_dir)?;
    }

    println!();
    println!("Model downloaded successfully!");

    Ok(model_dir)
}

pub fn get_dense_model_dir() -> Result<PathBuf> {
    let home =
        dirs::home_dir().ok_or_else(|| anyhow::anyhow!("Could not determine home directory"))?;
    let model_dir = home.join(".minni").join("models").join("all-MiniLM-L6-v2");
    Ok(model_dir)
}

pub fn is_dense_model_installed() -> bool {
    if let Ok(model_dir) = get_dense_model_dir() {
        DENSE_REQUIRED_FILES
            .iter()
            .all(|f| model_dir.join(f).exists())
    } else {
        false
    }
}

pub fn ensure_dense_model_available() -> Result<PathBuf> {
    let model_dir = get_dense_model_dir()?;

    if is_dense_model_installed() {
        return Ok(model_dir);
    }

    println!("Downloading dense retriever model (first run only)...");
    println!("Model: {}", DENSE_MODEL_ID);
    println!();

    fs::create_dir_all(&model_dir)?;

    for file_name in DENSE_REQUIRED_FILES {
        download_file_from_hf(DENSE_MODEL_ID, file_name, &model_dir)?;
    }

    println!();
    println!("Dense model downloaded successfully!");

    Ok(model_dir)
}

fn download_file_from_hf(model_id: &str, file_name: &str, dest_dir: &Path) -> Result<()> {
    // Cross-encoder ONNX files are in onnx/ subfolder, tokenizer is in root
    let url = if file_name == "model.onnx" {
        format!(
            "{}/{}/resolve/main/onnx/{}",
            HF_API_BASE, model_id, file_name
        )
    } else {
        // tokenizer.json and other files are in root
        format!("{}/{}/resolve/main/{}", HF_API_BASE, model_id, file_name)
    };

    let dest_path = dest_dir.join(file_name);

    println!("Downloading {}...", file_name);

    // Make request
    let response = ureq::get(&url)
        .call()
        .with_context(|| format!("Failed to download {}", file_name))?;

    // Get content length for progress bar
    let content_length: u64 = response
        .header("Content-Length")
        .and_then(|s| s.parse().ok())
        .unwrap_or(0);

    // Create progress bar
    let pb = ProgressBar::new(content_length);
    pb.set_style(
        ProgressStyle::default_bar()
            .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .unwrap()
            .progress_chars("#>-"),
    );

    // Read and write with progress
    let mut reader = response.into_reader();
    let mut file = File::create(&dest_path)
        .with_context(|| format!("Failed to create file: {:?}", dest_path))?;

    let mut buffer = [0u8; 8192];
    let mut downloaded: u64 = 0;

    loop {
        let bytes_read = reader.read(&mut buffer)?;
        if bytes_read == 0 {
            break;
        }
        file.write_all(&buffer[..bytes_read])?;
        downloaded += bytes_read as u64;
        pb.set_position(downloaded);
    }

    pb.finish_with_message("done");
    Ok(())
}

// Alternative: Download from a direct URL if HF is slow
#[allow(dead_code)]
pub fn download_from_url(url: &str, dest_path: &Path) -> Result<()> {
    let response = ureq::get(url)
        .call()
        .with_context(|| format!("Failed to download from {}", url))?;

    let content_length: u64 = response
        .header("Content-Length")
        .and_then(|s| s.parse().ok())
        .unwrap_or(0);

    let pb = ProgressBar::new(content_length);
    pb.set_style(
        ProgressStyle::default_bar()
            .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .unwrap()
            .progress_chars("#>-"),
    );

    let mut reader = response.into_reader();
    let mut file = File::create(dest_path)?;
    let mut buffer = [0u8; 8192];
    let mut downloaded: u64 = 0;

    loop {
        let bytes_read = reader.read(&mut buffer)?;
        if bytes_read == 0 {
            break;
        }
        file.write_all(&buffer[..bytes_read])?;
        downloaded += bytes_read as u64;
        pb.set_position(downloaded);
    }

    pb.finish_with_message("done");
    Ok(())
}