minni 0.1.1

use crate::search::BM25SearchResult;
use anyhow::{Context, Result};
use ndarray::Array2;
use ort::session::{builder::GraphOptimizationLevel, Session, SessionInputValue};
use ort::value::Value;
use std::borrow::Cow;
use std::path::Path;
use tokenizers::{EncodeInput, InputSequence, Tokenizer};

const MAX_LENGTH: usize = 512; // Max length for cross-encoder

/// MiniLM cross-encoder reranker.
pub struct Reranker {
    session: Session,
    tokenizer: Tokenizer,
}

impl Reranker {
    pub fn new(model_path: &Path) -> Result<Self> {
        let model_file = model_path.join("model.onnx");
        let tokenizer_file = model_path.join("tokenizer.json");

        let session = Session::builder()?
            .with_optimization_level(GraphOptimizationLevel::Level3)
            .map_err(|e| anyhow::anyhow!("ORT session builder error: {e}"))?
            .commit_from_file(&model_file)
            .context("Failed to load ONNX model")?;

        let tokenizer = Tokenizer::from_file(tokenizer_file)
            .map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e))?;

        Ok(Self { session, tokenizer })
    }

    /// Re-rank results by score.
    pub fn rerank(
        &mut self,
        query: &str,
        mut results: Vec<BM25SearchResult>,
        limit: usize,
    ) -> Result<Vec<crate::search::SearchResult>> {
        if results.is_empty() {
            return Ok(vec![]);
        }

        // Create query-doc pairs
        let pairs: Vec<(String, String)> = results
            .iter()
            .map(|r| {
                // Build structured doc text: file_path → signature → doc_comment → body remainder
                let mut doc_parts = vec![r.file_path.clone()];
                if let Some(sig) = &r.signature {
                    doc_parts.push(sig.clone());
                }
                if let Some(doc) = &r.doc_comment {
                    doc_parts.push(doc.clone());
                }
                // Fill the remaining ~500-char budget with raw content body.
                // Use chars().count() throughout so the budget is char-based, not byte-based.
                let prefix_len: usize =
                    doc_parts.iter().map(|s| s.chars().count()).sum::<usize>() + doc_parts.len(); // +1 per newline separator
                let remaining = 500usize.saturating_sub(prefix_len);
                if remaining > 50 {
                    let body: String = r.content.chars().take(remaining).collect();
                    doc_parts.push(body);
                }
                // Hard cap: ensure the final string never exceeds 500 chars regardless of
                // multi-byte content or separator arithmetic.
                let doc_text: String = doc_parts.join("\n").chars().take(500).collect();
                (query.to_string(), doc_text)
            })
            .collect();

        // Score all pairs in batches
        let scores = self.score_pairs(&pairs)?;

        // Update scores and sort
        for (result, score) in results.iter_mut().zip(scores.iter()) {
            result.score = *score;
        }

        results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());

        Ok(results.into_iter().take(limit).map(|r| r.into()).collect())
    }

    fn score_pairs(&mut self, pairs: &[(String, String)]) -> Result<Vec<f32>> {
        let batch_size = 8;
        let mut all_scores = Vec::new();

        for batch in pairs.chunks(batch_size) {
            // Use true pair encoding instead of injecting [SEP] manually.
            let inputs: Vec<EncodeInput> = batch
                .iter()
                .map(|(q, d)| {
                    EncodeInput::Dual(
                        InputSequence::Raw(q.clone().into()),
                        InputSequence::Raw(d.clone().into()),
                    )
                })
                .collect();

            let encodings = self
                .tokenizer
                .encode_batch(inputs, true)
                .map_err(|e| anyhow::anyhow!("Tokenization failed: {}", e))?;

            let batch_len = encodings.len();
            let max_len = encodings
                .iter()
                .map(|e| e.len())
                .max()
                .unwrap_or(0)
                .min(MAX_LENGTH);

            // Build input tensors (including token_type_ids for Xenova models)
            let mut input_ids_vec = vec![0i64; batch_len * max_len];
            let mut attention_mask_vec = vec![0i64; batch_len * max_len];
            let mut token_type_ids_vec = vec![0i64; batch_len * max_len];

            for (i, encoding) in encodings.iter().enumerate() {
                let ids = encoding.get_ids();
                let mask = encoding.get_attention_mask();
                let type_ids = encoding.get_type_ids();
                let len = ids.len().min(max_len);

                for j in 0..len {
                    input_ids_vec[i * max_len + j] = ids[j] as i64;
                    attention_mask_vec[i * max_len + j] = mask[j] as i64;
                    token_type_ids_vec[i * max_len + j] = type_ids[j] as i64;
                }
            }

            let input_ids: Array2<i64> =
                Array2::from_shape_vec((batch_len, max_len), input_ids_vec)?;
            let attention_mask: Array2<i64> =
                Array2::from_shape_vec((batch_len, max_len), attention_mask_vec)?;
            let token_type_ids: Array2<i64> =
                Array2::from_shape_vec((batch_len, max_len), token_type_ids_vec)?;

            let input_ids_value = Value::from_array(input_ids)?;
            let attention_mask_value = Value::from_array(attention_mask)?;
            let token_type_ids_value = Value::from_array(token_type_ids)?;

            // Build inputs vector manually with explicit types
            let inputs: Vec<(Cow<'_, str>, SessionInputValue<'_>)> = vec![
                (
                    Cow::Borrowed("input_ids"),
                    SessionInputValue::from(&input_ids_value),
                ),
                (
                    Cow::Borrowed("attention_mask"),
                    SessionInputValue::from(&attention_mask_value),
                ),
                (
                    Cow::Borrowed("token_type_ids"),
                    SessionInputValue::from(&token_type_ids_value),
                ),
            ];

            // Run inference
            let outputs = self.session.run(inputs)?;

            // Extract logits (cross-encoder outputs relevance scores)
            // Cross-encoder outputs [batch, 1] with relevance scores
            let output = outputs
                .get("logits")
                .or_else(|| outputs.get("output"))
                .context("Model output not found")?;

            let output_array = output.try_extract_array::<f32>()?;
            let shape: Vec<usize> = output_array.shape().to_vec();

            // Cross-encoder outputs [batch, 1] - single relevance score per pair
            if shape.len() == 2 && shape[1] == 1 {
                // Perfect - direct relevance scores
                for i in 0..batch_len {
                    all_scores.push(output_array[[i, 0]]);
                }
            } else if shape.len() == 1 {
                // 1D output [batch]
                for i in 0..batch_len {
                    all_scores.push(output_array[[i]]);
                }
            } else {
                // Fallback: use first column or mean
                for i in 0..batch_len {
                    let score = if shape.len() == 2 {
                        output_array[[i, 0]]
                    } else {
                        // Take mean across dimensions
                        let mut sum = 0.0f32;
                        let dim = shape.get(1).copied().unwrap_or(1);
                        for j in 0..dim {
                            sum += output_array[[i, j]];
                        }
                        sum / dim as f32
                    };
                    all_scores.push(score);
                }
            }
        }

        Ok(all_scores)
    }
}