//! HuggingFace Hub asset resolution.
//!
//! Resolves the three weight artifacts CSM needs:
//!   - sesame/csm-1b: model.safetensors (CSM dual-Llama weights)
//!   - kyutai/mimi:   model.safetensors (Mimi codec weights)
//!   - meta-llama/Llama-3.2-1B: tokenizer.json (BPE vocab)
//!
//! Llama tokenizer requires HF token acceptance; document via `HF_TOKEN` env var.

use crate::error::Result;
use hf_hub::api::sync::Api;
use std::path::PathBuf;

pub const REPO_CSM_1B: &str = "sesame/csm-1b";
pub const REPO_MIMI: &str = "kyutai/mimi";
pub const REPO_LLAMA_TOKENIZER: &str = "meta-llama/Llama-3.2-1B";
/// Public mirror — tokenizer.json is byte-identical to meta-llama's. Used as
/// fallback when the user hasn't been approved on Meta's gated form yet.
pub const REPO_LLAMA_TOKENIZER_FALLBACK: &str = "unsloth/Llama-3.2-1B";
/// AudioSeal watermark — public, no auth required.
pub const REPO_AUDIOSEAL: &str = "facebook/audioseal";
pub const FILE_AUDIOSEAL_GENERATOR: &str = "generator_base.pth";
pub const FILE_AUDIOSEAL_DETECTOR: &str = "detector_base.pth";
/// WavLM-Base+ Speaker Verification (X-vector head). Public, no auth.
pub const REPO_WAVLM_SV: &str = "microsoft/wavlm-base-plus-sv";
pub const FILE_WAVLM_SV_PT: &str = "pytorch_model.bin";

#[derive(Debug, Clone)]
pub struct CsmAssets {
    pub csm_weights: PathBuf,
    pub mimi_weights: PathBuf,
    pub tokenizer_json: PathBuf,
}

/// Resolve all three CSM-1B assets. Requires HF auth for `sesame/csm-1b` and
/// `meta-llama/Llama-3.2-1B`. Use `resolve_mimi()` alone if you only need the
/// codec (it's public).
pub fn resolve_csm_1b() -> Result<CsmAssets> {
    let csm_weights = resolve_csm_weights()?;
    let mimi_weights = resolve_mimi()?;
    let tokenizer_json = resolve_llama_tokenizer()?;
    Ok(CsmAssets {
        csm_weights,
        mimi_weights,
        tokenizer_json,
    })
}

/// Download `sesame/csm-1b/model.safetensors`. Requires HF auth + accepted repo terms.
pub fn resolve_csm_weights() -> Result<PathBuf> {
    let api = Api::new()?;
    Ok(api
        .model(REPO_CSM_1B.to_string())
        .get("model.safetensors")?)
}

/// Download `kyutai/mimi/model.safetensors`. Public, no auth required.
pub fn resolve_mimi() -> Result<PathBuf> {
    let api = Api::new()?;
    Ok(api.model(REPO_MIMI.to_string()).get("model.safetensors")?)
}

/// Download `facebook/audioseal/generator_base.pth`. Public, no auth required.
pub fn resolve_audioseal_generator() -> Result<PathBuf> {
    let api = Api::new()?;
    Ok(api
        .model(REPO_AUDIOSEAL.to_string())
        .get(FILE_AUDIOSEAL_GENERATOR)?)
}

/// Download `facebook/audioseal/detector_base.pth`. Public, no auth required.
pub fn resolve_audioseal_detector() -> Result<PathBuf> {
    let api = Api::new()?;
    Ok(api
        .model(REPO_AUDIOSEAL.to_string())
        .get(FILE_AUDIOSEAL_DETECTOR)?)
}

/// Download `microsoft/wavlm-base-plus-sv/pytorch_model.bin`. Public, no auth.
pub fn resolve_wavlm_sv() -> Result<PathBuf> {
    let api = Api::new()?;
    Ok(api.model(REPO_WAVLM_SV.to_string()).get(FILE_WAVLM_SV_PT)?)
}

/// Download Llama-3.2-1B `tokenizer.json`.
///
/// Tries the canonical `meta-llama/Llama-3.2-1B` first; on 403 (terms not
/// accepted) falls back to the public `unsloth/Llama-3.2-1B` mirror, which
/// hosts a byte-identical tokenizer.json.
pub fn resolve_llama_tokenizer() -> Result<PathBuf> {
    let api = Api::new()?;
    match api
        .model(REPO_LLAMA_TOKENIZER.to_string())
        .get("tokenizer.json")
    {
        Ok(p) => Ok(p),
        Err(e) => {
            tracing::warn!(
                "{REPO_LLAMA_TOKENIZER} tokenizer fetch failed ({e}); falling back to {REPO_LLAMA_TOKENIZER_FALLBACK}"
            );
            Ok(api
                .model(REPO_LLAMA_TOKENIZER_FALLBACK.to_string())
                .get("tokenizer.json")?)
        }
    }
}