//! Function score query: modify query scores using scoring functions.
//!
//! Wraps an inner query and applies one or more scoring functions
//! (weight, field_value_factor, random_score) to modify the final score.
//!
//! See [[elasticsearch-parity]].

use luci_core::{DocId, Result, ScoreMode, Scorer, TwoPhaseIterator, NO_MORE_DOCS};

use crate::query::ast::{ScoreFunction, FunctionScoreMode, FunctionBoostMode, FieldValueModifier};
use crate::query::{Query, ScorerSupplier, Weight};
use crate::search::searcher::Searcher;
use crate::segment::reader::SegmentReader;

pub struct FunctionScoreQuery {
    pub query: Box<dyn Query>,
    pub functions: Vec<ScoreFunction>,
    pub score_mode: FunctionScoreMode,
    pub boost_mode: FunctionBoostMode,
}

impl Query for FunctionScoreQuery {
    fn create_weight(&self, searcher: &Searcher, score_mode: ScoreMode) -> Result<Box<dyn Weight>> {
        let inner = self.query.create_weight(searcher, score_mode)?;
        Ok(Box::new(FunctionScoreWeight {
            inner,
            functions: self.functions.clone(),
            score_mode: self.score_mode.clone(),
            boost_mode: self.boost_mode.clone(),
        }))
    }
}

struct FunctionScoreWeight {
    inner: Box<dyn Weight>,
    functions: Vec<ScoreFunction>,
    score_mode: FunctionScoreMode,
    boost_mode: FunctionBoostMode,
}

impl Weight for FunctionScoreWeight {
    fn scorer_supplier(&self, reader: &SegmentReader) -> Result<Option<Box<dyn ScorerSupplier>>> {
        let inner = match self.inner.scorer_supplier(reader)? {
            Some(s) => s,
            None => return Ok(None),
        };

        // Pre-load field values for field_value_factor functions
        let mut field_values: Vec<Option<Vec<f64>>> = Vec::new();
        for func in &self.functions {
            match func {
                ScoreFunction::FieldValueFactor { field, missing, .. } => {
                    let field_id = reader.header().fields.iter()
                        .find(|f| f.field_name == *field)
                        .map(|f| f.field_id);
                    if let Some(fid) = field_id {
                        if let Some(col) = reader.column(fid) {
                            let doc_count = col.doc_count();
                            let vals: Vec<f64> = (0..doc_count)
                                .map(|i| col.numeric_value(i).unwrap_or(*missing))
                                .collect();
                            field_values.push(Some(vals));
                        } else {
                            field_values.push(None);
                        }
                    } else {
                        field_values.push(None);
                    }
                }
                _ => field_values.push(None),
            }
        }

        Ok(Some(Box::new(FunctionScoreScorerSupplier {
            inner,
            functions: self.functions.clone(),
            score_mode: self.score_mode.clone(),
            boost_mode: self.boost_mode.clone(),
            field_values,
        })))
    }
}

struct FunctionScoreScorerSupplier {
    inner: Box<dyn ScorerSupplier>,
    functions: Vec<ScoreFunction>,
    score_mode: FunctionScoreMode,
    boost_mode: FunctionBoostMode,
    field_values: Vec<Option<Vec<f64>>>,
}

impl ScorerSupplier for FunctionScoreScorerSupplier {
    fn cost(&self) -> u64 { self.inner.cost() }
    fn scorer(self: Box<Self>) -> Result<Box<dyn Scorer>> {
        let inner = self.inner.scorer()?;
        Ok(Box::new(FunctionScoreScorer {
            inner,
            functions: self.functions,
            score_mode: self.score_mode,
            boost_mode: self.boost_mode,
            field_values: self.field_values,
        }))
    }
}

struct FunctionScoreScorer {
    inner: Box<dyn Scorer>,
    functions: Vec<ScoreFunction>,
    score_mode: FunctionScoreMode,
    boost_mode: FunctionBoostMode,
    field_values: Vec<Option<Vec<f64>>>,
}

impl FunctionScoreScorer {
    fn compute_function_score(&self, doc_id: DocId) -> f32 {
        let mut scores: Vec<f32> = Vec::new();

        for (i, func) in self.functions.iter().enumerate() {
            let s = match func {
                ScoreFunction::Weight(w) => *w,
                ScoreFunction::FieldValueFactor { factor, modifier, missing, .. } => {
                    let val = self.field_values[i]
                        .as_ref()
                        .and_then(|vals| vals.get(doc_id.as_u32() as usize).copied())
                        .unwrap_or(*missing);
                    let modified = apply_modifier(val, modifier);
                    (modified * *factor as f64) as f32
                }
                ScoreFunction::RandomScore { seed } => {
                    // Simple deterministic hash
                    let h = seed.wrapping_mul(6364136223846793005)
                        .wrapping_add(doc_id.as_u32() as u64);
                    ((h >> 33) as f32) / (u32::MAX as f32)
                }
            };
            scores.push(s);
        }

        if scores.is_empty() {
            return 1.0;
        }

        match self.score_mode {
            FunctionScoreMode::Multiply => scores.iter().product(),
            FunctionScoreMode::Sum => scores.iter().sum(),
            FunctionScoreMode::Avg => scores.iter().sum::<f32>() / scores.len() as f32,
            FunctionScoreMode::First => scores[0],
            FunctionScoreMode::Max => scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max),
            FunctionScoreMode::Min => scores.iter().cloned().fold(f32::INFINITY, f32::min),
        }
    }
}

fn apply_modifier(val: f64, modifier: &FieldValueModifier) -> f64 {
    match modifier {
        FieldValueModifier::None => val,
        FieldValueModifier::Log1p => (1.0 + val).log10(),
        FieldValueModifier::Log2p => (2.0 + val).log10(),
        FieldValueModifier::Ln1p => (1.0 + val).ln(),
        FieldValueModifier::Ln2p => (2.0 + val).ln(),
        FieldValueModifier::Sqrt => val.sqrt(),
        FieldValueModifier::Square => val * val,
        FieldValueModifier::Reciprocal => 1.0 / val.max(f64::MIN_POSITIVE),
    }
}

impl Scorer for FunctionScoreScorer {
    fn doc_id(&self) -> DocId { self.inner.doc_id() }
    fn next(&mut self) -> DocId { self.inner.next() }
    fn advance(&mut self, target: DocId) -> DocId { self.inner.advance(target) }

    fn score(&mut self) -> f32 {
        let query_score = self.inner.score();
        let func_score = self.compute_function_score(self.inner.doc_id());

        match self.boost_mode {
            FunctionBoostMode::Multiply => query_score * func_score,
            FunctionBoostMode::Replace => func_score,
            FunctionBoostMode::Sum => query_score + func_score,
            FunctionBoostMode::Avg => (query_score + func_score) / 2.0,
            FunctionBoostMode::Max => query_score.max(func_score),
            FunctionBoostMode::Min => query_score.min(func_score),
        }
    }

    fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> { None }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::query::match_query::MatchQuery;
    use crate::segment::builder::SegmentBuilder;
    use crate::segment::reader::SegmentReader;
    use crate::columnar::writer::ColumnValue;
    use luci_analysis::Token;
    use luci_core::{FieldId, SegmentId};
    use luci_mapping::{FieldType, Mapping};

    fn make_tokens(terms: &[&str]) -> Vec<Token> {
        terms.iter().enumerate()
            .map(|(i, t)| Token::new(*t, 0, t.len(), i as u32))
            .collect()
    }

    #[test]
    fn function_score_weight() {
        let schema = Mapping::builder()
            .field("text", FieldType::Text)
            .build();
        let mut builder = SegmentBuilder::new(SegmentId::new(1), &schema);
        builder.add_document(&[(FieldId::new(0), make_tokens(&["hello", "world"]))], b"{}");
        let reader = SegmentReader::open(builder.build()).unwrap();
        let searcher = Searcher::new(vec![reader]);

        let query = FunctionScoreQuery {
            query: Box::new(MatchQuery {
                field: "text".into(),
                query_text: "hello".into(),
                analyzer: None,
            }),
            functions: vec![ScoreFunction::Weight(2.0)],
            score_mode: FunctionScoreMode::Multiply,
            boost_mode: FunctionBoostMode::Multiply,
        };

        let results = searcher.search_query(&query, 10, 0, &crate::search::SourceFilter::Enabled, &None, false, &[]);
        assert_eq!(results.total_hits.value, 1);
        // Score should be 2x the base BM25 score
        assert!(results.hits[0].score > 0.0);
    }

    #[test]
    fn function_score_field_value_factor() {
        let schema = Mapping::builder()
            .field("text", FieldType::Text)
            .field("popularity", FieldType::Integer)
            .build();
        let mut builder = SegmentBuilder::new(SegmentId::new(1), &schema);

        // Doc 0: popularity 10
        builder.add_document(&[(FieldId::new(0), make_tokens(&["search"]))], b"{}");
        builder.add_column_value(FieldId::new(1), ColumnValue::I64(10));

        // Doc 1: popularity 100
        builder.add_document(&[(FieldId::new(0), make_tokens(&["search"]))], b"{}");
        builder.add_column_value(FieldId::new(1), ColumnValue::I64(100));

        let reader = SegmentReader::open(builder.build()).unwrap();
        let searcher = Searcher::new(vec![reader]);

        let query = FunctionScoreQuery {
            query: Box::new(MatchQuery {
                field: "text".into(),
                query_text: "search".into(),
                analyzer: None,
            }),
            functions: vec![ScoreFunction::FieldValueFactor {
                field: "popularity".into(),
                factor: 1.0,
                modifier: FieldValueModifier::Log1p,
                missing: 1.0,
            }],
            score_mode: FunctionScoreMode::Multiply,
            boost_mode: FunctionBoostMode::Multiply,
        };

        let results = searcher.search_query(&query, 10, 0, &crate::search::SourceFilter::Enabled, &None, false, &[]);
        assert_eq!(results.total_hits.value, 2);
        // Doc 1 (popularity=100) should rank higher
        assert!(results.hits[0].score > results.hits[1].score);
    }
}
