//! Managed model acquisition helpers.

use super::track_managed_model_usage;
use crate::cli::output::{emit_event, interactive_tui_active, ModelProgressStatus, OutputEvent};
use crate::cli::terminal_progress::{start_spinner, SpinnerHandle};
use anyhow::{Context, Result};
use hf_hub::{DownloadEvent, Progress, ProgressEvent, ProgressHandler, RepoDownloadFileParams};
#[cfg(test)]
use std::collections::HashMap;
use std::io::Write;
use std::path::{Path, PathBuf};
#[cfg(test)]
use std::sync::LazyLock;
use std::sync::{Arc, Mutex};

/// Get the canonical managed model root (the Hugging Face hub cache).
pub fn models_dir() -> PathBuf {
    crate::models::huggingface_hub_cache_dir()
}

/// Parse a size string like "20GB", "4.4GB", "491MB" into GB as f64.
pub fn parse_size_gb(s: &str) -> f64 {
    let s = s.trim();
    if let Some(gb) = s.strip_suffix("GB") {
        gb.trim().parse().unwrap_or(0.0)
    } else if let Some(mb) = s.strip_suffix("MB") {
        mb.trim().parse::<f64>().unwrap_or(0.0) / 1000.0
    } else {
        0.0
    }
}

#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
struct HfAsset {
    repo: String,
    revision: String,
    file: String,
}

impl HfAsset {
    fn repo_parts(&self) -> (&str, &str) {
        self.repo
            .split_once('/')
            .unwrap_or(("", self.repo.as_str()))
    }
}

fn expand_split_asset(asset: &HfAsset) -> Result<Vec<HfAsset>> {
    let re = regex_lite::Regex::new(r"-00001-of-(\d{5})\.gguf$").unwrap();
    let Some(caps) = re.captures(&asset.file) else {
        return Ok(vec![asset.clone()]);
    };
    let count: u32 = caps[1].parse()?;
    Ok((1..=count)
        .map(|index| HfAsset {
            repo: asset.repo.clone(),
            revision: asset.revision.clone(),
            file: asset
                .file
                .replace("-00001-of-", &format!("-{index:05}-of-")),
        })
        .collect())
}

fn is_mlx_primary_asset(file: &str) -> bool {
    matches!(file, "model.safetensors" | "model.safetensors.index.json")
        || is_split_mlx_first_shard_file(file)
}

/// Returns true if `file` is the first shard of a sharded MLX safetensors set,
/// i.e. `model-00001-of-NNNNN.safetensors`.
fn is_split_mlx_first_shard_file(file: &str) -> bool {
    let Some(rest) = file.strip_prefix("model-") else {
        return false;
    };
    let Some(rest) = rest.strip_suffix(".safetensors") else {
        return false;
    };
    let Some((left, right)) = rest.split_once("-of-") else {
        return false;
    };
    left == "00001" && right.len() == 5 && right.bytes().all(|b| b.is_ascii_digit())
}

/// Expands a first-shard MLX ref (`model-00001-of-NNNNN.safetensors`) into the
/// full list of shard assets without needing to download the index.
fn expand_split_mlx_first_shard(asset: &HfAsset) -> Vec<HfAsset> {
    let Some(rest) = asset.file.strip_prefix("model-00001-of-") else {
        return Vec::new();
    };
    let Some(total_str) = rest.strip_suffix(".safetensors") else {
        return Vec::new();
    };
    if total_str.len() != 5 || !total_str.bytes().all(|b| b.is_ascii_digit()) {
        return Vec::new();
    }
    let Ok(count): Result<u32> = total_str.parse().map_err(anyhow::Error::from) else {
        return Vec::new();
    };
    (1..=count)
        .map(|index| HfAsset {
            repo: asset.repo.clone(),
            revision: asset.revision.clone(),
            file: format!("model-{index:05}-of-{total_str}.safetensors"),
        })
        .collect()
}

fn mlx_sidecar_assets(asset: &HfAsset) -> Vec<(bool, HfAsset)> {
    [
        (true, "tokenizer.json"),
        (false, "tokenizer_config.json"),
        (false, "chat_template.jinja"),
        (false, "chat_template.json"),
    ]
    .into_iter()
    .map(|(required, file)| {
        (
            required,
            HfAsset {
                repo: asset.repo.clone(),
                revision: asset.revision.clone(),
                file: file.to_string(),
            },
        )
    })
    .collect()
}

fn is_optional_metadata(required: bool, _asset: &HfAsset) -> bool {
    !required
}

fn parse_safetensors_index_shards(index: &serde_json::Value) -> Result<Vec<String>> {
    let weight_map = index["weight_map"]
        .as_object()
        .context("missing weight_map in safetensors index")?;
    let mut shards = std::collections::BTreeSet::new();
    for file in weight_map.values() {
        let file = file
            .as_str()
            .context("weight_map value in safetensors index is not a string")?;
        shards.insert(file.to_string());
    }
    Ok(shards.into_iter().collect())
}

fn ensure_cached_hf_asset(api: &hf_hub::HFClientSync, asset: &HfAsset) -> Result<PathBuf> {
    let (owner, name) = asset.repo_parts();
    api.model(owner, name)
        .download_file(
            &RepoDownloadFileParams::builder()
                .filename(asset.file.clone())
                .revision(asset.revision.clone())
                .build(),
        )
        .with_context(|| {
            format!(
                "Cache Hugging Face asset {}/{}@{}",
                asset.repo, asset.file, asset.revision
            )
        })
}

fn mlx_sharded_weight_assets(api: &hf_hub::HFClientSync, asset: &HfAsset) -> Result<Vec<HfAsset>> {
    if asset.file != "model.safetensors.index.json" {
        return Ok(Vec::new());
    }
    let index_path = ensure_cached_hf_asset(api, asset)?;
    let index_text = std::fs::read_to_string(&index_path)
        .with_context(|| format!("Read {}", index_path.display()))?;
    let index: serde_json::Value = serde_json::from_str(&index_text)
        .with_context(|| format!("Parse {}", index_path.display()))?;
    Ok(parse_safetensors_index_shards(&index)?
        .into_iter()
        .map(|file| HfAsset {
            repo: asset.repo.clone(),
            revision: asset.revision.clone(),
            file,
        })
        .collect())
}

#[cfg(test)]
type DownloadHfAssetsOverrideFn =
    Arc<dyn Fn(&str, Vec<HfAsset>) -> Result<Vec<PathBuf>> + Send + Sync>;

#[cfg(test)]
type DownloadPlanObserverFn = Arc<dyn Fn(&str, Vec<(bool, String)>) + Send + Sync>;

#[cfg(test)]
type DownloadHfAssetsLabelOverrideFn = Arc<dyn Fn(&str) -> Result<Vec<PathBuf>> + Send + Sync>;

#[cfg(test)]
static DOWNLOAD_HF_ASSETS_OVERRIDE: LazyLock<Mutex<HashMap<String, DownloadHfAssetsOverrideFn>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

#[cfg(test)]
static DOWNLOAD_PLAN_OBSERVER: LazyLock<Mutex<Option<DownloadPlanObserverFn>>> =
    LazyLock::new(|| Mutex::new(None));

#[cfg(test)]
pub(crate) struct DownloadHfAssetsOverrideGuard(String);

#[cfg(test)]
pub(crate) struct DownloadPlanObserverGuard;

#[cfg(test)]
impl DownloadHfAssetsOverrideGuard {
    fn set(label: String, func: DownloadHfAssetsOverrideFn) -> Self {
        let mut map = DOWNLOAD_HF_ASSETS_OVERRIDE.lock().unwrap();
        map.insert(label.clone(), func);
        DownloadHfAssetsOverrideGuard(label)
    }
}

#[cfg(test)]
pub(crate) fn set_download_hf_assets_label_override(
    label: String,
    func: DownloadHfAssetsLabelOverrideFn,
) -> DownloadHfAssetsOverrideGuard {
    DownloadHfAssetsOverrideGuard::set(
        label,
        Arc::new(move |label, assets| {
            if let Some(observer) = DOWNLOAD_PLAN_OBSERVER.lock().unwrap().clone() {
                let plan = initial_download_plan_for_assets(assets)?;
                observer(
                    label,
                    plan.into_iter()
                        .map(|(required, asset)| (required, asset.file))
                        .collect(),
                );
            }
            func(label)
        }),
    )
}

#[cfg(test)]
impl DownloadPlanObserverGuard {
    pub(crate) fn set(func: DownloadPlanObserverFn) -> Self {
        let mut slot = DOWNLOAD_PLAN_OBSERVER.lock().unwrap();
        *slot = Some(func);
        Self
    }
}

#[cfg(test)]
impl Drop for DownloadHfAssetsOverrideGuard {
    fn drop(&mut self) {
        let mut map = DOWNLOAD_HF_ASSETS_OVERRIDE.lock().unwrap();
        map.remove(&self.0);
    }
}

#[cfg(test)]
impl Drop for DownloadPlanObserverGuard {
    fn drop(&mut self) {
        let mut slot = DOWNLOAD_PLAN_OBSERVER.lock().unwrap();
        *slot = None;
    }
}

async fn download_hf_assets(
    label: &str,
    assets: Vec<HfAsset>,
    progress: bool,
) -> Result<Vec<PathBuf>> {
    let label = label.to_string();
    #[cfg(test)]
    {
        let func = DOWNLOAD_HF_ASSETS_OVERRIDE
            .lock()
            .unwrap()
            .get(&label)
            .cloned();
        if let Some(func) = func {
            return func(&label, assets);
        }
    }
    tokio::task::spawn_blocking(move || download_hf_assets_blocking(&label, assets, progress))
        .await
        .context("Join Hugging Face download task")?
}

struct MeshDownloadProgressState {
    filename: String,
    total: u64,
    downloaded: u64,
    bytes_per_sec: Option<f64>,
    last_draw: Option<std::time::Instant>,
}

struct MeshDownloadProgress {
    preflight_spinner: Mutex<Option<SpinnerHandle>>,
    state: Mutex<MeshDownloadProgressState>,
}

impl MeshDownloadProgress {
    fn new(filename: String) -> Self {
        let spinner_message = format!("Preparing download {}", filename);
        let preflight_spinner = if interactive_tui_active() {
            None
        } else {
            Some(start_spinner(&spinner_message))
        };
        Self {
            preflight_spinner: Mutex::new(preflight_spinner),
            state: Mutex::new(MeshDownloadProgressState {
                filename,
                total: 0,
                downloaded: 0,
                bytes_per_sec: None,
                last_draw: None,
            }),
        }
    }

    fn draw(state: &mut MeshDownloadProgressState, force: bool) {
        if !force && state.downloaded == 0 && state.total == 0 {
            return;
        }
        let now = std::time::Instant::now();
        if !force
            && state.last_draw.is_some_and(|last| {
                now.duration_since(last) < std::time::Duration::from_millis(150)
            })
        {
            return;
        }
        state.last_draw = Some(now);
        if interactive_tui_active() {
            emit_model_progress(
                &state.filename,
                Some(&state.filename),
                Some(state.downloaded),
                (state.total > 0).then_some(state.total),
                ModelProgressStatus::Downloading,
            );
            return;
        }
        let percent = if state.total == 0 {
            0
        } else {
            ((state.downloaded as f64 / state.total as f64) * 1000.0).round() as usize
        };
        let percent_major = (percent.min(1000)) / 10;
        let percent_minor = (percent.min(1000)) % 10;
        let speed_suffix = state
            .bytes_per_sec
            .filter(|bytes_per_sec| *bytes_per_sec > 0.0)
            .map(|bytes_per_sec| format!(" at {}/s", format_download_bytes(bytes_per_sec as u64)))
            .unwrap_or_default();
        eprint!(
            "\r\x1b[K   ⏬ {} {:>3}.{:01}% ({}/{}){}",
            state.filename,
            percent_major,
            percent_minor,
            format_download_bytes(state.downloaded),
            format_download_bytes(state.total),
            speed_suffix,
        );
        let _ = std::io::stderr().flush();
        if force {
            eprintln!();
        }
    }

    fn apply_download_event(state: &mut MeshDownloadProgressState, event: &DownloadEvent) {
        match event {
            DownloadEvent::Start { total_bytes, .. } => {
                if *total_bytes > 0 {
                    state.total = state.total.max(*total_bytes);
                }
            }
            DownloadEvent::Progress { files } => {
                if let Some(first) = files.first() {
                    if !first.filename.is_empty() {
                        state.filename = first.filename.clone();
                    }
                }
                if !files.is_empty() {
                    let reported_downloaded: u64 =
                        files.iter().map(|file| file.bytes_completed).sum();
                    state.downloaded = state.downloaded.max(reported_downloaded);
                    let reported_total: u64 = files.iter().map(|file| file.total_bytes).sum();
                    if reported_total > 0 {
                        state.total = state.total.max(reported_total);
                    }
                }
            }
            DownloadEvent::AggregateProgress {
                bytes_completed,
                total_bytes,
                bytes_per_sec,
            } => {
                state.downloaded = state.downloaded.max(*bytes_completed);
                if *total_bytes > 0 {
                    state.total = state.total.max(*total_bytes);
                }
                state.bytes_per_sec = *bytes_per_sec;
            }
            DownloadEvent::Complete => {
                if state.total > 0 {
                    state.downloaded = state.total;
                }
                state.bytes_per_sec = None;
            }
        }
    }

    fn showed_meaningful_progress(&self) -> bool {
        self.state
            .lock()
            .map(|state| state.downloaded > 0 || state.total > 0)
            .unwrap_or(false)
    }
}

impl ProgressHandler for MeshDownloadProgress {
    fn on_progress(&self, event: &ProgressEvent) {
        let ProgressEvent::Download(event) = event else {
            return;
        };
        let Ok(mut state) = self.state.lock() else {
            return;
        };
        Self::apply_download_event(&mut state, event);
        let should_show_progress = state.downloaded > 0 || state.total > 0;
        let force = matches!(event, DownloadEvent::Complete) && should_show_progress;
        if should_show_progress {
            if let Ok(mut spinner) = self.preflight_spinner.lock() {
                spinner.take();
            }
            Self::draw(&mut state, force);
        } else if matches!(event, DownloadEvent::Complete) {
            if let Ok(mut spinner) = self.preflight_spinner.lock() {
                spinner.take();
            }
        }
    }
}

impl Drop for MeshDownloadProgress {
    fn drop(&mut self) {
        if let Ok(mut spinner) = self.preflight_spinner.lock() {
            spinner.take();
        }
    }
}

fn format_download_bytes(bytes: u64) -> String {
    if bytes >= 1_000_000_000 {
        format!("{:.1}GB", bytes as f64 / 1e9)
    } else if bytes >= 1_000_000 {
        format!("{:.0}MB", bytes as f64 / 1e6)
    } else if bytes >= 1_000 {
        format!("{:.0}KB", bytes as f64 / 1e3)
    } else {
        format!("{bytes}B")
    }
}

fn emit_model_progress(
    label: &str,
    file: Option<&str>,
    downloaded_bytes: Option<u64>,
    total_bytes: Option<u64>,
    status: ModelProgressStatus,
) {
    let _ = emit_event(OutputEvent::ModelDownloadProgress {
        label: label.to_string(),
        file: file.map(ToOwned::to_owned),
        downloaded_bytes,
        total_bytes,
        status,
    });
}

fn emit_or_print_model_progress(
    label: &str,
    file: Option<&str>,
    downloaded_bytes: Option<u64>,
    total_bytes: Option<u64>,
    status: ModelProgressStatus,
    _fallback: impl FnOnce(),
) {
    emit_model_progress(label, file, downloaded_bytes, total_bytes, status);
}

fn download_hf_assets_blocking(
    label: &str,
    assets: Vec<HfAsset>,
    progress: bool,
) -> Result<Vec<PathBuf>> {
    let label = label.to_string();
    super::run_hf_sync(move || download_hf_assets_sync(&label, assets, progress))
}

fn download_hf_assets_sync(
    label: &str,
    assets: Vec<HfAsset>,
    progress: bool,
) -> Result<Vec<PathBuf>> {
    let api = super::build_hf_api(false)?;
    let mut download_plan = initial_download_plan_for_assets(assets)?;
    let current_plan: Vec<(bool, HfAsset)> = download_plan.iter().cloned().collect();
    for (_, asset) in current_plan {
        if !is_mlx_primary_asset(&asset.file) {
            continue;
        }
        for sidecar in mlx_sidecar_assets(&asset) {
            download_plan.insert(sidecar);
        }
        // Expand shards from an index file (downloads index to discover shard names)
        for shard in mlx_sharded_weight_assets(&api, &asset)? {
            download_plan.insert((true, shard));
        }
        // Expand shards from a first-shard ref without needing to download the index
        for shard in expand_split_mlx_first_shard(&asset) {
            download_plan.insert((true, shard));
        }
    }
    if progress {
        emit_or_print_model_progress(
            label,
            None,
            None,
            None,
            ModelProgressStatus::Ensuring,
            || eprintln!("📥 Ensuring {} is available locally...", label),
        );
    }

    #[cfg(test)]
    {
        if let Some(observer) = DOWNLOAD_PLAN_OBSERVER.lock().unwrap().clone() {
            observer(
                label,
                download_plan
                    .iter()
                    .map(|(required, asset)| (*required, asset.file.clone()))
                    .collect(),
            );
        }
    }

    let mut primary_paths = Vec::new();
    for (required, asset) in download_plan {
        let (owner, name) = asset.repo_parts();
        let api_repo = api.model(owner, name);
        if progress && required {
            emit_or_print_model_progress(
                label,
                Some(&asset.file),
                None,
                None,
                ModelProgressStatus::Ensuring,
                || eprintln!("   📥 Ensuring model {}", asset.file),
            );
        }
        let progress_tracker = if progress && required {
            Some(Arc::new(MeshDownloadProgress::new(asset.file.clone())))
        } else {
            None
        };
        let progress_handler: Progress = if let Some(tracker) = &progress_tracker {
            Some(tracker.clone())
        } else {
            None
        };
        let path = match api_repo.download_file(
            &RepoDownloadFileParams::builder()
                .filename(asset.file.clone())
                .revision(asset.revision.clone())
                .progress(progress_handler)
                .build(),
        ) {
            Ok(path) => {
                if progress {
                    if required {
                        let showed_progress = progress_tracker
                            .as_ref()
                            .is_some_and(|tracker| tracker.showed_meaningful_progress());
                        if showed_progress {
                            emit_or_print_model_progress(
                                label,
                                Some(&asset.file),
                                None,
                                None,
                                ModelProgressStatus::Ready,
                                || eprintln!("   ✅ Ready {}", asset.file),
                            );
                        } else if let Ok(meta) = std::fs::metadata(&path) {
                            emit_or_print_model_progress(
                                label,
                                Some(&asset.file),
                                Some(meta.len()),
                                Some(meta.len()),
                                ModelProgressStatus::Ready,
                                || {
                                    eprintln!(
                                        "   ✅ Ready {} ({})",
                                        asset.file,
                                        format_download_bytes(meta.len())
                                    )
                                },
                            );
                        } else {
                            emit_or_print_model_progress(
                                label,
                                Some(&asset.file),
                                None,
                                None,
                                ModelProgressStatus::Ready,
                                || eprintln!("   ✅ Ready {}", asset.file),
                            );
                        }
                    } else {
                        eprintln!("   🧾 Downloaded model metadata");
                    }
                }
                path
            }
            Err(_) if is_optional_metadata(required, &asset) => {
                continue;
            }
            Err(err) => {
                return Err(err).with_context(|| {
                    format!(
                        "Cache Hugging Face asset {}/{}@{}",
                        asset.repo, asset.file, asset.revision
                    )
                });
            }
        };
        if required && asset.file != "config.json" {
            primary_paths.push(path);
        }
    }

    Ok(primary_paths)
}

fn initial_download_plan_for_assets(
    assets: Vec<HfAsset>,
) -> Result<std::collections::BTreeSet<(bool, HfAsset)>> {
    let mut download_plan = std::collections::BTreeSet::new();
    let mut config_repos = std::collections::BTreeSet::new();

    for asset in assets {
        for expanded in expand_split_asset(&asset)? {
            config_repos.insert((expanded.repo.clone(), expanded.revision.clone()));
            download_plan.insert((true, expanded));
        }
    }

    for (repo, revision) in config_repos {
        download_plan.insert((
            false,
            HfAsset {
                repo,
                revision,
                file: "config.json".to_string(),
            },
        ));
    }

    Ok(download_plan)
}

#[cfg(test)]
pub async fn download_hf_repo_file(
    repo: &str,
    revision: Option<&str>,
    file: &str,
) -> Result<PathBuf> {
    download_hf_repo_file_with_progress(repo, revision, file, true).await
}

#[cfg(test)]
pub async fn download_hf_repo_file_with_progress(
    repo: &str,
    revision: Option<&str>,
    file: &str,
    progress: bool,
) -> Result<PathBuf> {
    download_hf_repo_file_with_progress_label(
        repo,
        revision,
        file,
        &format!("{repo}/{file}@{}", revision.unwrap_or("main")),
        progress,
    )
    .await
}

pub async fn download_hf_repo_file_with_progress_label(
    repo: &str,
    revision: Option<&str>,
    file: &str,
    label: &str,
    progress: bool,
) -> Result<PathBuf> {
    let revision = revision.unwrap_or("main").to_string();
    let asset = HfAsset {
        repo: repo.to_string(),
        revision: revision.clone(),
        file: file.to_string(),
    };
    let mut paths = download_hf_assets(label, vec![asset.clone()], progress).await?;
    paths.sort();
    let path = paths
        .into_iter()
        .find(|path| path_suffix_matches_ignore_case(path, &asset.file))
        .ok_or_else(|| {
            anyhow::anyhow!(
                "Downloaded Hugging Face asset not found in cache: {repo}/{file}@{revision}"
            )
        })?;
    let display_name = Path::new(&asset.file)
        .file_name()
        .and_then(|value| value.to_str())
        .unwrap_or(&asset.file);
    let model_ref = format!("{repo}@{revision}/{file}");
    if let Err(err) = track_managed_model_usage(
        &path,
        std::slice::from_ref(&path),
        display_name,
        Some(&model_ref),
        "huggingface",
    ) {
        tracing::warn!(
            "failed to record managed model usage for {}: {err}",
            path.display()
        );
    }
    Ok(path)
}

fn path_suffix_matches_ignore_case(path: &Path, expected: &str) -> bool {
    let expected_parts = expected
        .split(['/', '\\'])
        .filter(|part| !part.is_empty())
        .collect::<Vec<_>>();

    if expected_parts.is_empty() {
        return false;
    }

    let mut path_parts = path.iter().rev();

    for expected_part in expected_parts.iter().rev() {
        let Some(path_part) = path_parts.next() else {
            return false;
        };

        let Some(path_part) = path_part.to_str() else {
            return false;
        };

        if !path_part.eq_ignore_ascii_case(expected_part) {
            return false;
        }
    }

    true
}

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

    #[test]
    fn test_split_gguf_detection() {
        let re = regex_lite::Regex::new(r"-00001-of-(\d{5})\.gguf$").unwrap();

        // Should match split GGUFs
        let caps = re.captures("Model-Q4_K_M-00001-of-00004.gguf");
        assert!(caps.is_some());
        assert_eq!(&caps.unwrap()[1], "00004");

        let caps = re.captures("Qwen3-Coder-Next-Q4_K_M-00001-of-00004.gguf");
        assert!(caps.is_some());

        let caps = re.captures("MiniMax-M2.5-Q4_K_M-00001-of-00004.gguf");
        assert!(caps.is_some());

        // Should NOT match non-split or other parts
        assert!(re.captures("Model-Q4_K_M.gguf").is_none());
        assert!(re.captures("Model-Q4_K_M-00002-of-00004.gguf").is_none());
        assert!(re.captures("Model-Q4_K_M-00001-of-00004.bin").is_none());
    }

    #[test]
    fn test_split_url_generation() {
        let filename = "Model-Q4_K_M-00001-of-00003.gguf";
        let url = "https://huggingface.co/org/repo/resolve/main/Model-Q4_K_M-00001-of-00003.gguf";

        let mut files = Vec::new();
        for i in 1..=3u32 {
            let part_filename = filename.replace("-00001-of-", &format!("-{i:05}-of-"));
            let part_url = url.replace("-00001-of-", &format!("-{i:05}-of-"));
            files.push((part_filename, part_url));
        }

        assert_eq!(files.len(), 3);
        assert_eq!(files[0].0, "Model-Q4_K_M-00001-of-00003.gguf");
        assert_eq!(files[1].0, "Model-Q4_K_M-00002-of-00003.gguf");
        assert_eq!(files[2].0, "Model-Q4_K_M-00003-of-00003.gguf");
        assert!(files[0].1.contains("-00001-of-"));
        assert!(files[1].1.contains("-00002-of-"));
        assert!(files[2].1.contains("-00003-of-"));
    }

    #[test]
    fn path_file_name_matches_nested_path_ignore_case() {
        let path = Path::new("/tmp/cache/Subdir/Model.Q4_K_M.gguf");
        assert!(path_suffix_matches_ignore_case(
            path,
            "subdir/model.q4_k_m.gguf"
        ));
    }

    #[test]
    fn path_file_name_matches_rejects_wrong_suffix() {
        let path = Path::new("/tmp/cache/other/Model.Q4_K_M.gguf");
        assert!(!path_suffix_matches_ignore_case(
            path,
            "subdir/model.q4_k_m.gguf"
        ));
    }

    #[test]
    fn mlx_sidecars_include_required_tokenizer_and_optional_templates() {
        let asset = HfAsset {
            repo: "mlx-community/qwen2.5-0.5b-instruct-q2".to_string(),
            revision: "main".to_string(),
            file: "model.safetensors".to_string(),
        };
        let sidecars = mlx_sidecar_assets(&asset);
        assert_eq!(sidecars.len(), 4);
        assert!(sidecars[0].0);
        assert_eq!(sidecars[0].1.file, "tokenizer.json");
        assert!(sidecars
            .iter()
            .any(|(_, a)| a.file == "tokenizer_config.json"));
        assert!(sidecars
            .iter()
            .any(|(_, a)| a.file == "chat_template.jinja"));
        assert!(sidecars.iter().any(|(_, a)| a.file == "chat_template.json"));
    }

    #[test]
    fn parse_safetensors_index_shards_extracts_unique_shards() {
        let index = serde_json::json!({
            "weight_map": {
                "layer.0.q": "model-00001-of-00002.safetensors",
                "layer.0.k": "model-00001-of-00002.safetensors",
                "layer.1.q": "model-00002-of-00002.safetensors"
            }
        });
        let shards = parse_safetensors_index_shards(&index).unwrap();
        assert_eq!(
            shards,
            vec![
                "model-00001-of-00002.safetensors".to_string(),
                "model-00002-of-00002.safetensors".to_string()
            ]
        );
    }

    #[tokio::test]
    async fn download_hf_repo_file_matches_cache_file_case_insensitively() {
        let unique = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let cached_file = std::env::temp_dir()
            .join(format!("mesh-llm-hf-case-repo-{unique}"))
            .join("qwen2.5-coder-7b-instruct-q4_k_m.gguf");
        std::fs::create_dir_all(cached_file.parent().unwrap()).unwrap();
        std::fs::write(&cached_file, b"gguf").unwrap();

        {
            let label =
                "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF/Qwen2.5-Coder-7B-Instruct-Q4_K_M.gguf@main"
                    .to_string();
            let _guard = DownloadHfAssetsOverrideGuard::set(
                label,
                Arc::new({
                    let cached = cached_file.clone();
                    move |_, _| Ok(vec![cached.clone()])
                }),
            );
            let resolved = download_hf_repo_file(
                "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF",
                Some("main"),
                "Qwen2.5-Coder-7B-Instruct-Q4_K_M.gguf",
            )
            .await
            .unwrap();
            assert_eq!(resolved, cached_file);
        }

        {
            let label =
                "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF/Qwen2.5-Coder-7B-Instruct-Q4_K_M.gguf@main"
                    .to_string();
            let _guard = DownloadHfAssetsOverrideGuard::set(label, Arc::new(|_, _| Ok(Vec::new())));
            assert!(download_hf_repo_file(
                "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF",
                Some("main"),
                "Qwen2.5-Coder-7B-Instruct-Q4_K_M.gguf",
            )
            .await
            .is_err());
        }
    }

    #[tokio::test]
    async fn download_hf_repo_file_matches_nested_cache_path_case_insensitively() {
        let unique = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let cached_file = std::env::temp_dir()
            .join(format!("mesh-llm-hf-nested-repo-{unique}"))
            .join("nested")
            .join("Qwen2.5-Coder-7B-Instruct-Q4_K_M.gguf");
        std::fs::create_dir_all(cached_file.parent().unwrap()).unwrap();
        std::fs::write(&cached_file, b"gguf").unwrap();

        let label =
            "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF/nested/qwen2.5-coder-7b-instruct-q4_k_m.gguf@main"
                .to_string();
        let _guard = DownloadHfAssetsOverrideGuard::set(
            label,
            Arc::new({
                let cached = cached_file.clone();
                move |_, _| Ok(vec![cached.clone()])
            }),
        );

        let resolved = download_hf_repo_file(
            "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF",
            Some("main"),
            "nested/qwen2.5-coder-7b-instruct-q4_k_m.gguf",
        )
        .await
        .unwrap();
        assert_eq!(resolved, cached_file);
    }

    #[test]
    fn download_progress_state_merges_http_events_consistently() {
        let mut state = MeshDownloadProgressState {
            filename: "model.gguf".to_string(),
            total: 0,
            downloaded: 0,
            bytes_per_sec: None,
            last_draw: None,
        };

        MeshDownloadProgress::apply_download_event(
            &mut state,
            &DownloadEvent::Start {
                total_files: 1,
                total_bytes: 1_000,
            },
        );
        MeshDownloadProgress::apply_download_event(
            &mut state,
            &DownloadEvent::Progress {
                files: vec![hf_hub::FileProgress {
                    filename: "model.gguf".to_string(),
                    bytes_completed: 250,
                    total_bytes: 1_000,
                    status: hf_hub::FileStatus::InProgress,
                }],
            },
        );
        MeshDownloadProgress::apply_download_event(
            &mut state,
            &DownloadEvent::Progress {
                files: vec![hf_hub::FileProgress {
                    filename: "model.gguf".to_string(),
                    bytes_completed: 700,
                    total_bytes: 1_000,
                    status: hf_hub::FileStatus::InProgress,
                }],
            },
        );

        assert_eq!(state.filename, "model.gguf");
        assert_eq!(state.downloaded, 700);
        assert_eq!(state.total, 1_000);
        assert_eq!(state.bytes_per_sec, None);
    }

    #[test]
    fn download_progress_state_keeps_xet_progress_monotonic_when_per_file_lags() {
        let mut state = MeshDownloadProgressState {
            filename: "model.gguf".to_string(),
            total: 0,
            downloaded: 0,
            bytes_per_sec: None,
            last_draw: None,
        };

        MeshDownloadProgress::apply_download_event(
            &mut state,
            &DownloadEvent::AggregateProgress {
                bytes_completed: 32_000_000,
                total_bytes: 17_300_000_000,
                bytes_per_sec: Some(128_000_000.0),
            },
        );
        MeshDownloadProgress::apply_download_event(
            &mut state,
            &DownloadEvent::Progress {
                files: vec![hf_hub::FileProgress {
                    filename: "gemma-4-31B-it-Q4_0.gguf".to_string(),
                    bytes_completed: 4_000_000,
                    total_bytes: 17_300_000_000,
                    status: hf_hub::FileStatus::InProgress,
                }],
            },
        );

        assert_eq!(state.filename, "gemma-4-31B-it-Q4_0.gguf");
        assert_eq!(state.downloaded, 32_000_000);
        assert_eq!(state.total, 17_300_000_000);
        assert_eq!(state.bytes_per_sec, Some(128_000_000.0));
    }

    #[test]
    fn download_progress_state_clears_speed_and_finishes_at_total() {
        let mut state = MeshDownloadProgressState {
            filename: "model.gguf".to_string(),
            total: 1_000,
            downloaded: 700,
            bytes_per_sec: Some(42_000_000.0),
            last_draw: None,
        };

        MeshDownloadProgress::apply_download_event(&mut state, &DownloadEvent::Complete);

        assert_eq!(state.downloaded, 1_000);
        assert_eq!(state.total, 1_000);
        assert_eq!(state.bytes_per_sec, None);
    }

    #[test]
    fn is_split_mlx_first_shard_file_identifies_correct_patterns() {
        assert!(is_split_mlx_first_shard_file(
            "model-00001-of-00004.safetensors"
        ));
        assert!(is_split_mlx_first_shard_file(
            "model-00001-of-00048.safetensors"
        ));
        assert!(!is_split_mlx_first_shard_file(
            "model-00002-of-00004.safetensors"
        ));
        assert!(!is_split_mlx_first_shard_file("model.safetensors"));
        assert!(!is_split_mlx_first_shard_file(
            "model.safetensors.index.json"
        ));
        assert!(!is_split_mlx_first_shard_file("model-00001-of-00004.gguf"));
    }

    #[test]
    fn expand_split_mlx_first_shard_generates_all_shards() {
        let asset = HfAsset {
            repo: "org/repo".to_string(),
            revision: "main".to_string(),
            file: "model-00001-of-00003.safetensors".to_string(),
        };
        let shards = expand_split_mlx_first_shard(&asset);
        assert_eq!(shards.len(), 3);
        assert_eq!(shards[0].file, "model-00001-of-00003.safetensors");
        assert_eq!(shards[1].file, "model-00002-of-00003.safetensors");
        assert_eq!(shards[2].file, "model-00003-of-00003.safetensors");
        for shard in &shards {
            assert_eq!(shard.repo, "org/repo");
            assert_eq!(shard.revision, "main");
        }
    }

    #[test]
    fn expand_split_mlx_first_shard_returns_empty_for_non_first_shard() {
        let asset = HfAsset {
            repo: "org/repo".to_string(),
            revision: "main".to_string(),
            file: "model-00002-of-00003.safetensors".to_string(),
        };
        let shards = expand_split_mlx_first_shard(&asset);
        assert!(shards.is_empty());
    }

    #[test]
    fn gemma_bf16_first_shard_plans_full_split_download() {
        let plan = initial_download_plan_for_assets(vec![HfAsset {
            repo: "unsloth/gemma-4-31B-it-GGUF".to_string(),
            revision: "main".to_string(),
            file: "BF16/gemma-4-31B-it-BF16-00001-of-00002.gguf".to_string(),
        }])
        .unwrap();

        let files: Vec<_> = plan
            .into_iter()
            .map(|(required, asset)| (required, asset.file))
            .collect();

        assert_eq!(
            files,
            vec![
                (false, "config.json".to_string()),
                (
                    true,
                    "BF16/gemma-4-31B-it-BF16-00001-of-00002.gguf".to_string()
                ),
                (
                    true,
                    "BF16/gemma-4-31B-it-BF16-00002-of-00002.gguf".to_string()
                ),
            ]
        );
    }

    #[test]
    fn is_mlx_primary_asset_includes_first_shard() {
        assert!(is_mlx_primary_asset("model.safetensors"));
        assert!(is_mlx_primary_asset("model.safetensors.index.json"));
        assert!(is_mlx_primary_asset("model-00001-of-00048.safetensors"));
        assert!(!is_mlx_primary_asset("model-00002-of-00048.safetensors"));
        assert!(!is_mlx_primary_asset("model-00048-of-00048.safetensors"));
    }
}
