llama-cpp-v3 0.1.6

use crate::backend::Backend;
use serde_json::Value;
use std::fs::{self, File};
use std::io::{self, BufReader};
use std::path::PathBuf;

#[derive(Debug, thiserror::Error)]
pub enum DownloadError {
    #[error("IO error: {0}")]
    Io(#[from] io::Error),
    #[error("Network error: {0}")]
    Network(String),
    #[error("Failed to find appropriate release asset for this OS/Backend")]
    AssetNotFound,
    #[error("ZIP extraction error: {0}")]
    Zip(#[from] zip::result::ZipError),
    #[error("Missing DLL in ZIP")]
    MissingDll,
}

pub struct Downloader;

impl Downloader {
    pub fn ensure_dll(
        backend: Backend,
        app_name: &str,
        version: Option<&str>,
        cache_dir: Option<PathBuf>,
    ) -> Result<PathBuf, DownloadError> {
        let target_dir = if let Some(dir) = cache_dir {
            dir
        } else {
            let mut d = dirs::cache_dir().unwrap_or_else(|| PathBuf::from("."));
            d.push(app_name);
            d.push("llama-cpp-v3");
            d.push(backend.release_name_component());
            d.push(version.unwrap_or("latest"));
            d
        };

        fs::create_dir_all(&target_dir)?;

        let dll_path = target_dir.join(backend.dll_name());

        if dll_path.exists() {
            println!("DLL path resolved to: {}", dll_path.display());
            // Already cached/downloaded
            return Ok(dll_path);
        }

        println!(
            "Downloading llama.cpp {} backend...",
            backend.release_name_component()
        );

        // 1. Fetch release metadata
        let url = if let Some(v) = version {
            format!(
                "https://api.github.com/repos/ggml-org/llama.cpp/releases/tags/{}",
                v
            )
        } else {
            "https://api.github.com/repos/ggml-org/llama.cpp/releases/latest".to_string()
        };

        let response = ureq::get(&url)
            .header("User-Agent", "llama-cpp-v3-rust-wrapper")
            .call()
            .map_err(|e| DownloadError::Network(e.to_string()))?;

        let release: Value = response
            .into_body()
            .read_json()
            .map_err(|e| DownloadError::Network(format!("JSON parsing error: {}", e)))?;

        // Determine OS and Arch string matchers
        let os_str = if cfg!(windows) {
            "win"
        } else if cfg!(target_os = "macos") {
            "mac"
        } else {
            "ubuntu"
        };

        let arch_str = if cfg!(target_arch = "x86_64") {
            "x64"
        } else if cfg!(target_arch = "aarch64") {
            "arm64"
        } else {
            "x86"
        };

        let backend_str = backend.release_name_component();

        // Find the right asset zip
        let assets = release["assets"]
            .as_array()
            .ok_or(DownloadError::AssetNotFound)?;

        let mut download_url = None;
        for asset in assets {
            if let Some(name) = asset["name"].as_str() {
                let name = name.to_lowercase();
                if name.ends_with(".zip")
                    && name.contains(os_str)
                    && name.contains(arch_str)
                    && name.contains(backend_str)
                {
                    download_url = asset["browser_download_url"].as_str().map(String::from);
                    break;
                }
            }
        }

        let download_url = download_url.ok_or(DownloadError::AssetNotFound)?;

        // 2. Download the ZIP file
        let zip_path = target_dir.join("temp.zip");
        {
            let mut file = File::create(&zip_path)?;
            let mut response = ureq::get(&download_url)
                .call()
                .map_err(|e| DownloadError::Network(e.to_string()))?;
            io::copy(&mut response.body_mut().as_reader(), &mut file)?;
        }

        // 3. Extract the DLL from the ZIP
        {
            let file = File::open(&zip_path)?;
            let mut archive = zip::ZipArchive::new(BufReader::new(file))?;

            let mut found = false;
            for i in 0..archive.len() {
                let mut file = archive.by_index(i)?;
                let outpath = match file.enclosed_name() {
                    Some(path) => path.to_owned(),
                    None => continue,
                };

                if file.is_dir() {
                    let dir_path = target_dir.join(&outpath);
                    fs::create_dir_all(&dir_path)?;
                } else {
                    if let Some(file_name) = outpath.file_name() {
                        let extracted_path = target_dir.join(file_name);
                        let mut outfile = File::create(&extracted_path)?;
                        io::copy(&mut file, &mut outfile)?;

                        if file_name.to_string_lossy() == backend.dll_name() {
                            found = true;
                        }
                    }
                }
            }

            if !found {
                return Err(DownloadError::MissingDll);
            }
        }

        // Clean up the zip file
        let _ = fs::remove_file(&zip_path);

        println!("DLL path resolved to: {}", dll_path.display());
        Ok(dll_path)
    }
}