//! LLMHosts model configuration extracted from GGUF metadata.

use crate::gguf::{self, GgufHeader, GgufError};

#[derive(Debug, Clone)]
pub struct LLMHostsModelConfig {
    pub architecture: String,
    pub hidden_size: usize,
    pub num_attention_heads: usize,
    pub num_kv_heads: usize,
    pub num_layers: usize,
    pub intermediate_size: usize,
    pub vocab_size: usize,
    pub max_seq_len: usize,
    pub head_dim: usize,
    pub rope_freq_base: f32,
    pub rms_norm_eps: f32,
    pub eos_token_id: Option<u32>,
    pub bos_token_id: Option<u32>,
}

impl LLMHostsModelConfig {
    /// Extract model config from GGUF metadata.
    pub fn from_gguf(header: &GgufHeader) -> Result<Self, GgufError> {
        let arch = gguf::get_string(header, "general.architecture")
            .ok_or_else(|| GgufError::ConfigError("missing general.architecture".into()))?;

        let hidden_size = gguf::get_u32(header, &format!("{arch}.embedding_length"))
            .ok_or_else(|| GgufError::ConfigError("missing embedding_length".into()))? as usize;
        let num_attention_heads = gguf::get_u32(header, &format!("{arch}.attention.head_count"))
            .ok_or_else(|| GgufError::ConfigError("missing head_count".into()))? as usize;
        let num_kv_heads = gguf::get_u32(header, &format!("{arch}.attention.head_count_kv"))
            .unwrap_or(num_attention_heads as u32) as usize;
        let num_layers = gguf::get_u32(header, &format!("{arch}.block_count"))
            .ok_or_else(|| GgufError::ConfigError("missing block_count".into()))? as usize;
        let intermediate_size = gguf::get_u32(header, &format!("{arch}.feed_forward_length"))
            .ok_or_else(|| GgufError::ConfigError("missing feed_forward_length".into()))? as usize;
        let max_seq_len = gguf::get_u32(header, &format!("{arch}.context_length"))
            .unwrap_or(2048) as usize;

        // Vocab size: second dimension of token_embd.weight
        // GGUF stores 2D tensor dims as [cols, rows] (column-major)
        // token_embd.weight is [hidden_size, vocab_size]
        let vocab_size = header.tensors.iter()
            .find(|t| t.name == "token_embd.weight")
            .map(|t| if t.dimensions.len() >= 2 { t.dimensions[1] as usize } else { t.dimensions[0] as usize })
            .ok_or_else(|| GgufError::ConfigError("missing token_embd.weight".into()))?;

        let head_dim = hidden_size / num_attention_heads;

        let rope_freq_base = gguf::get_f32(header, &format!("{arch}.rope.freq_base"))
            .unwrap_or(10000.0);
        let rms_norm_eps = gguf::get_f32(header, &format!("{arch}.attention.layer_norm_rms_epsilon"))
            .unwrap_or(1e-5);

        let eos_token_id = gguf::get_u32(header, "tokenizer.ggml.eos_token_id");
        let bos_token_id = gguf::get_u32(header, "tokenizer.ggml.bos_token_id");

        Ok(Self {
            architecture: arch,
            hidden_size,
            num_attention_heads,
            num_kv_heads,
            num_layers,
            intermediate_size,
            vocab_size,
            max_seq_len,
            head_dim,
            rope_freq_base,
            rms_norm_eps,
            eos_token_id,
            bos_token_id,
        })
    }
}

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

    #[test]
    fn test_config_from_test_model() {
        let dir = tempfile::tempdir().unwrap();
        let gen_config = crate::test_utils::TestModelConfig::default();
        let path = crate::test_utils::generate_test_gguf(dir.path(), &gen_config);
        let header = crate::gguf::parse_header(&path).unwrap();
        let config = LLMHostsModelConfig::from_gguf(&header).unwrap();

        assert_eq!(config.architecture, "llama");
        assert_eq!(config.hidden_size, 64);
        assert_eq!(config.num_attention_heads, 4);
        assert_eq!(config.num_kv_heads, 2);
        assert_eq!(config.num_layers, 2);
        assert_eq!(config.intermediate_size, 128);
        assert_eq!(config.vocab_size, 128);
        assert_eq!(config.head_dim, 16); // 64 / 4
    }

    #[test]
    fn test_config_defaults() {
        let dir = tempfile::tempdir().unwrap();
        let gen_config = crate::test_utils::TestModelConfig::default();
        let path = crate::test_utils::generate_test_gguf(dir.path(), &gen_config);
        let header = crate::gguf::parse_header(&path).unwrap();
        let config = LLMHostsModelConfig::from_gguf(&header).unwrap();

        // These should use defaults since not in test GGUF metadata
        assert_eq!(config.rope_freq_base, 10000.0);
        assert_eq!(config.rms_norm_eps, 1e-5);
    }
}
