//! KNN query: approximate nearest neighbor search via HNSW.
//!
//! See [[architecture-overview#Step 4]].

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

use crate::query::{Query, ScorerSupplier, Weight};
use crate::search::searcher::Searcher;
use crate::segment::reader::SegmentReader;

/// Approximate k-nearest-neighbor query on a dense_vector field.
pub struct KnnQuery {
    pub field: String,
    pub query_vector: Vec<f32>,
    pub k: usize,
    pub num_candidates: usize,
}

impl Query for KnnQuery {
    fn create_weight(
        &self,
        _searcher: &Searcher,
        _score_mode: ScoreMode,
    ) -> Result<Box<dyn Weight>> {
        Ok(Box::new(KnnWeight {
            field: self.field.clone(),
            query_vector: self.query_vector.clone(),
            k: self.k,
            ef: self.num_candidates,
        }))
    }
}

struct KnnWeight {
    field: String,
    query_vector: Vec<f32>,
    k: usize,
    ef: usize,
}

impl Weight for KnnWeight {
    fn scorer_supplier(
        &self,
        reader: &SegmentReader,
    ) -> Result<Option<Box<dyn ScorerSupplier>>> {
        let field_id = match reader.header().fields.iter()
            .find(|f| f.field_name == self.field)
            .map(|f| f.field_id)
        {
            Some(id) => id,
            None => return Ok(None),
        };

        // Run kNN search now — results are pre-materialized
        let results = reader.knn_search(field_id, &self.query_vector, self.k, self.ef);
        if results.is_empty() {
            return Ok(None);
        }

        Ok(Some(Box::new(KnnScorerSupplier {
            results,
        })))
    }
}

struct KnnScorerSupplier {
    results: Vec<(u32, f32)>,
}

impl ScorerSupplier for KnnScorerSupplier {
    fn cost(&self) -> u64 {
        self.results.len() as u64
    }

    fn scorer(self: Box<Self>) -> Result<Box<dyn Scorer>> {
        Ok(Box::new(KnnScorer {
            results: self.results,
            pos: 0,
        }))
    }
}

/// Scorer that iterates over pre-materialized kNN results.
struct KnnScorer {
    results: Vec<(u32, f32)>, // (doc_id, distance) sorted by distance
    pos: usize,
}

impl Scorer for KnnScorer {
    fn doc_id(&self) -> DocId {
        if self.pos < self.results.len() {
            DocId::new(self.results[self.pos].0)
        } else {
            NO_MORE_DOCS
        }
    }

    fn next(&mut self) -> DocId {
        if self.pos < self.results.len() {
            self.pos += 1;
        }
        self.doc_id()
    }

    fn advance(&mut self, target: DocId) -> DocId {
        while self.pos < self.results.len() && self.results[self.pos].0 < target.as_u32() {
            self.pos += 1;
        }
        self.doc_id()
    }

    fn score(&mut self) -> f32 {
        if self.pos < self.results.len() {
            // Convert distance to score: lower distance → higher score
            1.0 / (1.0 + self.results[self.pos].1)
        } else {
            0.0
        }
    }

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