use std::collections::HashSet;

use nalgebra_glm::Vec3;

use super::primitives::Aabb;

const REBUILD_THRESHOLD: usize = 64;
const REMOVAL_RATIO_THRESHOLD: f32 = 0.25;

pub struct SdfEditBvh {
    nodes: Vec<BvhNode>,
    edit_indices: Vec<usize>,
    pending_adds: Vec<(usize, Aabb)>,
    removed_indices: HashSet<usize>,
    total_indexed: usize,
}

struct BvhNode {
    bounds: Aabb,
    left_first: u32,
    count: u32,
}

impl BvhNode {
    fn is_leaf(&self) -> bool {
        self.count > 0
    }
}

impl Default for SdfEditBvh {
    fn default() -> Self {
        Self::new()
    }
}

impl SdfEditBvh {
    pub fn new() -> Self {
        Self {
            nodes: Vec::new(),
            edit_indices: Vec::new(),
            pending_adds: Vec::new(),
            removed_indices: HashSet::new(),
            total_indexed: 0,
        }
    }

    pub fn build(&mut self, edit_bounds: &[Aabb]) {
        self.nodes.clear();
        self.edit_indices.clear();
        self.pending_adds.clear();
        self.removed_indices.clear();

        if edit_bounds.is_empty() {
            self.total_indexed = 0;
            return;
        }

        self.edit_indices = (0..edit_bounds.len()).collect();
        self.total_indexed = edit_bounds.len();

        let root_bounds = compute_bounds(edit_bounds, &self.edit_indices, 0, edit_bounds.len());
        self.nodes.push(BvhNode {
            bounds: root_bounds,
            left_first: 0,
            count: edit_bounds.len() as u32,
        });

        self.subdivide(0, edit_bounds);
    }

    pub fn add_edit(&mut self, index: usize, bounds: Aabb) {
        self.pending_adds.push((index, bounds));
        self.removed_indices.remove(&index);
    }

    pub fn remove_edit(&mut self, index: usize) {
        self.pending_adds.retain(|(idx, _)| *idx != index);
        if index < self.total_indexed {
            self.removed_indices.insert(index);
        }
    }

    pub fn update_edit(&mut self, index: usize, bounds: Aabb) {
        self.remove_edit(index);
        self.add_edit(index, bounds);
    }

    pub fn needs_rebuild(&self) -> bool {
        if self.pending_adds.len() >= REBUILD_THRESHOLD {
            return true;
        }
        if self.total_indexed > 0 {
            let removal_ratio = self.removed_indices.len() as f32 / self.total_indexed as f32;
            if removal_ratio >= REMOVAL_RATIO_THRESHOLD {
                return true;
            }
        }
        false
    }

    pub fn rebuild_if_needed(&mut self, edit_bounds: &[Aabb]) {
        if self.needs_rebuild() || (self.nodes.is_empty() && !edit_bounds.is_empty()) {
            self.build(edit_bounds);
        }
    }

    fn subdivide(&mut self, node_index: usize, edit_bounds: &[Aabb]) {
        let node = &self.nodes[node_index];
        let first = node.left_first as usize;
        let count = node.count as usize;

        if count <= 2 {
            return;
        }

        let extent = node.bounds.half_extents();
        let axis = if extent.x >= extent.y && extent.x >= extent.z {
            0
        } else if extent.y >= extent.z {
            1
        } else {
            2
        };

        let split_pos = node.bounds.center()[axis];

        let mut partition_point = first;
        let mut scan = first;
        while scan < first + count {
            let edit_idx = self.edit_indices[scan];
            let center = edit_bounds[edit_idx].center()[axis];
            if center < split_pos {
                self.edit_indices.swap(partition_point, scan);
                partition_point += 1;
            }
            scan += 1;
        }

        let left_count = partition_point - first;
        let right_count = count - left_count;

        if left_count == 0 || right_count == 0 {
            return;
        }

        let left_bounds = compute_bounds(edit_bounds, &self.edit_indices, first, left_count);
        let right_bounds = compute_bounds(
            edit_bounds,
            &self.edit_indices,
            partition_point,
            right_count,
        );

        let left_child_index = self.nodes.len();
        self.nodes.push(BvhNode {
            bounds: left_bounds,
            left_first: first as u32,
            count: left_count as u32,
        });

        let right_child_index = self.nodes.len();
        self.nodes.push(BvhNode {
            bounds: right_bounds,
            left_first: partition_point as u32,
            count: right_count as u32,
        });

        self.nodes[node_index].left_first = left_child_index as u32;
        self.nodes[node_index].count = 0;

        self.subdivide(left_child_index, edit_bounds);
        self.subdivide(right_child_index, edit_bounds);
    }

    pub fn query_point(&self, point: Vec3, result: &mut Vec<usize>) {
        result.clear();

        if !self.nodes.is_empty() {
            self.query_point_recursive(0, point, result);
        }

        for &(idx, ref bounds) in &self.pending_adds {
            if bounds.contains_point(point) && !result.contains(&idx) {
                result.push(idx);
            }
        }

        result.retain(|idx| !self.removed_indices.contains(idx));
    }

    fn query_point_recursive(&self, node_index: usize, point: Vec3, result: &mut Vec<usize>) {
        let node = &self.nodes[node_index];

        if !node.bounds.contains_point(point) {
            return;
        }

        if node.is_leaf() {
            let first = node.left_first as usize;
            let count = node.count as usize;
            for idx in first..first + count {
                result.push(self.edit_indices[idx]);
            }
        } else {
            let left_child = node.left_first as usize;
            let right_child = left_child + 1;
            self.query_point_recursive(left_child, point, result);
            self.query_point_recursive(right_child, point, result);
        }
    }

    pub fn query_bounds(&self, query: &Aabb, result: &mut Vec<usize>) {
        result.clear();

        if !self.nodes.is_empty() {
            self.query_bounds_recursive(0, query, result);
        }

        for &(idx, ref bounds) in &self.pending_adds {
            if bounds.intersects(query) && !result.contains(&idx) {
                result.push(idx);
            }
        }

        result.retain(|idx| !self.removed_indices.contains(idx));
    }

    fn query_bounds_recursive(&self, node_index: usize, query: &Aabb, result: &mut Vec<usize>) {
        let node = &self.nodes[node_index];

        if !node.bounds.intersects(query) {
            return;
        }

        if node.is_leaf() {
            let first = node.left_first as usize;
            let count = node.count as usize;
            for idx in first..first + count {
                result.push(self.edit_indices[idx]);
            }
        } else {
            let left_child = node.left_first as usize;
            let right_child = left_child + 1;
            self.query_bounds_recursive(left_child, query, result);
            self.query_bounds_recursive(right_child, query, result);
        }
    }

    pub fn is_empty(&self) -> bool {
        self.nodes.is_empty() && self.pending_adds.is_empty()
    }

    pub fn pending_count(&self) -> usize {
        self.pending_adds.len()
    }

    pub fn removed_count(&self) -> usize {
        self.removed_indices.len()
    }
}

fn compute_bounds(edit_bounds: &[Aabb], indices: &[usize], start: usize, count: usize) -> Aabb {
    let mut bounds = Aabb::default();
    for &edit_idx in indices.iter().skip(start).take(count) {
        bounds = bounds.union(&edit_bounds[edit_idx]);
    }
    bounds
}
