use super::selection::{
    Selection, selection_add, selection_clear, selection_contains, selection_set_single,
    selection_toggle,
};
use super::undo::{UndoHistory, UndoableOperation};
use nightshade::ecs::prefab::resources::mesh_cache_insert;
use nightshade::ecs::prefab::{import_gltf_from_path, spawn_prefab_node};
use nightshade::ecs::world::PREFAB_SOURCE;
use nightshade::prelude::*;
use std::collections::HashSet;
use std::path::Path;

pub struct TreeCache {
    roots: Vec<Entity>,
    editor_entities: HashSet<Entity>,
    last_entity_count: usize,
    dirty: bool,
}

impl Default for TreeCache {
    fn default() -> Self {
        Self {
            roots: Vec::new(),
            editor_entities: HashSet::new(),
            last_entity_count: 0,
            dirty: true,
        }
    }
}

impl TreeCache {
    fn needs_rebuild(&self, world: &World) -> bool {
        if self.dirty {
            return true;
        }
        let current_count = world
            .query_entities(
                nightshade::ecs::world::LOCAL_TRANSFORM | nightshade::ecs::world::GLOBAL_TRANSFORM,
            )
            .count();
        current_count != self.last_entity_count
    }

    fn rebuild(
        &mut self,
        world: &World,
        active_camera: Option<Entity>,
        gizmo_root: Option<Entity>,
    ) {
        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
            .query_entities(
                nightshade::ecs::world::LOCAL_TRANSFORM | nightshade::ecs::world::GLOBAL_TRANSFORM,
            )
            .filter(|entity| {
                !world.entity_has_parent(*entity) && !self.editor_entities.contains(entity)
            })
            .collect();
        self.roots.sort_by_key(|e| e.id);

        self.last_entity_count = world
            .query_entities(
                nightshade::ecs::world::LOCAL_TRANSFORM | nightshade::ecs::world::GLOBAL_TRANSFORM,
            )
            .count();
        self.dirty = false;
    }

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

fn get_sorted_children(world: &World, entity: Entity) -> Vec<Entity> {
    let mut children: Vec<Entity> = world
        .resources
        .children_cache
        .get(&entity)
        .cloned()
        .unwrap_or_default();
    children.sort_by_key(|e| e.id);
    children
}

fn get_entity_name(world: &World, entity: Entity) -> String {
    if let Some(Name(name)) = world.get_name(entity) {
        name.clone()
    } else {
        format!("Entity {}", entity.id)
    }
}

pub struct WorldTreeUi;

impl WorldTreeUi {
    pub fn ui_with_context(
        world: &mut World,
        selection: &mut Selection,
        undo_history: &mut UndoHistory,
        tree_cache: &mut TreeCache,
        gizmo_root: Option<Entity>,
        ui: &mut egui::Ui,
    ) -> bool {
        if tree_cache.needs_rebuild(world) {
            tree_cache.rebuild(world, world.resources.active_camera, gizmo_root);
        }

        let mut open_add_node_modal = false;

        ui.horizontal(|ui| {
            if ui.button("+ Add Node").clicked() {
                open_add_node_modal = true;
            }
        });

        ui.separator();

        egui::ScrollArea::vertical()
            .auto_shrink([false, false])
            .show(ui, |ui| {
                let roots = tree_cache.roots.clone();
                let editor_entities = tree_cache.editor_entities.clone();
                for root in roots {
                    Self::render_entity_recursive(
                        world,
                        selection,
                        undo_history,
                        tree_cache,
                        &editor_entities,
                        root,
                        ui,
                    );
                }
            });

        open_add_node_modal
    }

    fn render_entity_recursive(
        world: &mut World,
        selection: &mut Selection,
        undo_history: &mut UndoHistory,
        tree_cache: &mut TreeCache,
        editor_entities: &HashSet<Entity>,
        entity: Entity,
        ui: &mut egui::Ui,
    ) {
        if editor_entities.contains(&entity) {
            return;
        }

        let name = get_entity_name(world, entity);
        let children: Vec<Entity> = get_sorted_children(world, entity)
            .into_iter()
            .filter(|e| !editor_entities.contains(e))
            .collect();
        let has_children = !children.is_empty();
        let selected = selection_contains(selection, entity);

        if has_children {
            let id = ui.make_persistent_id(entity.id);
            egui::collapsing_header::CollapsingState::load_with_default_open(ui.ctx(), id, false)
                .show_header(ui, |ui| {
                    let response = ui.selectable_label(selected, &name);
                    Self::handle_click(selection, entity, &response, ui);
                    Self::show_context_menu(
                        world,
                        selection,
                        undo_history,
                        tree_cache,
                        entity,
                        &response,
                    );
                })
                .body(|ui| {
                    for child in children {
                        Self::render_entity_recursive(
                            world,
                            selection,
                            undo_history,
                            tree_cache,
                            editor_entities,
                            child,
                            ui,
                        );
                    }
                });
        } else {
            ui.horizontal(|ui| {
                ui.add_space(19.0);
                let response = ui.selectable_label(selected, &name);
                Self::handle_click(selection, entity, &response, ui);
                Self::show_context_menu(
                    world,
                    selection,
                    undo_history,
                    tree_cache,
                    entity,
                    &response,
                );
            });
        }
    }

    fn handle_click(
        selection: &mut Selection,
        entity: Entity,
        response: &egui::Response,
        ui: &egui::Ui,
    ) {
        if response.clicked() {
            let modifiers = ui.input(|i| i.modifiers);
            if modifiers.ctrl {
                selection_toggle(selection, entity);
            } else if modifiers.shift {
                selection_add(selection, entity);
            } else {
                selection_set_single(selection, entity);
            }
        }
    }

    fn show_context_menu(
        world: &mut World,
        selection: &mut Selection,
        undo_history: &mut UndoHistory,
        tree_cache: &mut TreeCache,
        entity: Entity,
        response: &egui::Response,
    ) {
        let has_prefab_source = world.get_prefab_source(entity).is_some();

        response.context_menu(|ui| {
            if ui.button("Add Entity").clicked() {
                Self::add_entity_internal(world, selection, undo_history, tree_cache, entity);
                ui.close();
            }

            if has_prefab_source && ui.button("Instantiate Prefab").clicked() {
                Self::instantiate_prefab(world, tree_cache, entity);
                ui.close();
            }

            if ui.button("Remove").clicked() {
                let hierarchy = Box::new(UndoHistory::capture_hierarchy(world, entity));
                let entity_name = world
                    .get_name(entity)
                    .map(|n| n.0.clone())
                    .unwrap_or_else(|| format!("Entity {}", entity.id));
                undo_history.push(
                    UndoableOperation::EntityDeleted {
                        hierarchy,
                        deleted_entity: entity,
                    },
                    format!("Delete {}", entity_name),
                );
                despawn_recursive_immediate(world, entity);
                selection_clear(selection);
                tree_cache.mark_dirty();
                ui.close();
            }
        });
    }

    fn add_entity_internal(
        world: &mut World,
        selection: &mut Selection,
        undo_history: &mut UndoHistory,
        tree_cache: &mut TreeCache,
        parent_entity: Entity,
    ) {
        let entities = EntityBuilder::new()
            .with_name(Name(String::new()))
            .with_local_transform(LocalTransform::default())
            .with_global_transform(GlobalTransform::default())
            .with_parent(Parent(Some(parent_entity)))
            .spawn(world, 1);
        let new_entity = entities[0];

        if let Some(name) = world.get_name_mut(new_entity) {
            *name = Name(format!("Entity {}", new_entity.id));
        }

        let hierarchy = Box::new(UndoHistory::capture_hierarchy(world, new_entity));
        undo_history.push(
            UndoableOperation::EntityCreated {
                hierarchy,
                current_entity: new_entity,
            },
            format!("Create Entity {}", new_entity.id),
        );

        tree_cache.mark_dirty();
        selection_set_single(selection, new_entity);
    }

    fn instantiate_prefab(world: &mut World, tree_cache: &mut TreeCache, entity: Entity) {
        let prefab_source = match world.get_prefab_source(entity) {
            Some(source) => source.clone(),
            None => return,
        };

        let Some(source_path) = &prefab_source.source_path else {
            tracing::warn!(
                "Cannot instantiate prefab '{}': no source path",
                prefab_source.prefab_name
            );
            return;
        };

        let gltf_path = Path::new(source_path);
        let result = match import_gltf_from_path(gltf_path) {
            Ok(result) => result,
            Err(error) => {
                tracing::error!("Failed to load prefab from '{}': {}", source_path, error);
                return;
            }
        };

        for (name, (rgba_data, width, height)) in result.textures {
            world.queue_command(WorldCommand::LoadTexture {
                name,
                rgba_data,
                width,
                height,
            });
        }

        for (name, mesh) in result.meshes {
            mesh_cache_insert(&mut world.resources.mesh_cache, name, mesh);
        }

        let mut prefabs = result.prefabs;
        let prefab_index = prefabs
            .iter()
            .position(|p| p.name == prefab_source.prefab_name);
        let prefab = if let Some(index) = prefab_index {
            prefabs.remove(index)
        } else if !prefabs.is_empty() {
            tracing::warn!(
                "Prefab '{}' not found in '{}', using first available",
                prefab_source.prefab_name,
                source_path
            );
            prefabs.remove(0)
        } else {
            tracing::error!("No prefabs found in '{}'", source_path);
            return;
        };

        for root_node in &prefab.root_nodes {
            spawn_prefab_node(world, root_node, nalgebra_glm::Vec3::zeros(), Some(entity));
        }

        world.remove_components(entity, PREFAB_SOURCE);
        tree_cache.mark_dirty();
        tracing::info!(
            "Instantiated prefab '{}' on entity {:?}",
            prefab_source.prefab_name,
            entity
        );
    }
}
