mod prefab;
mod rendering;

pub use rendering::WorldTreeUi;

use nightshade::prelude::*;
use std::collections::HashSet;

pub struct TreeCache {
    roots: Vec<Entity>,
    editor_entities: HashSet<Entity>,
    sorted_children: std::collections::HashMap<Entity, Vec<Entity>>,
    dirty: bool,
    last_gizmo_root: Option<Entity>,
    last_selection_count: usize,
}

impl Default for TreeCache {
    fn default() -> Self {
        Self {
            roots: Vec::new(),
            editor_entities: HashSet::new(),
            sorted_children: std::collections::HashMap::new(),
            dirty: true,
            last_gizmo_root: None,
            last_selection_count: 0,
        }
    }
}

impl TreeCache {
    pub fn needs_rebuild(&self, gizmo_root: Option<Entity>, selection_count: usize) -> bool {
        self.dirty
            || self.last_gizmo_root != gizmo_root
            || self.last_selection_count != selection_count
    }

    pub fn rebuild(
        &mut self,
        world: &World,
        active_camera: Option<Entity>,
        gizmo_root: Option<Entity>,
    ) {
        self.last_gizmo_root = gizmo_root;
        self.editor_entities.clear();

        if let Some(root) = gizmo_root {
            self.editor_entities.insert(root);
            for descendant in nightshade::ecs::transform::queries::query_descendants(world, root) {
                self.editor_entities.insert(descendant);
            }
        }

        if let Some(camera) = active_camera {
            self.editor_entities.insert(camera);
            for descendant in nightshade::ecs::transform::queries::query_descendants(world, camera)
            {
                self.editor_entities.insert(descendant);
            }
        }

        self.roots = world
            .core
            .query_entities(
                nightshade::ecs::world::LOCAL_TRANSFORM | nightshade::ecs::world::GLOBAL_TRANSFORM,
            )
            .filter(|entity| {
                !world.core.entity_has_parent(*entity) && !self.editor_entities.contains(entity)
            })
            .collect();
        self.roots.sort_by_key(|entity| entity.id);

        self.sorted_children.clear();
        for (parent, cached_children) in &world.resources.children_cache {
            let mut filtered: Vec<Entity> = cached_children
                .iter()
                .copied()
                .filter(|entity| !self.editor_entities.contains(entity))
                .collect();
            if !filtered.is_empty() {
                filtered.sort_by_key(|entity| entity.id);
                self.sorted_children.insert(*parent, filtered);
            }
        }

        self.dirty = false;
    }

    pub fn update_selection_count(&mut self, count: usize) {
        self.last_selection_count = count;
    }

    pub fn mark_dirty(&mut self) {
        self.dirty = true;
    }
}
