lucisearch 0.8.1

//! Geospatial query implementations: geo_distance, geo_bounding_box, geo_shape.
//!
//! See [[geospatial]] and [[feature-geo-shape]].

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

use super::geo::{GeoPoint, GeoPointStore, haversine_km};
use super::shape::GeoShapeStore;
use crate::query::ast::SpatialRelation;
use crate::query::{BoundQuery, Query, ScorerSupplier};
use crate::search::searcher::Searcher;
use crate::segment::reader::SegmentReader;

/// Match documents within a distance from a center point.
pub struct GeoDistanceQuery {
    pub field: String,
    pub center: GeoPoint,
    pub distance_km: f64,
}

impl Query for GeoDistanceQuery {
    fn bind(&self, _: &Searcher, _: ScoreMode) -> Result<Box<dyn BoundQuery>> {
        Ok(Box::new(BoundGeoDistanceQuery {
            field: self.field.clone(),
            center: self.center,
            distance_km: self.distance_km,
        }))
    }
}

struct BoundGeoDistanceQuery {
    field: String,
    center: GeoPoint,
    distance_km: f64,
}

impl BoundQuery for BoundGeoDistanceQuery {
    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),
        };

        let store = match reader.geo_points(field_id) {
            Some(s) => s,
            None => return Ok(None),
        };

        Ok(Some(Box::new(GeoDistanceScorerSupplier {
            center: self.center,
            distance_km: self.distance_km,
            store,
        })))
    }
}

struct GeoDistanceScorerSupplier {
    center: GeoPoint,
    distance_km: f64,
    store: GeoPointStore,
}

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

    fn scorer(self: Box<Self>) -> Result<Box<dyn Scorer>> {
        // Compute latitude bounds: 1 degree of latitude ≈ 111 km
        const KM_PER_DEG_LAT: f64 = 111.0;
        let lat_delta = self.distance_km / KM_PER_DEG_LAT;
        let min_lat = self.center.lat - lat_delta;
        let max_lat = self.center.lat + lat_delta;

        // Binary search for lat-bounded candidates, then exact haversine
        let candidates = self.store.docs_in_lat_range(min_lat, max_lat);
        let mut matches = Vec::new();
        for doc_id in candidates {
            if let Some(point) = self.store.get(doc_id) {
                let d = haversine_km(self.center, point);
                if d <= self.distance_km {
                    matches.push((doc_id, d));
                }
            }
        }

        Ok(Box::new(GeoScorer { matches, pos: 0 }))
    }
}

/// Match documents within a bounding box.
pub struct GeoBoundingBoxQuery {
    pub field: String,
    pub top_left: GeoPoint,
    pub bottom_right: GeoPoint,
}

impl Query for GeoBoundingBoxQuery {
    fn bind(&self, _: &Searcher, _: ScoreMode) -> Result<Box<dyn BoundQuery>> {
        Ok(Box::new(BoundGeoBBoxQuery {
            field: self.field.clone(),
            top_left: self.top_left,
            bottom_right: self.bottom_right,
        }))
    }
}

struct BoundGeoBBoxQuery {
    field: String,
    top_left: GeoPoint,
    bottom_right: GeoPoint,
}

impl BoundQuery for BoundGeoBBoxQuery {
    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),
        };

        let store = match reader.geo_points(field_id) {
            Some(s) => s,
            None => return Ok(None),
        };

        Ok(Some(Box::new(GeoBBoxScorerSupplier {
            top_left: self.top_left,
            bottom_right: self.bottom_right,
            store,
        })))
    }
}

struct GeoBBoxScorerSupplier {
    top_left: GeoPoint,
    bottom_right: GeoPoint,
    store: GeoPointStore,
}

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

    fn scorer(self: Box<Self>) -> Result<Box<dyn Scorer>> {
        // Lat range is already defined by the bounding box — only check lon
        let candidates = self
            .store
            .docs_in_lat_range(self.bottom_right.lat, self.top_left.lat);
        let mut matches = Vec::new();
        for doc_id in candidates {
            if let Some(point) = self.store.get(doc_id) {
                if point.lon >= self.top_left.lon && point.lon <= self.bottom_right.lon {
                    matches.push((doc_id, 0.0));
                }
            }
        }

        Ok(Box::new(GeoScorer { matches, pos: 0 }))
    }
}

/// Simple scorer over pre-materialized geo matches.
struct GeoScorer {
    matches: Vec<(u32, f64)>, // (doc_id, distance_km)
    pos: usize,
}

impl Scorer for GeoScorer {
    fn doc_id(&self) -> DocId {
        if self.pos < self.matches.len() {
            DocId::new(self.matches[self.pos].0)
        } else {
            NO_MORE_DOCS
        }
    }
    fn next(&mut self) -> DocId {
        if self.pos < self.matches.len() {
            self.pos += 1;
        }
        self.doc_id()
    }
    fn advance(&mut self, target: DocId) -> DocId {
        while self.pos < self.matches.len() && self.matches[self.pos].0 < target.as_u32() {
            self.pos += 1;
        }
        self.doc_id()
    }
    fn score(&mut self) -> f32 {
        1.0
    }
    fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> {
        None
    }
}

// ---------------------------------------------------------------------------
// geo_shape query
// ---------------------------------------------------------------------------

/// Check if a geometry is an axis-aligned rectangle (4 corners + closing point).
fn is_axis_aligned_rect(geom: &::geo::Geometry<f64>) -> bool {
    if let ::geo::Geometry::Polygon(p) = geom {
        if !p.interiors().is_empty() {
            return false;
        }
        let coords = &p.exterior().0;
        if coords.len() != 5 {
            return false;
        }
        // Check that all x values are one of two distinct values, same for y
        let xs: std::collections::BTreeSet<u64> =
            coords[..4].iter().map(|c| c.x.to_bits()).collect();
        let ys: std::collections::BTreeSet<u64> =
            coords[..4].iter().map(|c| c.y.to_bits()).collect();
        xs.len() == 2 && ys.len() == 2
    } else {
        false
    }
}

/// Match documents whose geo_shape field satisfies a spatial relation
/// with a query shape. Uses R-tree for candidate selection and the `geo`
/// crate for exact [[de-9im]] predicate evaluation.
///
/// See [[feature-geo-shape]].
pub struct GeoShapeQuery {
    pub field: String,
    pub query_shape: ::geo::Geometry<f64>,
    pub query_bbox: (f64, f64, f64, f64),
    pub relation: SpatialRelation,
}

impl Query for GeoShapeQuery {
    fn bind(&self, _: &Searcher, _: ScoreMode) -> Result<Box<dyn BoundQuery>> {
        Ok(Box::new(BoundGeoShapeQuery {
            field: self.field.clone(),
            query_shape: self.query_shape.clone(),
            query_bbox: self.query_bbox,
            relation: self.relation.clone(),
        }))
    }
}

struct BoundGeoShapeQuery {
    field: String,
    query_shape: ::geo::Geometry<f64>,
    query_bbox: (f64, f64, f64, f64),
    relation: SpatialRelation,
}

impl BoundQuery for BoundGeoShapeQuery {
    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),
        };

        let store = match reader.geo_shapes(field_id) {
            Some(s) => s,
            None => return Ok(None),
        };

        Ok(Some(Box::new(GeoShapeScorerSupplier {
            query_shape: self.query_shape.clone(),
            query_bbox: self.query_bbox,
            relation: self.relation.clone(),
            store,
        })))
    }
}

struct GeoShapeScorerSupplier {
    query_shape: ::geo::Geometry<f64>,
    query_bbox: (f64, f64, f64, f64),
    relation: SpatialRelation,
    store: GeoShapeStore,
}

impl GeoShapeScorerSupplier {
    /// Optimized disjoint: find the small set of docs that DO intersect
    /// via R-tree, verify with exact predicate, then return everything else.
    fn score_disjoint(self, rtree_data: &[u8]) -> Result<Box<dyn Scorer>> {
        use ::geo::algorithm::Intersects;

        let intersecting_candidates = GeoShapeStore::search_rtree(
            rtree_data,
            self.query_bbox,
            self.store.shape_offsets_ref(),
        );

        // Verify which candidates truly intersect (bbox overlap ≠ shape overlap)
        let mut intersecting = std::collections::HashSet::new();
        for doc_id in intersecting_candidates {
            if let Some(doc_shape) = self.store.get_shape(doc_id) {
                if self.query_shape.intersects(&doc_shape) {
                    intersecting.insert(doc_id);
                }
            }
        }

        // Return all non-null docs NOT in the intersecting set
        let mut matches = Vec::new();
        for doc_id in 0..self.store.len() as u32 {
            if self.store.get_bbox(doc_id).is_some() && !intersecting.contains(&doc_id) {
                matches.push((doc_id, 0.0));
            }
        }

        Ok(Box::new(GeoScorer { matches, pos: 0 }))
    }
}

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

    fn scorer(self: Box<Self>) -> Result<Box<dyn Scorer>> {
        use ::geo::algorithm::{Contains, Intersects, Relate};

        // Use pre-built R-tree from segment, fall back to building if empty
        let rtree_data_owned;
        let rtree_data = if !self.store.rtree_data().is_empty() {
            self.store.rtree_data()
        } else {
            rtree_data_owned = self.store.build_rtree();
            &rtree_data_owned
        };
        // For disjoint: find docs that DO intersect via R-tree, then return
        // the complement. Much faster than iterating all docs.
        if self.relation == SpatialRelation::Disjoint {
            let rtree_owned = if self.store.rtree_data().is_empty() {
                self.store.build_rtree()
            } else {
                self.store.rtree_data().to_vec()
            };
            return self.score_disjoint(&rtree_owned);
        }

        let candidates = GeoShapeStore::search_rtree(
            rtree_data,
            self.query_bbox,
            self.store.shape_offsets_ref(),
        );

        let (q_min_x, q_min_y, q_max_x, q_max_y) = self.query_bbox;

        // Check if query shape is axis-aligned rect (bbox == exact geometry).
        // When true, bbox containment proves exact containment.
        let query_is_rect = matches!(self.query_shape, ::geo::Geometry::Rect(_))
            || is_axis_aligned_rect(&self.query_shape);

        let mut matches = Vec::new();
        for doc_id in candidates {
            let doc_bbox = match self.store.get_bbox(doc_id) {
                Some(bb) => bb,
                None => continue,
            };
            let (d_min_x, d_min_y, d_max_x, d_max_y) = doc_bbox;

            // Bbox pre-filter: skip docs that can't possibly match
            match self.relation {
                SpatialRelation::Within | SpatialRelation::CoveredBy => {
                    if d_min_x < q_min_x
                        || d_min_y < q_min_y
                        || d_max_x > q_max_x
                        || d_max_y > q_max_y
                    {
                        continue;
                    }
                }
                SpatialRelation::Contains
                | SpatialRelation::Covers
                | SpatialRelation::ContainsProperly => {
                    if q_min_x < d_min_x
                        || q_min_y < d_min_y
                        || q_max_x > d_max_x
                        || q_max_y > d_max_y
                    {
                        continue;
                    }
                }
                _ => {}
            }

            // Fast path 1: rect query — if doc bbox is fully inside the
            // query bbox, then doc shape ⊆ doc bbox ⊆ query rect.
            if query_is_rect {
                match self.relation {
                    SpatialRelation::Within | SpatialRelation::CoveredBy => {
                        matches.push((doc_id, 0.0));
                        continue;
                    }
                    _ => {}
                }
            }

            // Fast path 2: rect doc — if query bbox is fully inside the
            // doc bbox and the doc is a rect, doc contains query.
            if self.store.is_rect_shape(doc_id) {
                let bbox_proves_contains = q_min_x >= d_min_x
                    && q_min_y >= d_min_y
                    && q_max_x <= d_max_x
                    && q_max_y <= d_max_y;
                if bbox_proves_contains {
                    match self.relation {
                        SpatialRelation::Contains | SpatialRelation::Covers => {
                            matches.push((doc_id, 0.0));
                            continue;
                        }
                        _ => {}
                    }
                }
            }

            // Fast path 3: non-rect query but all 4 doc bbox corners are
            // inside the query shape → doc shape ⊆ doc bbox ⊆ query shape.
            // Uses point-in-polygon which is much cheaper than shape-contains-shape.
            match self.relation {
                SpatialRelation::Within | SpatialRelation::CoveredBy => {
                    let corners = [
                        ::geo::Coord {
                            x: d_min_x,
                            y: d_min_y,
                        },
                        ::geo::Coord {
                            x: d_max_x,
                            y: d_min_y,
                        },
                        ::geo::Coord {
                            x: d_max_x,
                            y: d_max_y,
                        },
                        ::geo::Coord {
                            x: d_min_x,
                            y: d_max_y,
                        },
                    ];
                    let all_inside = corners
                        .iter()
                        .all(|c| self.query_shape.contains(&::geo::Point(*c)));
                    if all_inside {
                        matches.push((doc_id, 0.0));
                        continue;
                    }
                }
                _ => {}
            }

            let doc_shape = match self.store.get_shape(doc_id) {
                Some(s) => s,
                None => continue,
            };

            // Use fast-path traits for common predicates, fall back to
            // full DE-9IM Relate only when needed.
            let matched = match self.relation {
                SpatialRelation::Intersects => self.query_shape.intersects(&doc_shape),
                SpatialRelation::Disjoint => !self.query_shape.intersects(&doc_shape),
                SpatialRelation::Contains | SpatialRelation::Covers => {
                    doc_shape.contains(&self.query_shape)
                }
                SpatialRelation::Within | SpatialRelation::CoveredBy => {
                    self.query_shape.contains(&doc_shape)
                }
                // Full DE-9IM for predicates that need the intersection matrix
                _ => {
                    let im = doc_shape.relate(&self.query_shape);
                    match self.relation {
                        SpatialRelation::Touches => im.is_touches(),
                        SpatialRelation::Crosses => im.is_crosses(),
                        SpatialRelation::Overlaps => im.is_overlaps(),
                        SpatialRelation::Equals => im.is_equal_topo(),
                        SpatialRelation::Covers => im.is_covers(),
                        SpatialRelation::CoveredBy => im.is_coveredby(),
                        SpatialRelation::ContainsProperly => im.is_contains_properly(),
                        _ => unreachable!(),
                    }
                }
            };

            if matched {
                matches.push((doc_id, 0.0));
            }
        }

        matches.sort_by_key(|m| m.0);
        Ok(Box::new(GeoScorer { matches, pos: 0 }))
    }
}