use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::path::Path;

/// The subset of Hugging Face `config.json` needed for Llama.
#[derive(Debug, Clone, Deserialize)]
pub struct HfConfig {
    pub hidden_size: usize,
    pub intermediate_size: usize,
    pub num_hidden_layers: usize,
    pub num_attention_heads: usize,
    pub num_key_value_heads: usize,
    pub vocab_size: usize,
    pub rms_norm_eps: f64,
    pub rope_theta: f32,
    pub max_position_embeddings: usize,
    #[serde(default)]
    pub tie_word_embeddings: bool,
}

impl HfConfig {
    pub fn load(dir: impl AsRef<Path>) -> Result<Self> {
        let path = dir.as_ref().join("config.json");
        let raw = std::fs::read_to_string(&path)
            .map_err(|e| anyhow::anyhow!("{}: {e}", path.display()))?;
        let cfg: Self = serde_json::from_str(&raw)?;
        // the model has no separate lm_head: it reuses the embedding matrix
        if !cfg.tie_word_embeddings {
            anyhow::bail!("only models with tie_word_embeddings = true are supported");
        }
        Ok(cfg)
    }

    pub fn head_dim(&self) -> usize {
        self.hidden_size / self.num_attention_heads
    }

    /// Longest sequence the RoPE table covers. Capped so the table stays
    /// small on CPU; SmolLM2 declares 8192.
    pub fn max_seq(&self) -> usize {
        self.max_position_embeddings.min(2048)
    }
}

/// LoRA hyperparameters. Saved next to the adapter because the rank
/// defines the shape of A and B.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoraConfig {
    pub rank: usize,
    pub alpha: f64,
    /// Projections that get an adapter.
    pub targets: Vec<String>,
}

impl Default for LoraConfig {
    fn default() -> Self {
        Self {
            rank: 8,
            alpha: 16.0,
            targets: vec!["q".into(), "v".into()],
        }
    }
}

impl LoraConfig {
    pub fn scaling(&self) -> f64 {
        self.alpha / self.rank as f64
    }

    pub fn targets(&self, name: &str) -> bool {
        self.targets.iter().any(|t| t == name)
    }
}