//! Reranker framework — multi-strategy result re-ranking
//!
//! Provides a `Reranker` trait and multiple implementations for re-ranking
//! search results from hybrid queries:
//!
//! - **LinearReranker**: Weighted linear combination of scores
//! - **RRFReranker**: Reciprocal Rank Fusion (robust multi-signal fusion)
//! - **CrossEncoderReranker**: Neural cross-encoder scoring (placeholder for ONNX/API)
//!
//! # Architecture
//!
//! ```text
//!   HybridQueryEngine  ──▶  candidates + ScoreBreakdown
//!         │
//!         ▼
//!     Reranker::rerank(query, candidates) ──▶ RankedResult[]
//! ```

use std::collections::HashMap;
use serde::{Deserialize, Serialize};
#[cfg(feature = "candle")]
use std::sync::Arc;

// ═══════════════════════════════════════════════════════════════════
// Core types
// ═══════════════════════════════════════════════════════════════════

/// A candidate for re-ranking, carrying multi-channel scores.
#[derive(Debug, Clone)]
pub struct RerankCandidate {
    /// Unique node ID
    pub id: String,
    /// Original content text (for cross-encoder)
    pub content: String,
    /// Vector similarity score (0..1, higher = more similar)
    pub vector_score: Option<f64>,
    /// BM25 full-text search score
    pub bm25_score: Option<f64>,
    /// Legacy text match score
    pub text_score: Option<f64>,
    /// Importance score
    pub importance_score: f64,
    /// Recency score
    pub recency_score: f64,
    /// SiFu signal propagation energy score
    pub sifu_score: Option<f64>,
    /// Original rank in each channel (for RRF)
    pub channel_ranks: HashMap<String, usize>,
}

impl RerankCandidate {
    /// Create a new candidate with minimal info.
    pub fn new(id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            content: content.into(),
            vector_score: None,
            bm25_score: None,
            text_score: None,
            importance_score: 0.0,
            recency_score: 0.0,
            sifu_score: None,
            channel_ranks: HashMap::new(),
        }
    }

    /// Set vector score.
    pub fn with_vector_score(mut self, s: f64) -> Self { self.vector_score = Some(s); self }
    /// Set BM25 score.
    pub fn with_bm25_score(mut self, s: f64) -> Self { self.bm25_score = Some(s); self }
    /// Set text score.
    pub fn with_text_score(mut self, s: f64) -> Self { self.text_score = Some(s); self }
    /// Set importance score.
    pub fn with_importance(mut self, s: f64) -> Self { self.importance_score = s; self }
    /// Set recency score.
    pub fn with_recency(mut self, s: f64) -> Self { self.recency_score = s; self }
    /// Set SiFu score.
    pub fn with_sifu_score(mut self, s: f64) -> Self { self.sifu_score = Some(s); self }
    /// Set a channel rank (for RRF).
    pub fn with_rank(mut self, channel: impl Into<String>, rank: usize) -> Self {
        self.channel_ranks.insert(channel.into(), rank);
        self
    }
}

/// A re-ranked result with final score and provenance.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RankedResult {
    /// Node ID
    pub id: String,
    /// Final combined score (higher = better)
    pub score: f64,
    /// Original rank before reranking
    pub original_rank: usize,
    /// New rank after reranking
    pub new_rank: usize,
    /// Score breakdown by channel
    pub score_details: HashMap<String, f64>,
}

/// Error type for reranking operations.
#[derive(Debug, thiserror::Error)]
pub enum RerankError {
    #[error("Reranker configuration error: {0}")]
    Config(String),
    #[error("Reranker runtime error: {0}")]
    Runtime(String),
}

pub type RerankResult<T> = Result<T, RerankError>;

// ═══════════════════════════════════════════════════════════════════
// Reranker trait
// ═══════════════════════════════════════════════════════════════════

/// Trait for re-ranking search results.
///
/// Implementations receive the original query and a list of candidates
/// with multi-channel scores, and return results sorted by final score.
pub trait Reranker: Send + Sync {
    /// Re-rank candidates and return sorted results.
    fn rerank(
        &self,
        query: &str,
        candidates: Vec<RerankCandidate>,
    ) -> RerankResult<Vec<RankedResult>>;

    /// Return the reranker's name.
    fn name(&self) -> &str;
}

// ═══════════════════════════════════════════════════════════════════
// LinearReranker — weighted linear combination
// ═══════════════════════════════════════════════════════════════════

/// Reranker using weighted linear combination of scores.
///
/// final_score = w_vec * vector_score + w_bm25 * bm25_score
///             + w_text * text_score + w_imp * importance_score
///             + w_rec * recency_score + w_sifu * sifu_score
///
/// Weights are normalized automatically.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LinearReranker {
    /// Weight for vector similarity score
    pub w_vector: f64,
    /// Weight for BM25 score
    pub w_bm25: f64,
    /// Weight for text match score
    pub w_text: f64,
    /// Weight for importance score
    pub w_importance: f64,
    /// Weight for recency score
    pub w_recency: f64,
    /// Weight for SiFu score
    pub w_sifu: f64,
}

impl Default for LinearReranker {
    fn default() -> Self {
        Self {
            w_vector: 0.4,
            w_bm25: 0.25,
            w_text: 0.0,
            w_importance: 0.1,
            w_recency: 0.1,
            w_sifu: 0.15,
        }
    }
}

impl LinearReranker {
    /// Create with custom weights.
    pub fn new(w_vector: f64, w_bm25: f64, w_text: f64, w_importance: f64, w_recency: f64, w_sifu: f64) -> Self {
        Self { w_vector, w_bm25, w_text, w_importance, w_recency, w_sifu }
    }

    /// Compute the weighted score for a candidate.
    fn compute_score(&self, c: &RerankCandidate) -> (f64, HashMap<String, f64>) {
        let total_weight = self.w_vector + self.w_bm25 + self.w_text
            + self.w_importance + self.w_recency + self.w_sifu;
        let norm = if total_weight > 0.0 { total_weight } else { 1.0 };

        let mut details = HashMap::new();
        let mut score = 0.0;

        if let Some(vs) = c.vector_score {
            let s = (self.w_vector / norm) * vs;
            score += s;
            details.insert("vector".to_string(), vs);
        }
        if let Some(bs) = c.bm25_score {
            let s = (self.w_bm25 / norm) * bs;
            score += s;
            details.insert("bm25".to_string(), bs);
        }
        if let Some(ts) = c.text_score {
            let s = (self.w_text / norm) * ts;
            score += s;
            details.insert("text".to_string(), ts);
        }
        {
            let s = (self.w_importance / norm) * c.importance_score;
            score += s;
            details.insert("importance".to_string(), c.importance_score);
        }
        {
            let s = (self.w_recency / norm) * c.recency_score;
            score += s;
            details.insert("recency".to_string(), c.recency_score);
        }
        if let Some(ss) = c.sifu_score {
            let s = (self.w_sifu / norm) * ss;
            score += s;
            details.insert("sifu".to_string(), ss);
        }

        (score, details)
    }
}

impl Reranker for LinearReranker {
    fn rerank(
        &self,
        _query: &str,
        candidates: Vec<RerankCandidate>,
    ) -> RerankResult<Vec<RankedResult>> {
        let mut results: Vec<RankedResult> = candidates
            .iter()
            .enumerate()
            .map(|(i, c)| {
                let (score, details) = self.compute_score(c);
                RankedResult {
                    id: c.id.clone(),
                    score,
                    original_rank: i,
                    new_rank: 0, // will be set after sorting
                    score_details: details,
                }
            })
            .collect();

        results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));

        for (i, r) in results.iter_mut().enumerate() {
            r.new_rank = i;
        }

        Ok(results)
    }

    fn name(&self) -> &str {
        "linear"
    }
}

// ═══════════════════════════════════════════════════════════════════
// RRFReranker — Reciprocal Rank Fusion
// ═══════════════════════════════════════════════════════════════════

/// Reciprocal Rank Fusion reranker.
///
/// RRF Score = Σ 1 / (k + rank_i) for each ranking channel.
///
/// This is a robust, parameter-free fusion method that combines multiple
/// ranked lists without requiring score normalization.
///
/// Reference: Cormack et al., "Reciprocal Rank Fusion outperforms
/// Condorcet and individual Rank Learning Methods" (2009).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RRFReranker {
    /// Fusion constant (default: 60, standard value from the paper)
    pub k: f64,
    /// Channels to consider (e.g., ["vector", "bm25", "sifu"])
    /// If empty, all available channels are used.
    pub channels: Vec<String>,
}

impl Default for RRFReranker {
    fn default() -> Self {
        Self {
            k: 60.0,
            channels: Vec::new(), // use all available
        }
    }
}

impl RRFReranker {
    /// Create with custom k constant.
    pub fn with_k(k: f64) -> Self {
        Self { k, channels: Vec::new() }
    }

    /// Specify which channels to fuse.
    pub fn with_channels(mut self, channels: Vec<String>) -> Self {
        self.channels = channels;
        self
    }

    /// Build per-channel rankings from candidates.
    ///
    /// Two modes:
    /// 1. Pre-computed ranks: use `channel_ranks` directly (rank 1 = best).
    /// 2. Score-derived ranks: sort by score descending, assign rank 1..N.
    fn build_channel_rankings(
        &self,
        candidates: &[RerankCandidate],
    ) -> HashMap<String, Vec<(String, usize)>> {
        // Collect pre-computed ranks directly
        let mut precomputed: HashMap<String, Vec<(String, usize)>> = HashMap::new();
        // Collect scores for rank derivation
        let mut score_based: HashMap<String, Vec<(String, f64)>> = HashMap::new();

        for c in candidates {
            if !c.channel_ranks.is_empty() {
                // Use pre-computed ranks directly
                for (ch, rank) in &c.channel_ranks {
                    precomputed
                        .entry(ch.clone())
                        .or_default()
                        .push((c.id.clone(), *rank));
                }
            } else {
                // Derive ranks from scores
                if let Some(vs) = c.vector_score {
                    score_based
                        .entry("vector".to_string())
                        .or_default()
                        .push((c.id.clone(), vs));
                }
                if let Some(bs) = c.bm25_score {
                    score_based
                        .entry("bm25".to_string())
                        .or_default()
                        .push((c.id.clone(), bs));
                }
                if let Some(ss) = c.sifu_score {
                    score_based
                        .entry("sifu".to_string())
                        .or_default()
                        .push((c.id.clone(), ss));
                }
            }
        }

        let mut channel_rankings = HashMap::new();

        // Add pre-computed ranks (already in rank format)
        for (ch, ranks) in precomputed {
            if !self.channels.is_empty() && !self.channels.contains(&ch) {
                continue;
            }
            channel_rankings.insert(ch, ranks);
        }

        // Convert scores to ranks (sort by score descending, assign rank 1..N)
        for (ch, mut scores) in score_based {
            if !self.channels.is_empty() && !self.channels.contains(&ch) {
                continue;
            }
            scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
            let ranks: Vec<(String, usize)> = scores
                .iter()
                .enumerate()
                .map(|(rank, (id, _))| (id.clone(), rank + 1))
                .collect();
            channel_rankings.insert(ch, ranks);
        }

        channel_rankings
    }
}

impl Reranker for RRFReranker {
    fn rerank(
        &self,
        _query: &str,
        candidates: Vec<RerankCandidate>,
    ) -> RerankResult<Vec<RankedResult>> {
        if candidates.is_empty() {
            return Ok(Vec::new());
        }

        let channel_rankings = self.build_channel_rankings(&candidates);

        // Build RRF scores per candidate
        let mut rrf_scores: HashMap<String, f64> = HashMap::new();
        let mut rrf_details: HashMap<String, HashMap<String, f64>> = HashMap::new();

        for (channel, rankings) in &channel_rankings {
            for (id, rank) in rankings {
                let rrf_contribution = 1.0 / (self.k + *rank as f64);
                *rrf_scores.entry(id.clone()).or_insert(0.0) += rrf_contribution;
                rrf_details
                    .entry(id.clone())
                    .or_default()
                    .insert(format!("{}_rank", channel), *rank as f64);
            }
        }

        // Build original rank mapping
        let orig_rank_map: HashMap<String, usize> = candidates
            .iter()
            .enumerate()
            .map(|(i, c)| (c.id.clone(), i))
            .collect();

        let mut results: Vec<RankedResult> = rrf_scores
            .into_iter()
            .map(|(id, score)| {
                RankedResult {
                    original_rank: orig_rank_map.get(&id).copied().unwrap_or(0),
                    new_rank: 0,
                    score,
                    score_details: rrf_details.remove(&id).unwrap_or_default(),
                    id,
                }
            })
            .collect();

        results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));

        for (i, r) in results.iter_mut().enumerate() {
            r.new_rank = i;
        }

        Ok(results)
    }

    fn name(&self) -> &str {
        "rrf"
    }
}

// ═══════════════════════════════════════════════════════════════════
// CrossEncoderModel — local BERT inference via Candle (Phase-65)
// ═══════════════════════════════════════════════════════════════════

/// Local Cross-Encoder model using a BERT backbone + linear classifier.
///
/// Architecture: BERT → [CLS] hidden → Linear(hidden→1) → sigmoid → relevance.
///
/// Supported models (HuggingFace):
/// - `cross-encoder/ms-marco-MiniLM-L-6-v2` (default, 384-dim, 6 layers)
/// - `cross-encoder/ms-marco-MiniLM-L-12-v2` (384-dim, 12 layers)
/// - `BAAI/bge-reranker-base` (768-dim, 12 layers)
///
/// Falls back to heuristic scoring if model loading fails.
#[cfg(feature = "candle")]
pub struct CrossEncoderModel {
    model: candle_transformers::models::bert::BertModel,
    classifier: candle_nn::Linear,
    tokenizer: tokenizers::Tokenizer,
    device: candle_core::Device,
    max_length: usize,
}

#[cfg(feature = "candle")]
impl std::fmt::Debug for CrossEncoderModel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CrossEncoderModel")
            .field("device", &format!("{:?}", self.device))
            .field("max_length", &self.max_length)
            .finish()
    }
}

#[cfg(feature = "candle")]
impl CrossEncoderModel {
    /// Load a cross-encoder model from HuggingFace Hub.
    ///
    /// Uses `DeviceSelector` for automatic CPU/GPU device selection.
    pub fn from_pretrained(model_id: &str) -> Result<Self, RerankError> {
        use candle_core::DType;
        use candle_nn::VarBuilder;
        use hf_hub::{api::sync::ApiBuilder, Repo, RepoType};

        let selector = crate::device::global_selector();
        let (device, is_gpu) = selector.candle_device_with_fallback("cross-encoder");

        tracing::info!(
            model = model_id,
            is_gpu = is_gpu,
            "Loading Cross-Encoder model"
        );

        let api = ApiBuilder::from_env()
            .with_progress(true)
            .build()
            .map_err(|e| RerankError::Config(format!("HF Hub API init failed: {}", e)))?;
        let repo = api.repo(Repo::new(model_id.to_string(), RepoType::Model));

        // Download model weights
        let weights_path = repo
            .get("model.safetensors")
            .or_else(|_| repo.get("pytorch_model.bin"))
            .map_err(|e| RerankError::Config(format!("Failed to download weights: {}", e)))?;

        // Download config
        let config_path = repo
            .get("config.json")
            .map_err(|e| RerankError::Config(format!("Failed to download config: {}", e)))?;

        // Download tokenizer
        let tokenizer_path = repo
            .get("tokenizer.json")
            .map_err(|e| RerankError::Config(format!("Failed to download tokenizer: {}", e)))?;

        // Parse BERT config
        let config_json = std::fs::read_to_string(&config_path)
            .map_err(|e| RerankError::Config(format!("Read config failed: {}", e)))?;
        let bert_config: candle_transformers::models::bert::Config =
            serde_json::from_str(&config_json)
                .map_err(|e| RerankError::Config(format!("Parse config failed: {}", e)))?;

        let max_length = bert_config.max_position_embeddings.min(512);
        let hidden_size = bert_config.hidden_size;

        // Load weights
        let vb = if weights_path.extension().map_or(false, |e| e == "safetensors") {
            unsafe {
                VarBuilder::from_mmaped_safetensors(&[&weights_path], DType::F32, &device)
                    .map_err(|e| RerankError::Runtime(format!("Load safetensors failed: {}", e)))?
            }
        } else {
            VarBuilder::from_pth(&weights_path, DType::F32, &device)
                .map_err(|e| RerankError::Runtime(format!("Load weights failed: {}", e)))?
        };

        // Load BERT model
        let model =
            candle_transformers::models::bert::BertModel::load(vb.pp("bert"), &bert_config)
                .or_else(|_| {
                    candle_transformers::models::bert::BertModel::load(vb.clone(), &bert_config)
                })
                .map_err(|e| RerankError::Runtime(format!("Load BERT model failed: {}", e)))?;

        // Load classifier head: Linear(hidden_size → 1)
        // Cross-encoder models typically store this as "classifier" or "cls"
        let classifier = candle_nn::linear(hidden_size, 1, vb.pp("classifier"))
            .or_else(|_| candle_nn::linear(hidden_size, 1, vb.pp("cls")))
            .map_err(|e| {
                RerankError::Runtime(format!(
                    "Load classifier head failed (hidden_size={}): {}",
                    hidden_size, e
                ))
            })?;

        // Load tokenizer
        let tokenizer = tokenizers::Tokenizer::from_file(&tokenizer_path)
            .map_err(|e| RerankError::Config(format!("Load tokenizer failed: {}", e)))?;

        tracing::info!(
            hidden_size = hidden_size,
            max_length = max_length,
            device = ?device,
            "Cross-Encoder model loaded successfully"
        );

        Ok(Self {
            model,
            classifier,
            tokenizer,
            device,
            max_length,
        })
    }

    /// Score a (query, document) pair. Returns logit score (higher = more relevant).
    ///
    /// Pipeline: tokenize → BERT → [CLS] → classifier → score
    #[allow(dead_code)]
    fn score_pair(&self, query: &str, document: &str) -> Result<f64, RerankError> {
        use candle_core::{DType, Tensor};
        use candle_nn::Module;

        // Tokenize as pair: [CLS] query [SEP] document [SEP]
        let encoding = self
            .tokenizer
            .encode((query, document), true)
            .map_err(|e| RerankError::Runtime(format!("Tokenization failed: {}", e)))?;

        let mut ids: Vec<u32> = encoding.get_ids().to_vec();
        let mut type_ids: Vec<u32> = encoding.get_type_ids().to_vec();
        let mut attention: Vec<u32> = encoding.get_attention_mask().to_vec();

        // Truncate to max_length
        ids.truncate(self.max_length);
        type_ids.truncate(self.max_length);
        attention.truncate(self.max_length);

        let input_ids = Tensor::new(&ids[..], &self.device)
            .map_err(|e| RerankError::Runtime(e.to_string()))?
            .unsqueeze(0)
            .map_err(|e| RerankError::Runtime(e.to_string()))?;

        let token_type_ids = Tensor::new(&type_ids[..], &self.device)
            .map_err(|e| RerankError::Runtime(e.to_string()))?
            .unsqueeze(0)
            .map_err(|e| RerankError::Runtime(e.to_string()))?;

        let attention_mask = Tensor::new(&attention[..], &self.device)
            .map_err(|e| RerankError::Runtime(e.to_string()))?
            .unsqueeze(0)
            .map_err(|e| RerankError::Runtime(e.to_string()))?;

        // Attention mask needs to be f32 for the BERT model
        let attention_mask_f32 = attention_mask
            .to_dtype(DType::F32)
            .map_err(|e| RerankError::Runtime(e.to_string()))?;

        // Forward pass
        let output = self
            .model
            .forward(&input_ids, &token_type_ids, Some(&attention_mask_f32))
            .map_err(|e| RerankError::Runtime(format!("BERT forward failed: {}", e)))?;

        // Extract [CLS] token (first token)
        let cls_output = output
            .narrow(1, 0, 1)
            .map_err(|e| RerankError::Runtime(e.to_string()))?
            .squeeze(1)
            .map_err(|e| RerankError::Runtime(e.to_string()))?;

        // Classifier: [batch, hidden] → [batch, 1]
        let logits = self
            .classifier
            .forward(&cls_output)
            .map_err(|e| RerankError::Runtime(format!("Classifier forward failed: {}", e)))?;

        // Get scalar logit
        let logit: f32 = logits
            .to_device(&candle_core::Device::Cpu)
            .map_err(|e| RerankError::Runtime(e.to_string()))?
            .squeeze(0)
            .map_err(|e| RerankError::Runtime(e.to_string()))?
            .squeeze(0)
            .map_err(|e| RerankError::Runtime(e.to_string()))?
            .to_scalar()
            .map_err(|e| RerankError::Runtime(e.to_string()))?;

        Ok(logit as f64)
    }

    /// Score a batch of (query, document) pairs.
    ///
    /// Pads to the longest sequence in the batch for efficient GPU utilization.
    fn score_batch(
        &self,
        query: &str,
        documents: &[&str],
        batch_size: usize,
    ) -> Result<Vec<f64>, RerankError> {
        use candle_core::{DType, Tensor};
        use candle_nn::Module;

        let mut all_scores = Vec::with_capacity(documents.len());

        for chunk in documents.chunks(batch_size) {
            // Tokenize all pairs in this batch
            let mut batch_ids: Vec<Vec<u32>> = Vec::with_capacity(chunk.len());
            let mut batch_type_ids: Vec<Vec<u32>> = Vec::with_capacity(chunk.len());
            let mut batch_attention: Vec<Vec<u32>> = Vec::with_capacity(chunk.len());
            let mut max_len = 0usize;

            for doc in chunk {
                let encoding = self
                    .tokenizer
                    .encode((query, *doc), true)
                    .map_err(|e| RerankError::Runtime(format!("Tokenization failed: {}", e)))?;

                let mut ids: Vec<u32> = encoding.get_ids().to_vec();
                let mut type_ids: Vec<u32> = encoding.get_type_ids().to_vec();
                let mut attention: Vec<u32> = encoding.get_attention_mask().to_vec();

                ids.truncate(self.max_length);
                type_ids.truncate(self.max_length);
                attention.truncate(self.max_length);

                max_len = max_len.max(ids.len());

                batch_ids.push(ids);
                batch_type_ids.push(type_ids);
                batch_attention.push(attention);
            }

            // Pad to max_len
            for i in 0..chunk.len() {
                while batch_ids[i].len() < max_len {
                    batch_ids[i].push(0);
                    batch_type_ids[i].push(0);
                    batch_attention[i].push(0);
                }
            }

            let batch_n = chunk.len();

            // Flatten and create tensors
            let flat_ids: Vec<u32> = batch_ids.into_iter().flatten().collect();
            let flat_type_ids: Vec<u32> = batch_type_ids.into_iter().flatten().collect();
            let flat_attention: Vec<u32> = batch_attention.into_iter().flatten().collect();

            let input_ids =
                Tensor::from_vec(flat_ids, (batch_n, max_len), &self.device)
                    .map_err(|e| RerankError::Runtime(e.to_string()))?;
            let token_type_ids =
                Tensor::from_vec(flat_type_ids, (batch_n, max_len), &self.device)
                    .map_err(|e| RerankError::Runtime(e.to_string()))?;
            let attention_mask =
                Tensor::from_vec(flat_attention, (batch_n, max_len), &self.device)
                    .map_err(|e| RerankError::Runtime(e.to_string()))?;
            let attention_mask_f32 = attention_mask
                .to_dtype(DType::F32)
                .map_err(|e| RerankError::Runtime(e.to_string()))?;

            // Forward pass (batched)
            let output = self
                .model
                .forward(&input_ids, &token_type_ids, Some(&attention_mask_f32))
                .map_err(|e| RerankError::Runtime(format!("BERT batch forward failed: {}", e)))?;

            // [CLS] token → classifier
            let cls_output = output
                .narrow(1, 0, 1)
                .map_err(|e| RerankError::Runtime(e.to_string()))?
                .squeeze(1)
                .map_err(|e| RerankError::Runtime(e.to_string()))?;

            let logits = self
                .classifier
                .forward(&cls_output)
                .map_err(|e| {
                    RerankError::Runtime(format!("Classifier batch forward failed: {}", e))
                })?
                .squeeze(1)
                .map_err(|e| RerankError::Runtime(e.to_string()))?;

            // Convert to CPU and collect scores
            let scores: Vec<f32> = logits
                .to_device(&candle_core::Device::Cpu)
                .map_err(|e| RerankError::Runtime(e.to_string()))?
                .to_vec1()
                .map_err(|e| RerankError::Runtime(e.to_string()))?;

            all_scores.extend(scores.into_iter().map(|s| s as f64));
        }

        Ok(all_scores)
    }
}

// ═══════════════════════════════════════════════════════════════════
// CrossEncoderReranker — neural cross-encoder with heuristic fallback
// ═══════════════════════════════════════════════════════════════════

/// Scoring mode for the cross-encoder reranker.
enum CrossEncoderMode {
    /// Heuristic word-overlap scoring (no model required)
    Heuristic,
    /// Local BERT model inference via Candle (GPU/CPU)
    #[cfg(feature = "candle")]
    Model(Arc<CrossEncoderModel>),
}

impl std::fmt::Debug for CrossEncoderMode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CrossEncoderMode::Heuristic => write!(f, "Heuristic"),
            #[cfg(feature = "candle")]
            CrossEncoderMode::Model(_) => write!(f, "Model(BERT)"),
        }
    }
}

/// Cross-encoder reranker with BERT model inference and heuristic fallback.
///
/// Scores (query, document) pairs using a pre-trained cross-encoder BERT model.
/// If the model fails to load (e.g., no internet, no GPU), falls back to
/// word-overlap heuristic scoring.
///
/// # Supported Models
///
/// - `cross-encoder/ms-marco-MiniLM-L-6-v2` — fast, 6 layers, good quality
/// - `cross-encoder/ms-marco-MiniLM-L-12-v2` — balanced, 12 layers
/// - `BAAI/bge-reranker-base` — 12 layers, strong multilingual
///
/// # Example
///
/// ```no_run
/// use aifastdb_core::query::reranker::{CrossEncoderReranker, Reranker, RerankCandidate};
///
/// // Default: heuristic mode (no model download)
/// let reranker = CrossEncoderReranker::default();
///
/// // With model: downloads from HuggingFace Hub
/// let reranker = CrossEncoderReranker::with_model("cross-encoder/ms-marco-MiniLM-L-6-v2");
/// ```
pub struct CrossEncoderReranker {
    /// Model name (for logging/identification)
    pub model_name: String,
    /// Batch size for model inference
    pub batch_size: usize,
    /// Blend weight for cross-encoder score vs vector score
    /// final = ce_weight * ce_score + (1 - ce_weight) * vector_score
    pub ce_weight: f64,
    /// Scoring mode
    mode: CrossEncoderMode,
}

impl std::fmt::Debug for CrossEncoderReranker {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CrossEncoderReranker")
            .field("model_name", &self.model_name)
            .field("batch_size", &self.batch_size)
            .field("ce_weight", &self.ce_weight)
            .field("mode", &self.mode)
            .finish()
    }
}

impl Default for CrossEncoderReranker {
    fn default() -> Self {
        Self {
            model_name: "heuristic-overlap".to_string(),
            batch_size: 32,
            ce_weight: 0.4,
            mode: CrossEncoderMode::Heuristic,
        }
    }
}

impl CrossEncoderReranker {
    /// Create with a specific model name (heuristic mode).
    pub fn new(model_name: impl Into<String>) -> Self {
        Self {
            model_name: model_name.into(),
            batch_size: 32,
            ce_weight: 0.4,
            mode: CrossEncoderMode::Heuristic,
        }
    }

    /// Create with a BERT cross-encoder model from HuggingFace Hub.
    ///
    /// Falls back to heuristic mode if model loading fails.
    /// Uses `DeviceSelector` for automatic CPU/GPU selection.
    #[cfg(feature = "candle")]
    pub fn with_model(model_id: &str) -> Self {
        match CrossEncoderModel::from_pretrained(model_id) {
            Ok(model) => {
                tracing::info!(
                    model = model_id,
                    "Cross-encoder model loaded, using neural scoring"
                );
                Self {
                    model_name: model_id.to_string(),
                    batch_size: 32,
                    ce_weight: 0.7, // Higher weight when using real model
                    mode: CrossEncoderMode::Model(Arc::new(model)),
                }
            }
            Err(e) => {
                tracing::warn!(
                    model = model_id,
                    error = %e,
                    "Failed to load cross-encoder model, falling back to heuristic"
                );
                Self {
                    model_name: format!("heuristic({})", model_id),
                    batch_size: 32,
                    ce_weight: 0.4,
                    mode: CrossEncoderMode::Heuristic,
                }
            }
        }
    }

    /// Set the batch size for model inference.
    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size.max(1);
        self
    }

    /// Set the cross-encoder weight for blending with vector score.
    pub fn with_ce_weight(mut self, weight: f64) -> Self {
        self.ce_weight = weight.clamp(0.0, 1.0);
        self
    }

    /// Check if using neural model inference.
    pub fn is_model_mode(&self) -> bool {
        #[cfg(feature = "candle")]
        {
            matches!(self.mode, CrossEncoderMode::Model(_))
        }
        #[cfg(not(feature = "candle"))]
        {
            false
        }
    }

    /// Check if using heuristic fallback.
    pub fn is_heuristic_mode(&self) -> bool {
        matches!(self.mode, CrossEncoderMode::Heuristic)
    }

    /// Score a (query, content) pair using simple overlap heuristic.
    /// Returns a score in [0, 1].
    fn heuristic_score(query: &str, content: &str) -> f64 {
        let query_lower = query.to_lowercase();
        let query_terms: Vec<&str> = query_lower.split_whitespace().collect();
        let content_lower = content.to_lowercase();

        if query_terms.is_empty() {
            return 0.0;
        }

        let matches = query_terms
            .iter()
            .filter(|t| content_lower.contains(**t))
            .count();

        matches as f64 / query_terms.len() as f64
    }

    /// Apply sigmoid to a logit value.
    fn sigmoid(x: f64) -> f64 {
        1.0 / (1.0 + (-x).exp())
    }

    /// Re-rank using heuristic scoring.
    fn rerank_heuristic(
        &self,
        query: &str,
        candidates: &[RerankCandidate],
    ) -> Vec<(f64, HashMap<String, f64>)> {
        candidates
            .iter()
            .map(|c| {
                let ce_score = Self::heuristic_score(query, &c.content);
                let mut details = HashMap::new();
                details.insert("cross_encoder".to_string(), ce_score);
                details.insert("mode".to_string(), 0.0); // 0 = heuristic

                let final_score = if let Some(vs) = c.vector_score {
                    self.ce_weight * ce_score + (1.0 - self.ce_weight) * vs
                } else {
                    ce_score
                };
                details.insert("final".to_string(), final_score);

                (final_score, details)
            })
            .collect()
    }

    /// Re-rank using model inference (batched).
    #[cfg(feature = "candle")]
    fn rerank_model(
        &self,
        model: &CrossEncoderModel,
        query: &str,
        candidates: &[RerankCandidate],
    ) -> Result<Vec<(f64, HashMap<String, f64>)>, RerankError> {
        let documents: Vec<&str> = candidates.iter().map(|c| c.content.as_str()).collect();

        let logits = model.score_batch(query, &documents, self.batch_size)?;

        Ok(logits
            .into_iter()
            .zip(candidates.iter())
            .map(|(logit, c)| {
                let ce_score = Self::sigmoid(logit);
                let mut details = HashMap::new();
                details.insert("cross_encoder".to_string(), ce_score);
                details.insert("logit".to_string(), logit);
                details.insert("mode".to_string(), 1.0); // 1 = model

                let final_score = if let Some(vs) = c.vector_score {
                    self.ce_weight * ce_score + (1.0 - self.ce_weight) * vs
                } else {
                    ce_score
                };
                details.insert("final".to_string(), final_score);

                (final_score, details)
            })
            .collect())
    }
}

impl Reranker for CrossEncoderReranker {
    fn rerank(
        &self,
        query: &str,
        candidates: Vec<RerankCandidate>,
    ) -> RerankResult<Vec<RankedResult>> {
        if candidates.is_empty() {
            return Ok(Vec::new());
        }

        // Score all candidates
        let scored: Vec<(f64, HashMap<String, f64>)> = match &self.mode {
            CrossEncoderMode::Heuristic => self.rerank_heuristic(query, &candidates),
            #[cfg(feature = "candle")]
            CrossEncoderMode::Model(model) => {
                // Try model inference; fall back to heuristic on error
                match self.rerank_model(model, query, &candidates) {
                    Ok(results) => results,
                    Err(e) => {
                        tracing::warn!(
                            error = %e,
                            "Cross-encoder model inference failed, falling back to heuristic"
                        );
                        // Log fallback event
                        {
                            let selector = crate::device::global_selector();
                            selector.emit_fallback_public(crate::device::GpuFallbackEvent {
                                reason: crate::device::FallbackReason::RuntimeError,
                                detail: format!("Model inference failed: {}", e),
                                operation: "cross-encoder-rerank".to_string(),
                                requested_device: crate::device::DevicePreference::Auto,
                            });
                        }
                        self.rerank_heuristic(query, &candidates)
                    }
                }
            }
        };

        // Build sorted results
        let mut results: Vec<RankedResult> = scored
            .into_iter()
            .enumerate()
            .map(|(i, (score, details))| RankedResult {
                id: candidates[i].id.clone(),
                score,
                original_rank: i,
                new_rank: 0,
                score_details: details,
            })
            .collect();

        results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));

        for (i, r) in results.iter_mut().enumerate() {
            r.new_rank = i;
        }

        Ok(results)
    }

    fn name(&self) -> &str {
        "cross-encoder"
    }
}

// Implement Send + Sync (model Arc is Send+Sync safe)
unsafe impl Send for CrossEncoderReranker {}
unsafe impl Sync for CrossEncoderReranker {}

// ═══════════════════════════════════════════════════════════════════
// Tests
// ═══════════════════════════════════════════════════════════════════

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

    fn make_candidates() -> Vec<RerankCandidate> {
        vec![
            RerankCandidate::new("a", "Rust programming language guide")
                .with_vector_score(0.9)
                .with_bm25_score(0.3)
                .with_importance(0.8)
                .with_recency(0.5),
            RerankCandidate::new("b", "Python machine learning tutorial")
                .with_vector_score(0.5)
                .with_bm25_score(0.95)
                .with_importance(0.6)
                .with_recency(0.7),
            RerankCandidate::new("c", "Database design fundamentals")
                .with_vector_score(0.7)
                .with_bm25_score(0.6)
                .with_sifu_score(0.85)
                .with_importance(0.9)
                .with_recency(0.3),
            RerankCandidate::new("d", "Rust async runtime internals")
                .with_vector_score(0.85)
                .with_bm25_score(0.7)
                .with_sifu_score(0.4)
                .with_importance(0.4)
                .with_recency(0.9),
        ]
    }

    #[test]
    fn test_linear_reranker_default() {
        let reranker = LinearReranker::default();
        let results = reranker.rerank("rust", make_candidates()).unwrap();

        assert_eq!(results.len(), 4);
        // All results should have scores > 0
        for r in &results {
            assert!(r.score > 0.0);
            assert!(!r.score_details.is_empty());
        }
        // First result should have highest score
        assert!(results[0].score >= results[1].score);
        assert!(results[1].score >= results[2].score);
    }

    #[test]
    fn test_linear_reranker_custom_weights() {
        // Heavy BM25 weight should favor candidate "b" (highest BM25)
        let reranker = LinearReranker::new(0.0, 1.0, 0.0, 0.0, 0.0, 0.0);
        let results = reranker.rerank("python", make_candidates()).unwrap();

        assert_eq!(results[0].id, "b");
    }

    #[test]
    fn test_linear_reranker_vector_heavy() {
        // Heavy vector weight should favor candidate "a" (highest vector_score)
        let reranker = LinearReranker::new(1.0, 0.0, 0.0, 0.0, 0.0, 0.0);
        let results = reranker.rerank("rust", make_candidates()).unwrap();

        assert_eq!(results[0].id, "a");
    }

    #[test]
    fn test_linear_reranker_empty() {
        let reranker = LinearReranker::default();
        let results = reranker.rerank("test", vec![]).unwrap();
        assert!(results.is_empty());
    }

    #[test]
    fn test_rrf_reranker_default() {
        let reranker = RRFReranker::default();
        let results = reranker.rerank("test", make_candidates()).unwrap();

        assert_eq!(results.len(), 4);
        for r in &results {
            assert!(r.score > 0.0);
        }
        // Results should be sorted by score descending
        for i in 1..results.len() {
            assert!(results[i - 1].score >= results[i].score);
        }
    }

    #[test]
    fn test_rrf_reranker_with_ranks() {
        // Test with pre-computed channel ranks
        // RRF(x) = 1/(60+1) + 1/(60+3) = 0.01639 + 0.01587 = 0.03226
        // RRF(y) = 1/(60+3) + 1/(60+1) = 0.01587 + 0.01639 = 0.03226
        // RRF(z) = 1/(60+2) + 1/(60+2) = 0.01613 + 0.01613 = 0.03226
        // All three have identical RRF scores! Adjust ranks for a clear winner.
        let candidates = vec![
            RerankCandidate::new("x", "doc x")
                .with_rank("vector", 1)
                .with_rank("bm25", 3),
            RerankCandidate::new("y", "doc y")
                .with_rank("vector", 3)
                .with_rank("bm25", 2),
            RerankCandidate::new("z", "doc z")
                .with_rank("vector", 2)
                .with_rank("bm25", 1),
        ];

        let reranker = RRFReranker::default();
        let results = reranker.rerank("test", candidates).unwrap();

        assert_eq!(results.len(), 3);
        // "x": 1/(60+1) + 1/(60+3) = 0.016393 + 0.015873 = 0.032267
        // "z": 1/(60+2) + 1/(60+1) = 0.016129 + 0.016393 = 0.032522  ← highest
        // "y": 1/(60+3) + 1/(60+2) = 0.015873 + 0.016129 = 0.032003
        assert_eq!(results[0].id, "z");
    }

    #[test]
    fn test_rrf_reranker_empty() {
        let reranker = RRFReranker::default();
        let results = reranker.rerank("test", vec![]).unwrap();
        assert!(results.is_empty());
    }

    #[test]
    fn test_rrf_with_specific_channels() {
        let reranker = RRFReranker::default()
            .with_channels(vec!["vector".to_string()]);

        let results = reranker.rerank("test", make_candidates()).unwrap();
        // Only vector channel used, so ranking by vector score
        assert_eq!(results[0].id, "a"); // highest vector score
    }

    #[test]
    fn test_cross_encoder_reranker_heuristic() {
        let reranker = CrossEncoderReranker::default();
        assert!(reranker.is_heuristic_mode());
        assert!(!reranker.is_model_mode());

        let candidates = vec![
            RerankCandidate::new("a", "Rust programming language systems")
                .with_vector_score(0.5),
            RerankCandidate::new("b", "Python data science toolkit")
                .with_vector_score(0.8),
            RerankCandidate::new("c", "Rust async concurrency patterns")
                .with_vector_score(0.6),
        ];

        let results = reranker.rerank("Rust programming", candidates).unwrap();
        assert_eq!(results.len(), 3);

        // "a" has both "Rust" and "programming" matching → highest CE overlap
        // Check that cross_encoder detail exists
        assert!(results[0].score_details.contains_key("cross_encoder"));
        // Verify mode=0 (heuristic)
        assert_eq!(*results[0].score_details.get("mode").unwrap(), 0.0);
    }

    #[test]
    fn test_cross_encoder_empty_query() {
        let reranker = CrossEncoderReranker::default();
        let candidates = vec![
            RerankCandidate::new("a", "some content").with_vector_score(0.9),
        ];
        let results = reranker.rerank("", candidates).unwrap();
        // With empty query, CE score is 0, so vector score dominates
        assert_eq!(results.len(), 1);
    }

    #[test]
    fn test_cross_encoder_empty_candidates() {
        let reranker = CrossEncoderReranker::default();
        let results = reranker.rerank("test query", vec![]).unwrap();
        assert!(results.is_empty());
    }

    #[test]
    fn test_cross_encoder_heuristic_scoring_quality() {
        // Full match: all query terms present
        let score_full = CrossEncoderReranker::heuristic_score(
            "rust async programming",
            "Rust is a systems programming language with async support",
        );
        // Partial match: only some terms
        let score_partial = CrossEncoderReranker::heuristic_score(
            "rust async programming",
            "Python is a great language",
        );
        // No match
        let score_none = CrossEncoderReranker::heuristic_score(
            "rust async programming",
            "The quick brown fox jumps",
        );

        assert!(score_full > score_partial, "Full match > partial: {} > {}", score_full, score_partial);
        assert!(score_partial >= score_none, "Partial >= none: {} >= {}", score_partial, score_none);
        assert!(score_full > 0.0);
    }

    #[test]
    fn test_cross_encoder_ce_weight() {
        // Custom weight: CE-only scoring
        let reranker = CrossEncoderReranker::new("test")
            .with_ce_weight(1.0);

        let candidates = vec![
            RerankCandidate::new("a", "rust programming")
                .with_vector_score(0.1), // Low vector score
            RerankCandidate::new("b", "unrelated topic")
                .with_vector_score(0.99), // High vector score
        ];

        let results = reranker.rerank("rust programming", candidates).unwrap();
        // With ce_weight=1.0, CE score dominates — "a" should win
        assert_eq!(results[0].id, "a");
    }

    #[test]
    fn test_cross_encoder_vector_weight() {
        // Custom weight: vector-only scoring
        let reranker = CrossEncoderReranker::new("test")
            .with_ce_weight(0.0);

        let candidates = vec![
            RerankCandidate::new("a", "rust programming")
                .with_vector_score(0.1),
            RerankCandidate::new("b", "unrelated topic")
                .with_vector_score(0.99),
        ];

        let results = reranker.rerank("rust programming", candidates).unwrap();
        // With ce_weight=0.0, vector score dominates — "b" should win
        assert_eq!(results[0].id, "b");
    }

    #[test]
    fn test_cross_encoder_sigmoid() {
        assert!((CrossEncoderReranker::sigmoid(0.0) - 0.5).abs() < 1e-6);
        assert!(CrossEncoderReranker::sigmoid(10.0) > 0.99);
        assert!(CrossEncoderReranker::sigmoid(-10.0) < 0.01);
    }

    #[test]
    fn test_cross_encoder_batch_size() {
        let reranker = CrossEncoderReranker::new("test")
            .with_batch_size(4);
        assert_eq!(reranker.batch_size, 4);

        // Minimum 1
        let reranker = CrossEncoderReranker::new("test")
            .with_batch_size(0);
        assert_eq!(reranker.batch_size, 1);
    }

    /// Test with real cross-encoder model (requires model download).
    /// Run with: cargo test --features candle -- --ignored test_cross_encoder_model
    #[cfg(feature = "candle")]
    #[test]
    #[ignore]
    fn test_cross_encoder_model_inference() {
        let reranker =
            CrossEncoderReranker::with_model("cross-encoder/ms-marco-MiniLM-L-6-v2");
        assert!(reranker.is_model_mode(), "Model should load successfully");

        let candidates = vec![
            RerankCandidate::new("a", "Rust is a systems programming language")
                .with_vector_score(0.5),
            RerankCandidate::new("b", "The weather is nice today")
                .with_vector_score(0.5),
            RerankCandidate::new("c", "Rust programming with async/await patterns")
                .with_vector_score(0.5),
        ];

        let results = reranker.rerank("What is Rust programming?", candidates).unwrap();
        assert_eq!(results.len(), 3);

        // Model should recognize that "a" and "c" are more relevant
        // Both should score higher than "b" (irrelevant weather topic)
        let b_rank = results.iter().find(|r| r.id == "b").unwrap().new_rank;
        assert_eq!(b_rank, 2, "Irrelevant doc should be ranked last");

        // Verify mode=1 (model)
        assert_eq!(*results[0].score_details.get("mode").unwrap(), 1.0);
        // Verify logit exists
        assert!(results[0].score_details.contains_key("logit"));
    }

    /// Test model fallback to heuristic when model doesn't exist.
    #[cfg(feature = "candle")]
    #[test]
    fn test_cross_encoder_model_fallback() {
        let reranker = CrossEncoderReranker::with_model("nonexistent/model-that-wont-load");
        // Should fall back to heuristic
        assert!(reranker.is_heuristic_mode());
    }

    #[test]
    fn test_reranker_preserves_original_rank() {
        let reranker = LinearReranker::default();
        let results = reranker.rerank("test", make_candidates()).unwrap();

        // Verify original_rank was correctly set
        for r in &results {
            assert!(r.original_rank < 4);
        }
    }

    #[test]
    fn test_reranker_trait_object() {
        // Verify Reranker can be used as trait object
        let rerankers: Vec<Box<dyn Reranker>> = vec![
            Box::new(LinearReranker::default()),
            Box::new(RRFReranker::default()),
            Box::new(CrossEncoderReranker::default()),
        ];

        for reranker in &rerankers {
            let results = reranker.rerank("test", make_candidates()).unwrap();
            assert!(!results.is_empty());
            assert!(!reranker.name().is_empty());
        }
    }
}
