use serde::Deserialize;

#[derive(Debug, Clone, Deserialize)]
pub struct LlamaConfig {
    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 hidden_act: String,
    pub rms_norm_eps: f64,
    pub rope_theta: f64,
    pub max_position_embeddings: usize,
    pub bos_token_id: u32,
    pub eos_token_id: EosTokenId,
    pub pad_token_id: u32,
    pub tie_word_embeddings: bool,
    pub head_dim: Option<usize>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub enum EosTokenId {
    Single(u32),
    Multiple(Vec<u32>),
}

impl EosTokenId {
    pub fn first(&self) -> u32 {
        match self {
            EosTokenId::Single(id) => *id,
            EosTokenId::Multiple(ids) => ids[0],
        }
    }

    pub fn contains(&self, token: u32) -> bool {
        match self {
            EosTokenId::Single(id) => *id == token,
            EosTokenId::Multiple(ids) => ids.contains(&token),
        }
    }

    pub fn as_slice(&self) -> &[u32] {
        match self {
            EosTokenId::Single(id) => std::slice::from_ref(id),
            EosTokenId::Multiple(ids) => ids.as_slice(),
        }
    }
}

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

    pub fn from_json(path: &str) -> anyhow::Result<Self> {
        let content = std::fs::read_to_string(path)?;
        let config: LlamaConfig = serde_json::from_str(&content)?;
        Ok(config)
    }
}