//! EmbeddingGemma 模型配置
//!
//! 定义 EmbeddingGemma 模型的架构参数和配置选项。

use serde::{Deserialize, Serialize};

/// EmbeddingGemma 模型配置
///
/// 基于 Google EmbeddingGemma-300M 的默认配置。
/// 该模型使用 Gemma3 架构，但修改为双向注意力用于 embedding 任务。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingGemmaConfig {
    /// 词汇表大小 (default: 262144 for EmbeddingGemma)
    #[serde(default = "default_vocab_size")]
    pub vocab_size: usize,

    /// 隐藏层维度 (default: 768)
    #[serde(default = "default_hidden_size")]
    pub hidden_size: usize,

    /// FFN 中间层维度 (default: 1152 for EmbeddingGemma)
    #[serde(default = "default_intermediate_size")]
    pub intermediate_size: usize,

    /// Transformer 层数 (default: 24)
    #[serde(default = "default_num_hidden_layers")]
    pub num_hidden_layers: usize,

    /// Query 注意力头数 (default: 3 for EmbeddingGemma)
    #[serde(default = "default_num_attention_heads")]
    pub num_attention_heads: usize,

    /// Key-Value 注意力头数，用于 GQA (default: 1 for EmbeddingGemma)
    #[serde(default = "default_num_key_value_heads")]
    pub num_key_value_heads: usize,

    /// 每个注意力头的维度 (default: 256 for EmbeddingGemma)
    #[serde(default = "default_head_dim")]
    pub head_dim: usize,

    /// 最大位置编码长度 (default: 2048)
    #[serde(default = "default_max_position_embeddings")]
    pub max_position_embeddings: usize,

    /// RMSNorm epsilon (default: 1e-6)
    #[serde(default = "default_rms_norm_eps")]
    pub rms_norm_eps: f64,

    /// RoPE theta 参数 (default: 1000000.0 for EmbeddingGemma)
    #[serde(default = "default_rope_theta")]
    pub rope_theta: f64,

    /// 是否使用注意力偏置 (default: false)
    #[serde(default)]
    pub attention_bias: bool,

    /// 投影层中间维度 (default: 0, 表示不使用投影层)
    /// 用于最终的 embedding 投影: hidden_size -> projection_dim -> hidden_size
    #[serde(default)]
    pub projection_dim: usize,

    /// Query pre-attention scalar (default: 256)
    #[serde(default = "default_query_pre_attn_scalar")]
    pub query_pre_attn_scalar: usize,
}

fn default_vocab_size() -> usize { 262144 }
fn default_hidden_size() -> usize { 768 }
fn default_intermediate_size() -> usize { 1152 }
fn default_num_hidden_layers() -> usize { 24 }
fn default_num_attention_heads() -> usize { 3 }
fn default_num_key_value_heads() -> usize { 1 }
fn default_head_dim() -> usize { 256 }
fn default_max_position_embeddings() -> usize { 2048 }
fn default_rms_norm_eps() -> f64 { 1e-6 }
fn default_rope_theta() -> f64 { 1000000.0 }
fn default_query_pre_attn_scalar() -> usize { 256 }

impl Default for EmbeddingGemmaConfig {
    fn default() -> Self {
        Self {
            vocab_size: 262144,
            hidden_size: 768,
            intermediate_size: 1152,
            num_hidden_layers: 24,
            num_attention_heads: 3,
            num_key_value_heads: 1,
            head_dim: 256,
            max_position_embeddings: 2048,
            rms_norm_eps: 1e-6,
            rope_theta: 1000000.0,
            attention_bias: false,
            projection_dim: 0, // EmbeddingGemma doesn't use projection by default
            query_pre_attn_scalar: 256,
        }
    }
}

impl EmbeddingGemmaConfig {
    /// 创建一个新的配置
    pub fn new() -> Self {
        Self::default()
    }

    /// 从 JSON 字符串加载配置
    pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
        serde_json::from_str(json)
    }

    /// 转换为 JSON 字符串
    pub fn to_json(&self) -> Result<String, serde_json::Error> {
        serde_json::to_string_pretty(self)
    }

    /// 设置隐藏层维度
    pub fn with_hidden_size(mut self, hidden_size: usize) -> Self {
        self.hidden_size = hidden_size;
        self
    }

    /// 设置层数
    pub fn with_num_layers(mut self, num_layers: usize) -> Self {
        self.num_hidden_layers = num_layers;
        self
    }

    /// 设置最大位置编码长度
    pub fn with_max_position_embeddings(mut self, max_pos: usize) -> Self {
        self.max_position_embeddings = max_pos;
        self
    }

    /// 获取 Query 投影维度 (num_attention_heads * head_dim)
    pub fn q_dim(&self) -> usize {
        self.num_attention_heads * self.head_dim
    }

    /// 获取 KV 投影维度 (num_key_value_heads * head_dim)
    pub fn kv_dim(&self) -> usize {
        self.num_key_value_heads * self.head_dim
    }

    /// 验证配置的一致性
    pub fn validate(&self) -> Result<(), String> {
        // EmbeddingGemma uses GQA where hidden_size = num_attention_heads * head_dim
        if self.hidden_size != self.num_attention_heads * self.head_dim {
            return Err(format!(
                "hidden_size ({}) must equal num_attention_heads ({}) * head_dim ({})",
                self.hidden_size, self.num_attention_heads, self.head_dim
            ));
        }

        if self.vocab_size == 0 {
            return Err("vocab_size must be greater than 0".to_string());
        }

        if self.num_hidden_layers == 0 {
            return Err("num_hidden_layers must be greater than 0".to_string());
        }

        if self.num_key_value_heads > self.num_attention_heads {
            return Err(format!(
                "num_key_value_heads ({}) must be <= num_attention_heads ({})",
                self.num_key_value_heads, self.num_attention_heads
            ));
        }

        Ok(())
    }
}

/// 支持的 Matryoshka 截断维度
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MatryoshkaDim {
    /// 完整维度 (768)
    Full,
    /// 512 维
    Dim512,
    /// 256 维
    Dim256,
    /// 128 维
    Dim128,
}

impl MatryoshkaDim {
    /// 获取实际的维度值
    pub fn dim(&self) -> usize {
        match self {
            MatryoshkaDim::Full => 768,
            MatryoshkaDim::Dim512 => 512,
            MatryoshkaDim::Dim256 => 256,
            MatryoshkaDim::Dim128 => 128,
        }
    }

    /// 从维度值创建
    pub fn from_dim(dim: usize) -> Option<Self> {
        match dim {
            768 => Some(MatryoshkaDim::Full),
            512 => Some(MatryoshkaDim::Dim512),
            256 => Some(MatryoshkaDim::Dim256),
            128 => Some(MatryoshkaDim::Dim128),
            _ => None,
        }
    }
}

impl Default for MatryoshkaDim {
    fn default() -> Self {
        MatryoshkaDim::Full
    }
}

/// 任务类型前缀
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TaskType {
    /// 搜索查询
    SearchQuery,
    /// 搜索文档
    SearchDocument,
    /// 聚类
    Clustering,
    /// 分类
    Classification,
    /// 语义相似度
    SemanticSimilarity,
    /// 摘要
    Summarization,
}

impl TaskType {
    /// 获取任务对应的前缀字符串
    pub fn prefix(&self) -> &'static str {
        match self {
            TaskType::SearchQuery => "task: search result | query: ",
            TaskType::SearchDocument => "title: none | text: ",
            TaskType::Clustering => "task: clustering | query: ",
            TaskType::Classification => "task: classification | query: ",
            TaskType::SemanticSimilarity => "task: sentence similarity | query: ",
            TaskType::Summarization => "task: summarization | query: ",
        }
    }

    /// 从字符串解析任务类型
    pub fn from_str(s: &str) -> Option<Self> {
        match s.to_lowercase().as_str() {
            "query" | "search_query" | "retrieval-query" => Some(TaskType::SearchQuery),
            "document" | "search_document" | "retrieval-document" => Some(TaskType::SearchDocument),
            "clustering" => Some(TaskType::Clustering),
            "classification" => Some(TaskType::Classification),
            "sts" | "similarity" | "semantic_similarity" => Some(TaskType::SemanticSimilarity),
            "summarization" => Some(TaskType::Summarization),
            _ => None,
        }
    }
}

impl Default for TaskType {
    fn default() -> Self {
        TaskType::SearchQuery
    }
}

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

    #[test]
    fn test_default_config() {
        let config = EmbeddingGemmaConfig::default();
        // EmbeddingGemma actual model values
        assert_eq!(config.vocab_size, 262144);
        assert_eq!(config.hidden_size, 768);
        assert_eq!(config.intermediate_size, 1152);
        assert_eq!(config.num_hidden_layers, 24);
        assert_eq!(config.num_attention_heads, 3);
        assert_eq!(config.num_key_value_heads, 1);
        assert_eq!(config.head_dim, 256);
        assert_eq!(config.max_position_embeddings, 2048);
        assert_eq!(config.rope_theta, 1000000.0);
        assert_eq!(config.projection_dim, 0);
    }

    #[test]
    fn test_config_validation() {
        let config = EmbeddingGemmaConfig::default();
        assert!(config.validate().is_ok());

        let invalid_config = EmbeddingGemmaConfig {
            hidden_size: 512, // 不匹配 12 * 64 = 768
            ..Default::default()
        };
        assert!(invalid_config.validate().is_err());
    }

    #[test]
    fn test_matryoshka_dim() {
        assert_eq!(MatryoshkaDim::Full.dim(), 768);
        assert_eq!(MatryoshkaDim::Dim512.dim(), 512);
        assert_eq!(MatryoshkaDim::Dim256.dim(), 256);
        assert_eq!(MatryoshkaDim::Dim128.dim(), 128);

        assert_eq!(MatryoshkaDim::from_dim(768), Some(MatryoshkaDim::Full));
        assert_eq!(MatryoshkaDim::from_dim(512), Some(MatryoshkaDim::Dim512));
        assert_eq!(MatryoshkaDim::from_dim(100), None);
    }

    #[test]
    fn test_task_type() {
        assert_eq!(TaskType::SearchQuery.prefix(), "task: search result | query: ");
        assert_eq!(TaskType::SearchDocument.prefix(), "title: none | text: ");

        assert_eq!(TaskType::from_str("query"), Some(TaskType::SearchQuery));
        assert_eq!(TaskType::from_str("document"), Some(TaskType::SearchDocument));
        assert_eq!(TaskType::from_str("unknown"), None);
    }

    #[test]
    fn test_config_json() {
        let config = EmbeddingGemmaConfig::default();
        let json = config.to_json().unwrap();
        let parsed = EmbeddingGemmaConfig::from_json(&json).unwrap();
        assert_eq!(parsed.hidden_size, config.hidden_size);
        assert_eq!(parsed.num_hidden_layers, config.num_hidden_layers);
    }
}
