xz-rerank 0.1.1

use async_trait::async_trait;
use std::collections::HashMap;
use std::time::Instant;
use tokio::sync::Mutex;

use crate::error::RerankError;
use crate::traits::{Reranker, RerankerBackendType, RerankerInfo};
use crate::types::{RerankCandidate, RerankConfig, RerankHit, RerankResult, RerankStats};

/// 测试用 Mock Reranker
#[derive(Debug)]
pub struct MockReranker {
    info: RerankerInfo,
    mock_result: Mutex<Option<RerankResult>>,
    should_error: Mutex<Option<RerankError>>,
}

impl MockReranker {
    /// 创建 Mock Reranker
    ///
    /// `name` 用于标识此 mock 实例，会体现在 [`RerankerInfo`] 中。
    pub fn new(name: &str) -> Self {
        Self {
            info: RerankerInfo {
                name: name.to_string(),
                display_name: format!("Mock {name}"),
                backend_type: RerankerBackendType::Local,
                supports_batch: true,
                max_candidates: None,
                pricing: None,
            },
            mock_result: Mutex::new(None),
            should_error: Mutex::new(None),
        }
    }

    /// 设置模拟返回结果
    ///
    /// 调用后，下次 `rerank` 调用将返回该预设结果。
    /// 可用于测试特定场景下的重排序行为。
    pub fn set_result(&mut self, result: RerankResult) {
        *self.mock_result.get_mut() = Some(result);
    }

    /// 设置模拟错误
    ///
    /// 调用后，下次 `rerank` 调用将返回该预设错误。
    /// 可用于测试错误处理逻辑。
    pub fn set_error(&mut self, error: RerankError) {
        *self.should_error.get_mut() = Some(error);
    }
}

#[async_trait]
impl Reranker for MockReranker {
    async fn rerank(
        &self,
        _query: &str,
        candidates: Vec<RerankCandidate>,
        config: &RerankConfig,
    ) -> Result<RerankResult, RerankError> {
        if let Some(ref err) = *self.should_error.lock().await {
            return Err(RerankError::Engine(format!("Mock error: {err}")));
        }

        if let Some(ref result) = *self.mock_result.lock().await {
            return Ok(result.clone());
        }

        // 默认行为：原样返回，分数 = retrieval_score or 0.5
        let now = Instant::now();
        let hits: Vec<RerankHit> = candidates
            .iter()
            .take(config.top_k)
            .map(|c| RerankHit {
                candidate_id: c.id.clone(),
                score: c.retrieval_score.unwrap_or(0.5),
                score_breakdown: None,
                candidate: c.clone(),
            })
            .collect();

        Ok(RerankResult {
            hits,
            stats: RerankStats {
                total_candidates: candidates.len(),
                filtered_out: 0,
                max_score: 1.0,
                min_score: 0.0,
                avg_score: 0.5,
                median_score: 0.5,
                signal_timings: HashMap::new(),
            },
            reranker: self.info.name.clone(),
            latency_ms: now.elapsed().as_millis() as u64,
        })
    }

    fn reranker_info(&self) -> &RerankerInfo {
        &self.info
    }
}