xz-rerank 0.1.1

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::error::RerankError;

/// 信号权重配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SignalWeights {
    /// 关键词重叠度权重
    pub keyword_overlap: f32,
    /// 向量相似度权重
    pub vector_similarity: f32,
    /// 元数据匹配权重
    pub metadata_match: f32,
    /// 内容质量权重
    pub content_quality: f32,
    /// 时间近因性权重
    pub recency: f32,
    /// 自定义信号权重（按信号名称索引）
    #[serde(default)]
    pub custom_weights: HashMap<String, f32>,
}

impl Default for SignalWeights {
    /// 返回默认权重配置
    ///
    /// 权重分配：keyword_overlap=0.30, vector_similarity=0.25,
    /// metadata_match=0.20, content_quality=0.10, recency=0.15
    fn default() -> Self {
        Self {
            keyword_overlap: 0.30,
            vector_similarity: 0.25,
            metadata_match: 0.20,
            content_quality: 0.10,
            recency: 0.15,
            custom_weights: HashMap::new(),
        }
    }
}

impl SignalWeights {
    /// 验证权重和为 1.0
    pub fn validate(&self) -> Result<(), RerankError> {
        let sum = self.keyword_overlap
            + self.vector_similarity
            + self.metadata_match
            + self.content_quality
            + self.recency
            + self.custom_weights.values().sum::<f32>();

        if (sum - 1.0).abs() > 0.01 {
            return Err(RerankError::WeightSumInvalid(sum));
        }
        Ok(())
    }

    /// 按名称查找权重（用于信号名称映射）
    ///
    /// 支持的名称：`keyword_overlap`, `vector_similarity`, `metadata_match`,
    /// `content_quality`, `recency`。如果以上均不匹配，则查询自定义权重；
    /// 未找到返回 0.0。
    pub fn get_weight_by_name(&self, name: &str) -> f32 {
        match name {
            "keyword_overlap" => self.keyword_overlap,
            "vector_similarity" => self.vector_similarity,
            "metadata_match" => self.metadata_match,
            "content_quality" => self.content_quality,
            "recency" => self.recency,
            _ => self.custom_weights.get(name).copied().unwrap_or(0.0),
        }
    }

    /// 设置自定义信号权重并返回 self
    pub fn with_custom_weight(mut self, key: &str, weight: f32) -> Self {
        self.custom_weights.insert(key.to_string(), weight);
        self
    }

    /// 自动归一化所有权重（含自定义），使其和为 1.0
    ///
    /// 如果当前权重之和为 0，则返回默认权重配置。
    pub fn normalize(&self) -> Self {
        let sum = self.keyword_overlap
            + self.vector_similarity
            + self.metadata_match
            + self.content_quality
            + self.recency
            + self.custom_weights.values().sum::<f32>();

        if sum == 0.0 {
            return Self::default();
        }

        let custom_weights =
            self.custom_weights.iter().map(|(k, v)| (k.clone(), v / sum)).collect();

        Self {
            keyword_overlap: self.keyword_overlap / sum,
            vector_similarity: self.vector_similarity / sum,
            metadata_match: self.metadata_match / sum,
            content_quality: self.content_quality / sum,
            recency: self.recency / sum,
            custom_weights,
        }
    }
}