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;

struct TreeNode {
    entity: Entity,
    name: String,
    depth: u32,
    has_children: bool,
    is_gizmo: bool,
}

pub struct TreeCache {
    nodes: Vec<TreeNode>,
    expanded: HashSet<Entity>,
    visible_rows: Vec<usize>,
    last_entity_count: usize,
    dirty: bool,
}

impl Default for TreeCache {
    fn default() -> Self {
        Self {
            nodes: Vec::new(),
            expanded: HashSet::new(),
            visible_rows: Vec::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) {
        self.nodes.clear();

        let mut roots: Vec<Entity> = world
            .query_entities(
                nightshade::ecs::world::LOCAL_TRANSFORM | nightshade::ecs::world::GLOBAL_TRANSFORM,
            )
            .filter(|entity| !world.entity_has_parent(*entity))
            .collect();
        roots.sort_by_key(|e| e.id);

        for root in roots {
            self.build_subtree(world, root, 0);
        }

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

    fn build_subtree(&mut self, world: &World, entity: Entity, depth: u32) {
        let name = if let Some(Name(name)) = world.get_name(entity) {
            name.clone()
        } else {
            format!("Entity {}", entity.id)
        };

        let is_gizmo = name.ends_with("Gizmo");

        let children = self.get_sorted_children(world, entity);
        let has_children = !children.is_empty();

        self.nodes.push(TreeNode {
            entity,
            name,
            depth,
            has_children,
            is_gizmo,
        });

        for child in children {
            self.build_subtree(world, child, depth + 1);
        }
    }

    fn get_sorted_children(&self, 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 rebuild_visible_rows(&mut self) {
        self.visible_rows.clear();

        let mut skip_until_depth: Option<u32> = None;

        for (index, node) in self.nodes.iter().enumerate() {
            if node.is_gizmo {
                continue;
            }

            if let Some(skip_depth) = skip_until_depth {
                if node.depth > skip_depth {
                    continue;
                } else {
                    skip_until_depth = None;
                }
            }

            self.visible_rows.push(index);

            if node.has_children && !self.expanded.contains(&node.entity) {
                skip_until_depth = Some(node.depth);
            }
        }
    }

    fn toggle_expanded(&mut self, entity: Entity) {
        if self.expanded.contains(&entity) {
            self.expanded.remove(&entity);
        } else {
            self.expanded.insert(entity);
        }
        self.rebuild_visible_rows();
    }

    fn set_expanded(&mut self, entity: Entity, expanded: bool) {
        if expanded {
            self.expanded.insert(entity);
        } else {
            self.expanded.remove(&entity);
        }
    }

    fn expand_all(&mut self) {
        for node in &self.nodes {
            if node.has_children && !node.is_gizmo {
                self.expanded.insert(node.entity);
            }
        }
        self.rebuild_visible_rows();
    }

    fn collapse_all(&mut self) {
        self.expanded.clear();
        self.rebuild_visible_rows();
    }

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

pub struct WorldTreeUi;

const ROW_HEIGHT: f32 = 20.0;
const INDENT_SIZE: f32 = 16.0;

impl WorldTreeUi {
    pub fn ui_with_context(
        world: &mut World,
        selected_entity: &mut Option<Entity>,
        undo_history: &mut UndoHistory,
        tree_cache: &mut TreeCache,
        ui: &mut egui::Ui,
    ) {
        if tree_cache.needs_rebuild(world) {
            tree_cache.rebuild(world);
        }

        ui.horizontal(|ui| {
            if ui.button("New Entity").clicked() {
                let entities = EntityBuilder::new()
                    .with_local_transform(LocalTransform::default())
                    .with_global_transform(GlobalTransform::default())
                    .spawn(world, 1);
                let new_entity = entities[0];
                world.set_name(new_entity, 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,
                    },
                    "Create entity".to_string(),
                );
                tree_cache.mark_dirty();
            }
            if ui.button("Expand All").clicked() {
                tree_cache.expand_all();
            }
            if ui.button("Collapse All").clicked() {
                tree_cache.collapse_all();
            }
        });

        let total_rows = tree_cache.visible_rows.len();

        egui::ScrollArea::vertical()
            .auto_shrink([false, false])
            .show_rows(ui, ROW_HEIGHT, total_rows, |ui, row_range| {
                for row_index in row_range {
                    if row_index >= tree_cache.visible_rows.len() {
                        break;
                    }

                    let node_index = tree_cache.visible_rows[row_index];
                    let node = &tree_cache.nodes[node_index];

                    let entity = node.entity;
                    let depth = node.depth;
                    let has_children = node.has_children;
                    let name = node.name.clone();
                    let is_expanded = tree_cache.expanded.contains(&entity);

                    let indent = depth as f32 * INDENT_SIZE;

                    ui.horizontal(|ui| {
                        ui.add_space(indent);

                        if has_children {
                            let icon = if is_expanded { "v" } else { ">" };
                            if ui.small_button(icon).clicked() {
                                tree_cache.toggle_expanded(entity);
                            }
                        } else {
                            ui.add_space(18.0);
                        }

                        let selected = *selected_entity == Some(entity);
                        let response = ui.selectable_label(selected, &name);

                        if response.clicked() {
                            *selected_entity = Some(entity);
                        }

                        Self::show_context_menu(
                            world,
                            selected_entity,
                            undo_history,
                            tree_cache,
                            entity,
                            &response,
                        );
                    });
                }
            });
    }

    fn show_context_menu(
        world: &mut World,
        selected_entity: &mut Option<Entity>,
        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, selected_entity, 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);
                *selected_entity = None;
                tree_cache.mark_dirty();
                ui.close();
            }
        });
    }

    fn add_entity_internal(
        world: &mut World,
        selected_entity: &mut Option<Entity>,
        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.set_expanded(parent_entity, true);
        tree_cache.mark_dirty();
        *selected_entity = Some(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
        );
    }
}
