xz-rerank 0.1.1

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::fmt::Debug;

use crate::error::RerankError;
use crate::types::{RerankCandidate, RerankConfig, RerankResult};

/// 重排序器后端类型
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RerankerBackendType {
    /// 本地信号融合重排序
    Local,
    /// 远程 API 重排序
    Remote,
    /// 插件式重排序
    Plugin,
}

/// 重排序器信息
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RerankerInfo {
    /// 内部名称标识
    pub name: String,
    /// 对外展示名称
    pub display_name: String,
    /// 后端类型
    pub backend_type: RerankerBackendType,
    /// 是否支持批量重排序
    pub supports_batch: bool,
    /// 最大支持的候选项数（`None` 表示无限制）
    pub max_candidates: Option<usize>,
    /// 定价信息
    pub pricing: Option<RerankerPricing>,
}

/// 重排序器定价
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RerankerPricing {
    /// 单次搜索的成本（美元）
    pub cost_per_search: f64,
}

/// 统一的重排序接口
#[async_trait]
pub trait Reranker: Send + Sync + Debug {
    /// 对候选项重新排序
    async fn rerank(
        &self,
        query: &str,
        candidates: Vec<RerankCandidate>,
        config: &RerankConfig,
    ) -> Result<RerankResult, RerankError>;

    /// 对候选项重新排序（简化版，使用默认配置）
    async fn rerank_default(
        &self,
        query: &str,
        candidates: Vec<RerankCandidate>,
    ) -> Result<RerankResult, RerankError> {
        self.rerank(query, candidates, &RerankConfig::default()).await
    }

    /// 重排序器信息
    fn reranker_info(&self) -> &RerankerInfo;
}

/// 信号插件 trait — 可自定义打分信号
#[async_trait]
pub trait SignalPlugin: Send + Sync + Debug {
    /// 信号名称（用于分数分解）
    fn name(&self) -> &str;

    /// 权重键名（对应 SignalWeights 中的字段名，用于名称映射）
    fn weight_key(&self) -> &'static str;

    /// 对单个候选项打分（返回 [0, 1] 区间的分数）
    async fn score(&self, query: &str, candidate: &RerankCandidate) -> Result<f32, RerankError>;

    /// 批量打分
    async fn score_batch(
        &self,
        query: &str,
        candidates: &[RerankCandidate],
    ) -> Result<Vec<f32>, RerankError> {
        use futures::future::join_all;

        let futures: Vec<_> =
            candidates.iter().map(|candidate| self.score(query, candidate)).collect();
        let results = join_all(futures).await;

        let mut scores = Vec::with_capacity(results.len());
        for result in results {
            scores.push(result?);
        }

        Ok(scores)
    }
}