use bevy::{
    math::bounding::{Aabb2d, BoundingVolume, IntersectsVolume},
    prelude::Vec2,
};

use crate::geom::Point;

#[derive(Debug, Clone)]
pub enum Entry {
    Leaf { mbr: Aabb2d, object: Point },
    Node { mbr: Aabb2d, child: TreeNode },
}

impl Entry {
    /// Returns a reference to the minimum bounding volume for this entry.
    pub fn mbr(&self) -> &Aabb2d {
        match self {
            Entry::Leaf { mbr, .. } | Entry::Node { mbr, .. } => mbr,
        }
    }
}

#[derive(Debug, Clone)]
pub struct TreeNode {
    pub entries: Vec<Entry>,
    pub is_leaf: bool,
}

#[derive(Debug, Clone)]
pub struct RStarTree {
    root: TreeNode,
    max_entries: usize,
}

impl Entry {
    fn as_leaf_obj(&self) -> Option<&Point> {
        match self {
            Entry::Leaf { object, .. } => Some(object),
            _ => None,
        }
    }
    fn child(&self) -> Option<&TreeNode> {
        match self {
            Entry::Node { child, .. } => Some(child),
            _ => None,
        }
    }
}

impl TreeNode {
    fn is_leaf(&self) -> bool {
        self.is_leaf
    }

    fn entries(&self) -> &[Entry] {
        &self.entries
    }

    fn range_search_bbox<'s>(&'s self, bbox: &Aabb2d, out: &mut Vec<&'s Point>) {
        if self.is_leaf() {
            for entry in self.entries() {
                if let Some(obj) = entry.as_leaf_obj()
                    && entry.mbr().intersects(bbox)
                {
                    out.push(obj);
                }
            }
        } else {
            for entry in self.entries() {
                if let Some(child) = entry.child()
                    && entry.mbr().intersects(bbox)
                {
                    child.range_search_bbox(bbox, out);
                }
            }
        }
    }

    fn mbr(&self) -> Option<Aabb2d> {
        entries_mbr(self.entries())
    }
}

impl RStarTree {
    pub fn new(max_entries: usize) -> Self {
        let max_entries = max_entries.max(4);
        RStarTree {
            root: TreeNode {
                entries: Vec::new(),
                is_leaf: true,
            },
            max_entries,
        }
    }

    pub fn range_search(&self, query_point: &Point, radius: f32) -> Vec<&Point> {
        let query_bbox = Aabb2d::new(query_point.point, Vec2::splat(radius));
        let r2 = radius.powi(2);
        let candidates = self.range_search_bbox(&query_bbox);
        candidates
            .into_iter()
            .filter(|&other| query_point.point.distance_squared(other.point) <= r2)
            .collect()
    }

    pub fn range_search_bbox(&self, query_bbox: &Aabb2d) -> Vec<&Point> {
        let mut res = Vec::with_capacity(self.root.entries.len() / 10);
        self.root.range_search_bbox(query_bbox, &mut res);
        res
    }

    pub fn insert_bulk(&mut self, objects: Vec<Point>) {
        if objects.is_empty() {
            return;
        }

        let mut entries: Vec<Entry> = objects
            .into_iter()
            .map(|point| Entry::Leaf {
                mbr: point.mbr(),
                object: point,
            })
            .collect();

        while entries.len() > self.max_entries {
            let mut new_level_entries = Vec::new();
            let chunks = entries.chunks(self.max_entries);

            for chunk in chunks {
                let child = TreeNode {
                    entries: chunk.to_vec(),
                    is_leaf: self.root.is_leaf,
                };
                if let Some(mbr) = child.mbr() {
                    new_level_entries.push(Entry::Node { mbr, child });
                }
            }
            entries = new_level_entries;
            self.root.is_leaf = false;
        }

        self.root.entries.extend(entries);
    }
}

fn entries_mbr(entries: &[Entry]) -> Option<Aabb2d> {
    let mut iter = entries.iter();
    let first = *iter.next()?.mbr();
    Some(iter.fold(first, |acc, entry| acc.merge(entry.mbr())))
}