// Copyright (c) 2023-2025 ParadeDB, Inc.
//
// This file is part of ParadeDB - Postgres for Search and Analytics
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

use std::ops::{Bound, RangeBounds};
use tantivy::query::{EmptyScorer, EnableScoring, Explanation, Query, QueryClone, Scorer, Weight};
use tantivy::schema::Field;
use tantivy::{DocId, DocSet, Score, SegmentReader, Term, TERMINATED};

#[derive(Debug)]
pub struct ScoreFilter {
    bounds: Vec<(Bound<f32>, Bound<f32>)>,
    query: Box<dyn Query>,
}

impl QueryClone for ScoreFilter {
    fn box_clone(&self) -> Box<dyn Query> {
        Box::new(Self {
            bounds: self.bounds.clone(),
            query: self.query.box_clone(),
        })
    }
}

impl ScoreFilter {
    pub fn new(bounds: Vec<(Bound<f32>, Bound<f32>)>, query: Box<dyn Query>) -> Self {
        Self { bounds, query }
    }
}

struct ScoreFilterWeight {
    bounds: Vec<(Bound<f32>, Bound<f32>)>,
    weight: Box<dyn Weight>,
}

impl Weight for ScoreFilterWeight {
    fn scorer(&self, reader: &SegmentReader, boost: Score) -> tantivy::Result<Box<dyn Scorer>> {
        let mut scorer = self.weight.scorer(reader, boost)?;

        while !self
            .bounds
            .iter()
            .any(|&(lower, upper)| (lower, upper).contains(&scorer.score()))
        {
            if scorer.advance() == TERMINATED {
                scorer = Box::new(EmptyScorer);
                break;
            }
        }

        Ok(Box::new(ScoreFilterScorer {
            bounds: self.bounds.clone(),
            scorer,
        }))
    }

    fn explain(&self, reader: &SegmentReader, doc: DocId) -> tantivy::Result<Explanation> {
        self.weight.explain(reader, doc)
    }
}

struct ScoreFilterScorer {
    bounds: Vec<(Bound<f32>, Bound<f32>)>,
    scorer: Box<dyn Scorer>,
}

impl Scorer for ScoreFilterScorer {
    fn score(&mut self) -> Score {
        self.scorer.score()
    }
}

impl DocSet for ScoreFilterScorer {
    fn advance(&mut self) -> DocId {
        loop {
            let doc = self.scorer.advance();
            if doc == TERMINATED
                || self
                    .bounds
                    .iter()
                    .any(|&(lower, upper)| (lower, upper).contains(&self.scorer.score()))
            {
                return doc;
            }
        }
    }

    fn doc(&self) -> DocId {
        self.scorer.doc()
    }

    fn size_hint(&self) -> u32 {
        self.scorer.size_hint()
    }
}

impl Query for ScoreFilter {
    fn weight(&self, enable_scoring: EnableScoring<'_>) -> tantivy::Result<Box<dyn Weight>> {
        Ok(Box::new(ScoreFilterWeight {
            bounds: self.bounds.clone(),
            weight: self.query.weight(enable_scoring)?,
        }))
    }

    fn query_terms(
        &self,
        field: Field,
        segment_reader: &SegmentReader,
        visitor: &mut dyn FnMut(&Term, bool),
    ) {
        self.query.query_terms(field, segment_reader, visitor)
    }
}
