plozone 0.1.2

//! R-tree-indexed zone store.
//!
//! [`ZoneStore`] holds every zone as a `Box<dyn ZoneShape>` keyed by id, indexed
//! by ENU AABB in an R-tree. Point queries prune candidates by bounding box
//! first, then run the exact `contains_enu` test only on survivors.

use std::collections::HashMap;

use rstar::{AABB, PointDistance, RTree, RTreeObject};
use smallvec::SmallVec;

use crate::coord::CoordSystem;
use crate::zone::{Zone, ZoneEntry, ZoneShape, zone_to_shape};

/// Placeholder shape used only to build a removal key — its geometry is never
/// queried because [`ZoneRecord`] compares by id alone.
struct DummyShape;
impl ZoneShape for DummyShape {
    fn contains_enu(&self, _: [f64; 3]) -> bool {
        false
    }
    fn aabb_enu(&self) -> [f64; 6] {
        [0.0; 6]
    }
}

/// Internal R-tree record: an ENU shape plus its original id, metadata, and cached AABB.
struct ZoneRecord {
    id: u32,
    priority: u8,
    layer: u8,
    shape: Box<dyn ZoneShape>,
    aabb: [f64; 6],
}

// Equality by id only, so `RTree::remove` can match on a key record without
// reconstructing the real shape.
impl PartialEq for ZoneRecord {
    fn eq(&self, o: &Self) -> bool {
        self.id == o.id
    }
}
impl Eq for ZoneRecord {}

impl RTreeObject for ZoneRecord {
    type Envelope = AABB<[f64; 3]>;
    fn envelope(&self) -> Self::Envelope {
        AABB::from_corners(
            [self.aabb[0], self.aabb[1], self.aabb[2]],
            [self.aabb[3], self.aabb[4], self.aabb[5]],
        )
    }
}

// Required so `locate_all_at_point` can prune by envelope. The default
// `contains_point` (AABB membership) is exactly the pruning we want; we only
// need a `distance_2` to the bounding box.
impl PointDistance for ZoneRecord {
    fn distance_2(&self, point: &[f64; 3]) -> f64 {
        let mut d2 = 0.0;
        for (i, &v) in point.iter().enumerate() {
            let (lo, hi) = (self.aabb[i], self.aabb[i + 3]);
            if v < lo {
                d2 += (lo - v) * (lo - v);
            } else if v > hi {
                d2 += (v - hi) * (v - hi);
            }
        }
        d2
    }
}

/// A spatial index of zones, queryable by ENU point, geodetic point, or region.
pub struct ZoneStore {
    index: RTree<ZoneRecord>,
    /// Side index of id → AABB so removals can rebuild the R-tree key in O(1).
    id_to_aabb: HashMap<u32, [f64; 6]>,
    /// Side index of id → (priority, layer) for filtered/sorted queries.
    id_to_meta: HashMap<u32, (u8, u8)>,
}

impl ZoneStore {
    /// Build from a list of serializable entries, converting each into the
    /// store's internal metric space via `conv`.
    pub fn from_entries<C: CoordSystem + ?Sized>(entries: &[ZoneEntry], conv: &C) -> Self {
        let mut id_to_aabb = HashMap::new();
        let records: Vec<ZoneRecord> = entries
            .iter()
            .map(|e| {
                let shape = zone_to_shape(&e.zone, conv);
                let aabb = shape.aabb_enu();
                id_to_aabb.insert(e.id, aabb);
                ZoneRecord { id: e.id, priority: e.priority, layer: e.layer, shape, aabb }
            })
            .collect();
        let id_to_meta: HashMap<u32, (u8, u8)> = entries
            .iter()
            .map(|e| (e.id, (e.priority, e.layer)))
            .collect();
        Self { index: RTree::bulk_load(records), id_to_aabb, id_to_meta }
    }

    /// Add a built-in [`Zone`], converting it via `conv`.
    pub fn add_zone<C: CoordSystem + ?Sized>(&mut self, id: u32, zone: &Zone, conv: &C) {
        let shape = zone_to_shape(zone, conv);
        self.add_record(id, 0, 0, shape);
    }

    /// Add a [`ZoneEntry`] (with priority and layer) converting it via `conv`.
    pub fn add_entry<C: CoordSystem + ?Sized>(&mut self, entry: &ZoneEntry, conv: &C) {
        let shape = zone_to_shape(&entry.zone, conv);
        self.add_record(entry.id, entry.priority, entry.layer, shape);
    }

    /// Add a custom, already-ENU [`ZoneShape`] (local-only, not wire-serializable).
    pub fn add_custom(&mut self, id: u32, shape: Box<dyn ZoneShape>) {
        self.add_record(id, 0, 0, shape);
    }

    /// Add a custom shape with explicit priority and layer.
    pub fn add_custom_with_meta(
        &mut self,
        id: u32,
        priority: u8,
        layer: u8,
        shape: Box<dyn ZoneShape>,
    ) {
        self.add_record(id, priority, layer, shape);
    }

    fn add_record(&mut self, id: u32, priority: u8, layer: u8, shape: Box<dyn ZoneShape>) {
        if self.id_to_aabb.contains_key(&id) {
            self.remove(id);
        }
        let aabb = shape.aabb_enu();
        self.id_to_aabb.insert(id, aabb);
        self.id_to_meta.insert(id, (priority, layer));
        self.index.insert(ZoneRecord { id, priority, layer, shape, aabb });
    }

    /// Remove a zone by id in O(log n). No-op if the id is unknown.
    pub fn remove(&mut self, id: u32) {
        let Some(&aabb) = self.id_to_aabb.get(&id) else {
            return;
        };
        let (priority, layer) = self.id_to_meta.get(&id).copied().unwrap_or((0, 0));
        let key = ZoneRecord { id, priority, layer, shape: Box::new(DummyShape), aabb };
        self.index.remove(&key);
        self.id_to_aabb.remove(&id);
        self.id_to_meta.remove(&id);
    }

    /// All zone ids containing the ENU point `p`.
    #[must_use]
    pub fn query_enu(&self, p: [f64; 3]) -> SmallVec<[u32; 8]> {
        self.index
            .locate_all_at_point(p)
            .filter(|r| r.shape.contains_enu(p))
            .map(|r| r.id)
            .collect()
    }

    /// All zone ids containing the user-space point, converted via `conv`.
    #[must_use]
    pub fn query_geodetic<C: CoordSystem + ?Sized>(
        &self,
        lat: f64,
        lon: f64,
        alt: f64,
        conv: &C,
    ) -> SmallVec<[u32; 8]> {
        self.query_enu(conv.to_internal([lat, lon, alt]))
    }

    /// All zone ids whose AABB overlaps the user-space bounding box `[min, max]`.
    /// This is a coarse, AABB-level test — it does not run `contains_enu`.
    #[must_use]
    pub fn query_region<C: CoordSystem + ?Sized>(
        &self,
        min_geo: [f64; 3],
        max_geo: [f64; 3],
        conv: &C,
    ) -> SmallVec<[u32; 8]> {
        let lo = conv.to_internal(min_geo);
        let hi = conv.to_internal(max_geo);
        let corner_lo = std::array::from_fn(|i| lo[i].min(hi[i]));
        let corner_hi = std::array::from_fn(|i| lo[i].max(hi[i]));
        self.index
            .locate_in_envelope_intersecting(AABB::from_corners(corner_lo, corner_hi))
            .map(|r| r.id)
            .collect()
    }

    /// All zone ids containing the ENU point `p` that belong to `layer`.
    #[must_use]
    pub fn query_enu_by_layer(&self, p: [f64; 3], layer: u8) -> SmallVec<[u32; 8]> {
        self.index
            .locate_all_at_point(p)
            .filter(|r| r.layer == layer && r.shape.contains_enu(p))
            .map(|r| r.id)
            .collect()
    }

    /// Id of the highest-priority zone containing `p`, or `None` if none.
    ///
    /// Ties on priority are broken by smaller id.
    #[must_use]
    pub fn query_enu_top_priority(&self, p: [f64; 3]) -> Option<u32> {
        self.index
            .locate_all_at_point(p)
            .filter(|r| r.shape.contains_enu(p))
            .max_by(|a, b| a.priority.cmp(&b.priority).then_with(|| b.id.cmp(&a.id)))
            .map(|r| r.id)
    }

    /// Number of zones in the store.
    pub fn len(&self) -> usize {
        self.index.size()
    }

    /// True if the store holds no zones.
    pub fn is_empty(&self) -> bool {
        self.index.size() == 0
    }

    /// Every zone id currently stored.
    #[must_use]
    pub fn ids(&self) -> Vec<u32> {
        self.index.iter().map(|r| r.id).collect()
    }

    /// Does `zone_id` contain the ENU point? Single-zone membership test.
    pub fn zone_contains(&self, zone_id: u32, p: [f64; 3]) -> bool {
        self.index
            .locate_all_at_point(p)
            .any(|r| r.id == zone_id && r.shape.contains_enu(p))
    }

    /// The ENU AABB of `zone_id`, if present.
    pub fn zone_aabb(&self, zone_id: u32) -> Option<[f64; 6]> {
        self.id_to_aabb.get(&zone_id).copied()
    }
}

/// A serializable change to a zone collection — the wire/diff primitive.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq)]
pub enum ZoneDiff {
    /// A new zone was added (carries the full entry).
    Add(ZoneEntry),
    /// The zone with this id was removed.
    Remove { id: u32 },
    /// The zone with this id was replaced by a new shape.
    Modify { id: u32, zone: Zone },
    /// A hole scan for this zone has finished and results are available.
    ScanReady { zone_id: u32 },
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::coord::EnuConverter;

    fn conv() -> EnuConverter {
        EnuConverter::new(10.7626, 106.6601, 0.0)
    }

    fn sample_store() -> (ZoneStore, EnuConverter) {
        let c = conv();
        let store = ZoneStore::from_entries(
            &[ZoneEntry::new(
                1,
                Zone::Cylinder {
                    center: [10.7626, 106.6601],
                    radius_m: 50.0,
                    z_min: 0.0,
                    z_max: 20.0,
                },
            )],
            &c,
        );
        (store, c)
    }

    #[test]
    fn query_geodetic_hits_zone() {
        let (store, c) = sample_store();
        let expected: SmallVec<[u32; 8]> = smallvec::smallvec![1];
        assert_eq!(store.query_geodetic(10.7626, 106.6601, 5.0, &c), expected);
        // Well above the cylinder — no hit.
        assert!(store.query_geodetic(10.7626, 106.6601, 50.0, &c).is_empty());
    }

    #[test]
    fn add_and_remove() {
        let (mut store, c) = sample_store();
        assert_eq!(store.len(), 1);
        store.add_zone(
            2,
            &Zone::Aabb { min: [10.7620, 106.6595, 0.0], max: [10.7630, 106.6605, 30.0] },
            &c,
        );
        assert_eq!(store.len(), 2);
        let mut ids = store.ids();
        ids.sort();
        assert_eq!(ids, vec![1, 2]);

        store.remove(1);
        assert_eq!(store.len(), 1);
        assert_eq!(store.ids(), vec![2]);
        // Removing an unknown id is a no-op.
        store.remove(999);
        assert_eq!(store.len(), 1);
    }

    #[test]
    fn re_adding_same_id_replaces() {
        let (mut store, c) = sample_store();
        store.add_zone(
            1,
            &Zone::Cylinder { center: [10.7626, 106.6601], radius_m: 5.0, z_min: 0.0, z_max: 2.0 },
            &c,
        );
        // Still a single zone with id 1, now the smaller cylinder.
        assert_eq!(store.len(), 1);
        assert!(store.query_geodetic(10.7626, 106.6601, 10.0, &c).is_empty(), "shrunk z range");
    }

    #[test]
    fn zone_contains_and_aabb() {
        let (store, c) = sample_store();
        let center = c.to_enu(10.7626, 106.6601, 5.0);
        assert!(store.zone_contains(1, center));
        assert!(!store.zone_contains(1, [1000.0, 1000.0, 5.0]));
        let bb = store.zone_aabb(1).unwrap();
        assert!(bb[3] - bb[0] >= 99.0, "cylinder diameter ≈ 100 m: {bb:?}");
        assert!(store.zone_aabb(7).is_none());
    }

    #[test]
    fn custom_shape_via_store() {
        use crate::zone::ShrinkingZone;
        use std::sync::Arc;
        use std::sync::atomic::AtomicU32;

        let (mut store, _c) = sample_store();
        store.add_custom(
            2,
            Box::new(ShrinkingZone {
                cx: 0.0,
                cy: 0.0,
                z_min: 0.0,
                z_max: 50.0,
                radius: Arc::new(AtomicU32::new(100_000)), // 100 m
            }),
        );
        assert!(store.query_enu([10.0, 10.0, 5.0]).contains(&2));
    }

    #[test]
    fn query_by_layer_filters_correctly() {
        let c = conv();
        let entries = vec![
            ZoneEntry::new(1, Zone::Cylinder { center: [10.7626, 106.6601], radius_m: 50.0, z_min: 0.0, z_max: 10.0 })
                .with_layer(0),
            ZoneEntry::new(2, Zone::Cylinder { center: [10.7626, 106.6601], radius_m: 50.0, z_min: 0.0, z_max: 10.0 })
                .with_layer(1),
            ZoneEntry::new(3, Zone::Cylinder { center: [10.7626, 106.6601], radius_m: 50.0, z_min: 0.0, z_max: 10.0 })
                .with_layer(0),
        ];
        let store = ZoneStore::from_entries(&entries, &c);
        let p = c.to_enu(10.7626, 106.6601, 5.0);
        let mut layer0 = store.query_enu_by_layer(p, 0);
        layer0.sort();
        let expected: SmallVec<[u32; 8]> = smallvec::smallvec![1, 3];
        assert_eq!(layer0, expected);
        let expected: SmallVec<[u32; 8]> = smallvec::smallvec![2];
        assert_eq!(store.query_enu_by_layer(p, 1), expected);
        assert_eq!(store.query_enu_by_layer(p, 99), SmallVec::<[u32; 8]>::new());
    }

    #[test]
    fn top_priority_returns_highest() {
        let c = conv();
        let entries = vec![
            ZoneEntry::new(10, Zone::Cylinder { center: [10.7626, 106.6601], radius_m: 50.0, z_min: 0.0, z_max: 10.0 })
                .with_priority(1),
            ZoneEntry::new(20, Zone::Cylinder { center: [10.7626, 106.6601], radius_m: 50.0, z_min: 0.0, z_max: 10.0 })
                .with_priority(5),
            ZoneEntry::new(30, Zone::Cylinder { center: [10.7626, 106.6601], radius_m: 50.0, z_min: 0.0, z_max: 10.0 })
                .with_priority(3),
        ];
        let store = ZoneStore::from_entries(&entries, &c);
        let p = c.to_enu(10.7626, 106.6601, 5.0);
        assert_eq!(store.query_enu_top_priority(p), Some(20));
    }

    #[test]
    fn top_priority_tie_breaks_by_smaller_id() {
        let c = conv();
        let entries = vec![
            ZoneEntry::new(5, Zone::Cylinder { center: [10.7626, 106.6601], radius_m: 50.0, z_min: 0.0, z_max: 10.0 })
                .with_priority(10),
            ZoneEntry::new(2, Zone::Cylinder { center: [10.7626, 106.6601], radius_m: 50.0, z_min: 0.0, z_max: 10.0 })
                .with_priority(10),
        ];
        let store = ZoneStore::from_entries(&entries, &c);
        let p = c.to_enu(10.7626, 106.6601, 5.0);
        assert_eq!(store.query_enu_top_priority(p), Some(2));
    }

    #[test]
    fn top_priority_none_when_no_zone_contains() {
        let c = conv();
        let store = ZoneStore::from_entries(&[], &c);
        assert_eq!(store.query_enu_top_priority([0.0, 0.0, 0.0]), None);
    }

    #[test]
    fn add_entry_preserves_priority_and_layer() {
        let c = conv();
        let mut store = ZoneStore::from_entries(&[], &c);
        store.add_entry(
            &ZoneEntry::new(1, Zone::Cylinder { center: [10.7626, 106.6601], radius_m: 20.0, z_min: 0.0, z_max: 5.0 })
                .with_priority(7)
                .with_layer(3),
            &c,
        );
        let p = c.to_enu(10.7626, 106.6601, 2.5);
        let expected: SmallVec<[u32; 8]> = smallvec::smallvec![1];
        assert_eq!(store.query_enu_by_layer(p, 3), expected);
        assert_eq!(store.query_enu_by_layer(p, 0), SmallVec::<[u32; 8]>::new());
        assert_eq!(store.query_enu_top_priority(p), Some(1));
    }
}