nightshade 0.13.3

use crate::ecs::ui::types::Rect;
use crate::ecs::world::World;

const DOUBLE_CLICK_TIME: f64 = 0.3;
const DOUBLE_CLICK_DISTANCE: f32 = 4.0;
const PICKING_CELL_SIZE: f32 = 64.0;

struct HitTestResult {
    entity: Option<freecs::Entity>,
    disabled: bool,
}

struct PickingEntry {
    entity: freecs::Entity,
    depth: f32,
    computed_rect: Rect,
    clip_rect: Option<Rect>,
    disabled: bool,
}

pub struct PickingGrid {
    cells: Vec<Vec<PickingEntry>>,
    cell_size: f32,
    columns: usize,
    rows: usize,
}

impl PickingGrid {
    fn new(viewport_width: f32, viewport_height: f32, cell_size: f32) -> Self {
        let columns = ((viewport_width / cell_size).ceil() as usize).max(1);
        let rows = ((viewport_height / cell_size).ceil() as usize).max(1);
        Self {
            cells: (0..columns * rows).map(|_| Vec::new()).collect(),
            cell_size,
            columns,
            rows,
        }
    }

    fn insert(&mut self, entry: PickingEntry) {
        let effective_rect = if let Some(clip) = entry.clip_rect {
            match entry.computed_rect.intersect(&clip) {
                Some(clipped) => clipped,
                None => return,
            }
        } else {
            entry.computed_rect
        };

        let epsilon = 0.01;
        let col_start =
            (((effective_rect.min.x - epsilon) / self.cell_size).floor() as isize).max(0) as usize;
        let col_end =
            (((effective_rect.max.x + epsilon) / self.cell_size).ceil() as usize).min(self.columns);
        let row_start =
            (((effective_rect.min.y - epsilon) / self.cell_size).floor() as isize).max(0) as usize;
        let row_end =
            (((effective_rect.max.y + epsilon) / self.cell_size).ceil() as usize).min(self.rows);

        for row in row_start..row_end {
            for col in col_start..col_end {
                let cell_index = row * self.columns + col;
                self.cells[cell_index].push(PickingEntry {
                    entity: entry.entity,
                    depth: entry.depth,
                    computed_rect: entry.computed_rect,
                    clip_rect: entry.clip_rect,
                    disabled: entry.disabled,
                });
            }
        }
    }

    fn sort_cells(&mut self) {
        for cell in &mut self.cells {
            cell.sort_by(|a, b| {
                b.depth
                    .partial_cmp(&a.depth)
                    .unwrap_or(std::cmp::Ordering::Equal)
            });
        }
    }

    fn query(&self, point: nalgebra_glm::Vec2) -> Option<&[PickingEntry]> {
        let col = (point.x / self.cell_size).floor() as isize;
        let row = (point.y / self.cell_size).floor() as isize;
        if col < 0 || row < 0 || col >= self.columns as isize || row >= self.rows as isize {
            return None;
        }
        let cell_index = row as usize * self.columns + col as usize;
        Some(&self.cells[cell_index])
    }
}

fn is_descendant_of(world: &World, entity: freecs::Entity, ancestor: freecs::Entity) -> bool {
    let mut current = entity;
    loop {
        if current == ancestor {
            return true;
        }
        if let Some(parent) = world.core.get_parent(current)
            && let Some(parent_entity) = parent.0
        {
            current = parent_entity;
        } else {
            return false;
        }
    }
}

fn hit_test_grid(grid: &PickingGrid, mouse_position: nalgebra_glm::Vec2) -> HitTestResult {
    let Some(entries) = grid.query(mouse_position) else {
        return HitTestResult {
            entity: None,
            disabled: false,
        };
    };

    for entry in entries {
        if let Some(clip) = entry.clip_rect
            && !clip.contains(mouse_position)
        {
            continue;
        }

        if entry.computed_rect.contains(mouse_position) {
            return HitTestResult {
                entity: Some(entry.entity),
                disabled: entry.disabled,
            };
        }
    }

    HitTestResult {
        entity: None,
        disabled: false,
    }
}

pub fn build_picking_grid(world: &World) -> PickingGrid {
    let viewport = world
        .resources
        .window
        .cached_viewport_size
        .map(|(width, height)| (width as f32, height as f32))
        .unwrap_or((800.0, 600.0));

    let mut grid = PickingGrid::new(viewport.0, viewport.1, PICKING_CELL_SIZE);

    for &(entity, depth) in &world.resources.retained_ui.z_sorted_nodes {
        let Some(node) = world.ui.get_ui_layout_node(entity) else {
            continue;
        };
        if !node.visible || !node.pointer_events {
            continue;
        }
        if world.ui.get_ui_node_interaction(entity).is_none() {
            continue;
        }

        let disabled = world
            .ui
            .get_ui_node_interaction(entity)
            .map(|interaction| interaction.disabled)
            .unwrap_or(false);

        grid.insert(PickingEntry {
            entity,
            depth,
            computed_rect: node.computed_rect,
            clip_rect: node.computed_clip_rect,
            disabled,
        });
    }

    grid.sort_cells();
    grid
}

fn hit_test_sorted_nodes(
    world: &World,
    sorted_nodes: &[(freecs::Entity, f32)],
    mouse_position: nalgebra_glm::Vec2,
) -> HitTestResult {
    for &(entity, _) in sorted_nodes.iter().rev() {
        let (has_pointer_events, visible, computed_rect, computed_clip_rect) = {
            if let Some(node) = world.ui.get_ui_layout_node(entity) {
                (
                    node.pointer_events,
                    node.visible,
                    node.computed_rect,
                    node.computed_clip_rect,
                )
            } else {
                continue;
            }
        };

        if !visible || !has_pointer_events {
            continue;
        }

        if world.ui.get_ui_node_interaction(entity).is_none() {
            continue;
        }

        if let Some(clip) = computed_clip_rect
            && !clip.contains(mouse_position)
        {
            continue;
        }

        if computed_rect.contains(mouse_position) {
            let disabled = world
                .ui
                .get_ui_node_interaction(entity)
                .map(|interaction| interaction.disabled)
                .unwrap_or(false);
            return HitTestResult {
                entity: Some(entity),
                disabled,
            };
        }
    }
    HitTestResult {
        entity: None,
        disabled: false,
    }
}

pub fn ui_layout_picking_system(world: &mut World) {
    if !world.resources.retained_ui.enabled {
        return;
    }

    let mouse_position = world.resources.input.mouse.position;
    let mouse_state = world.resources.input.mouse.state;
    let mouse_just_pressed =
        mouse_state.contains(crate::ecs::input::resources::MouseState::LEFT_JUST_PRESSED);
    let mouse_just_released =
        mouse_state.contains(crate::ecs::input::resources::MouseState::LEFT_JUST_RELEASED);
    let mouse_down = mouse_state.contains(crate::ecs::input::resources::MouseState::LEFT_CLICKED);

    let previous_active = world.resources.retained_ui.active_entity;
    let current_time = world.resources.retained_ui.current_time;

    let sorted_nodes = std::mem::take(&mut world.resources.retained_ui.z_sorted_nodes);

    let interaction_entities: Vec<freecs::Entity> = world
        .ui
        .query_entities(crate::ecs::world::UI_NODE_INTERACTION)
        .collect();

    for entity in &interaction_entities {
        if let Some(interaction) = world.ui.get_ui_node_interaction_mut(*entity) {
            interaction.hovered = false;
            interaction.pressed = false;
            interaction.clicked = false;
            interaction.focused = false;
            interaction.dragging = false;
            interaction.double_clicked = false;
            interaction.right_clicked = false;
        }
    }

    let hit = if let Some(ref grid) = world.resources.retained_ui.picking_grid {
        hit_test_grid(grid, mouse_position)
    } else {
        hit_test_sorted_nodes(world, &sorted_nodes, mouse_position)
    };

    if let Some(entity) = hit.entity
        && !hit.disabled
        && let Some(interaction) = world.ui.get_ui_node_interaction_mut(entity)
    {
        interaction.hovered = true;
    }

    if mouse_just_pressed {
        world.resources.retained_ui.focus_ring_visible = false;

        if let Some(entity) = hit.entity
            && !hit.disabled
        {
            world.resources.retained_ui.active_entity = Some(entity);
            world.resources.retained_ui.focused_entity = Some(entity);

            if let Some(interaction) = world.ui.get_ui_node_interaction_mut(entity) {
                interaction.pressed = true;
                interaction.drag_start = Some(mouse_position);
            }
        } else if hit.entity.is_none() {
            world.resources.retained_ui.focused_entity = None;
        }
    }

    if mouse_down
        && let Some(active) = previous_active
        && let Some(interaction) = world.ui.get_ui_node_interaction_mut(active)
    {
        interaction.pressed = true;
        if interaction.drag_start.is_some() {
            interaction.dragging = true;
        }
    }

    if mouse_just_released {
        if let Some(active) = previous_active {
            if hit.entity == Some(active) {
                if let Some(interaction) = world.ui.get_ui_node_interaction_mut(active) {
                    interaction.clicked = true;
                }

                let is_double = if let Some((last_entity, last_time, last_pos)) =
                    world.resources.retained_ui.last_click
                {
                    last_entity == active
                        && (current_time - last_time) < DOUBLE_CLICK_TIME
                        && (mouse_position - last_pos).magnitude() < DOUBLE_CLICK_DISTANCE
                } else {
                    false
                };

                if is_double {
                    if let Some(interaction) = world.ui.get_ui_node_interaction_mut(active) {
                        interaction.double_clicked = true;
                    }
                    world.resources.retained_ui.last_click = None;
                } else {
                    world.resources.retained_ui.last_click =
                        Some((active, current_time, mouse_position));
                }
            }
            if let Some(interaction) = world.ui.get_ui_node_interaction_mut(active) {
                interaction.drag_start = None;
                interaction.dragging = false;
            }
        }
        world.resources.retained_ui.active_entity = None;
    }

    let right_just_released =
        mouse_state.contains(crate::ecs::input::resources::MouseState::RIGHT_JUST_RELEASED);
    if right_just_released
        && !hit.disabled
        && let Some(entity) = hit.entity
        && let Some(interaction) = world.ui.get_ui_node_interaction_mut(entity)
    {
        interaction.right_clicked = true;
    }

    world.resources.retained_ui.requested_cursor = if hit.disabled {
        None
    } else {
        hit.entity.and_then(|entity| {
            world
                .ui
                .get_ui_node_interaction(entity)
                .and_then(|interaction| interaction.cursor_icon)
        })
    };

    world.resources.retained_ui.hovered_entity = if hit.disabled { None } else { hit.entity };

    let frame_keys = std::mem::take(&mut world.resources.retained_ui.frame_keys);
    let shift_held = world.resources.retained_ui.shift_held;

    for (key, is_pressed) in &frame_keys {
        if !is_pressed {
            continue;
        }
        match key {
            winit::keyboard::KeyCode::Tab => {
                let mut tabbable: Vec<(freecs::Entity, i32)> = Vec::new();
                for entity in &interaction_entities {
                    if let Some(interaction) = world.ui.get_ui_node_interaction(*entity)
                        && let Some(tab_index) = interaction.tab_index
                        && !interaction.disabled
                        && let Some(node) = world.ui.get_ui_layout_node(*entity)
                        && node.visible
                    {
                        tabbable.push((*entity, tab_index));
                    }
                }
                tabbable.sort_by_key(|(_, index)| *index);

                if let Some(modal) = world.resources.retained_ui.active_modal {
                    tabbable.retain(|(entity, _)| is_descendant_of(world, *entity, modal));
                } else if !world.resources.retained_ui.popup_entities.is_empty() {
                    let popups = world.resources.retained_ui.popup_entities.clone();
                    tabbable.retain(|(entity, _)| {
                        popups
                            .iter()
                            .any(|popup| is_descendant_of(world, *entity, *popup))
                    });
                }

                if !tabbable.is_empty() {
                    let focused = world.resources.retained_ui.focused_entity;
                    let current_index = focused.and_then(|focused_entity| {
                        tabbable
                            .iter()
                            .position(|(entity, _)| *entity == focused_entity)
                    });

                    let next_index = if shift_held {
                        match current_index {
                            Some(0) => tabbable.len() - 1,
                            Some(index) => index - 1,
                            None => tabbable.len() - 1,
                        }
                    } else {
                        match current_index {
                            Some(index) => (index + 1) % tabbable.len(),
                            None => 0,
                        }
                    };

                    let focused = tabbable[next_index].0;
                    world.resources.retained_ui.focused_entity = Some(focused);
                    world.resources.retained_ui.focus_ring_visible = true;

                    if let Some(interaction) = world.ui.get_ui_node_interaction(focused) {
                        let role_name =
                            interaction.accessible_role.as_ref().map(|role| match role {
                                crate::ecs::ui::components::AccessibleRole::Button => "button",
                                crate::ecs::ui::components::AccessibleRole::Slider => "slider",
                                crate::ecs::ui::components::AccessibleRole::Checkbox => "checkbox",
                                crate::ecs::ui::components::AccessibleRole::Toggle => "toggle",
                                crate::ecs::ui::components::AccessibleRole::TextInput => {
                                    "text input"
                                }
                                crate::ecs::ui::components::AccessibleRole::TextArea => "text area",
                                crate::ecs::ui::components::AccessibleRole::Dropdown => "dropdown",
                                crate::ecs::ui::components::AccessibleRole::Tab => "tab",
                                crate::ecs::ui::components::AccessibleRole::TabPanel => "tab panel",
                                crate::ecs::ui::components::AccessibleRole::Tree => "tree",
                                crate::ecs::ui::components::AccessibleRole::TreeItem => "tree item",
                                crate::ecs::ui::components::AccessibleRole::Grid => "grid",
                                crate::ecs::ui::components::AccessibleRole::GridCell => "grid cell",
                                crate::ecs::ui::components::AccessibleRole::Dialog => "dialog",
                                crate::ecs::ui::components::AccessibleRole::Alert => "alert",
                                crate::ecs::ui::components::AccessibleRole::ProgressBar => {
                                    "progress bar"
                                }
                                crate::ecs::ui::components::AccessibleRole::Menu => "menu",
                                crate::ecs::ui::components::AccessibleRole::MenuItem => "menu item",
                            });
                        let label = interaction.accessible_label.as_deref();
                        let announcement = match (label, role_name) {
                            (Some(label), Some(role)) => format!("{label}, {role}"),
                            (Some(label), None) => label.to_string(),
                            (None, Some(role)) => role.to_string(),
                            (None, None) => String::new(),
                        };
                        if !announcement.is_empty() {
                            world
                                .resources
                                .retained_ui
                                .announce_queue
                                .push(announcement);
                        }
                    }
                }
            }
            winit::keyboard::KeyCode::Enter | winit::keyboard::KeyCode::Space => {
                if world.resources.retained_ui.focus_ring_visible
                    && let Some(focused) = world.resources.retained_ui.focused_entity
                    && let Some(interaction) = world.ui.get_ui_node_interaction_mut(focused)
                {
                    interaction.clicked = true;
                }
            }
            _ => {}
        }
    }

    if let Some(focused) = world.resources.retained_ui.focused_entity
        && let Some(interaction) = world.ui.get_ui_node_interaction_mut(focused)
    {
        interaction.focused = true;
    }

    let dpi_scale = world.resources.window.cached_scale_factor;
    let drag_threshold = 8.0f32 * dpi_scale;
    if mouse_down
        && let Some(active) = previous_active
        && world.resources.retained_ui.active_drag.is_none()
        && let Some(interaction) = world.ui.get_ui_node_interaction(active)
        && let Some(start) = interaction.drag_start
        && (mouse_position - start).magnitude() > drag_threshold
        && let Some(source) = world.ui.get_ui_drag_source(active)
    {
        let payload = source.payload.clone();
        world.resources.retained_ui.active_drag = Some(crate::ecs::ui::resources::ActiveDrag {
            source: active,
            payload: payload.clone(),
            start_pos: start,
            current_pos: mouse_position,
        });
        world.resources.retained_ui.frame_events.push(
            crate::ecs::ui::resources::UiEvent::DragStarted {
                source: active,
                payload,
            },
        );
    }

    if let Some(ref mut drag) = world.resources.retained_ui.active_drag {
        drag.current_pos = mouse_position;
    }

    if let Some(ref drag) = world.resources.retained_ui.active_drag {
        let drag_source = drag.source;
        let drag_payload = drag.payload.clone();

        let current_over = hit.entity.filter(|&entity| {
            world
                .ui
                .get_ui_drop_target(entity)
                .is_some_and(|target| target.accepted && target.filter.accepts(&drag_payload))
        });
        let previous_over = world.resources.retained_ui.drag_over_entity;

        if current_over != previous_over {
            if let Some(prev) = previous_over {
                world.resources.retained_ui.frame_events.push(
                    crate::ecs::ui::resources::UiEvent::DragLeave {
                        target: prev,
                        source: drag_source,
                    },
                );
            }
            if let Some(curr) = current_over {
                world.resources.retained_ui.frame_events.push(
                    crate::ecs::ui::resources::UiEvent::DragEnter {
                        target: curr,
                        source: drag_source,
                        payload: drag_payload.clone(),
                    },
                );
            }
            world.resources.retained_ui.drag_over_entity = current_over;
        }

        if let Some(over) = current_over {
            world.resources.retained_ui.frame_events.push(
                crate::ecs::ui::resources::UiEvent::DragOver {
                    target: over,
                    source: drag_source,
                    position: mouse_position,
                },
            );
        }
    }

    if mouse_just_released && world.resources.retained_ui.active_drag.is_some() {
        world.resources.retained_ui.drag_over_entity = None;
        let drag = world.resources.retained_ui.active_drag.take().unwrap();
        if let Some(target_entity) = hit.entity
            && world
                .ui
                .get_ui_drop_target(target_entity)
                .is_some_and(|t| t.accepted && t.filter.accepts(&drag.payload))
        {
            world.resources.retained_ui.frame_events.push(
                crate::ecs::ui::resources::UiEvent::DragDropped {
                    source: drag.source,
                    target: target_entity,
                    payload: drag.payload,
                },
            );
        } else {
            world.resources.retained_ui.frame_events.push(
                crate::ecs::ui::resources::UiEvent::DragCancelled {
                    source: drag.source,
                },
            );
        }
    }

    for &(key, pressed) in &frame_keys {
        if pressed
            && key == winit::keyboard::KeyCode::Escape
            && world.resources.retained_ui.active_drag.is_some()
        {
            let drag = world.resources.retained_ui.active_drag.take().unwrap();
            world.resources.retained_ui.frame_events.push(
                crate::ecs::ui::resources::UiEvent::DragCancelled {
                    source: drag.source,
                },
            );
        }
    }

    if let Some(entity) = hit.entity
        && let Some(interaction) = world.ui.get_ui_node_interaction(entity)
        && (interaction.tooltip_text.is_some() || interaction.tooltip_entity.is_some())
    {
        let text = interaction.tooltip_text.clone().unwrap_or_default();
        let tooltip_entity = interaction.tooltip_entity;
        let should_reset = world
            .resources
            .retained_ui
            .tooltip_state
            .as_ref()
            .is_none_or(|state| state.entity != entity);
        if should_reset {
            world.resources.retained_ui.tooltip_state =
                Some(crate::ecs::ui::resources::TooltipActiveState {
                    entity,
                    hover_start: current_time,
                    text,
                    tooltip_entity,
                });
        }
    } else {
        if let Some(tooltip_entity) = world
            .resources
            .retained_ui
            .tooltip_state
            .as_ref()
            .and_then(|prev| prev.tooltip_entity)
            && let Some(node) = world.ui.get_ui_layout_node_mut(tooltip_entity)
        {
            node.visible = false;
        }
        world.resources.retained_ui.tooltip_state = None;
    }

    world.resources.retained_ui.frame_keys = frame_keys;
    world.resources.retained_ui.z_sorted_nodes = sorted_nodes;

    let secondary_window_indices: Vec<usize> = world
        .resources
        .retained_ui
        .secondary_buffers
        .keys()
        .copied()
        .collect();
    for window_index in secondary_window_indices {
        let sec_nodes = world
            .resources
            .retained_ui
            .secondary_buffers
            .get_mut(&window_index)
            .map(|b| std::mem::take(&mut b.z_sorted_nodes))
            .unwrap_or_default();
        if sec_nodes.is_empty() {
            continue;
        }

        let sec_input = world
            .resources
            .secondary_windows
            .states
            .iter()
            .find(|s| s.index == window_index)
            .map(|s| (s.input.mouse_position, s.input.mouse_state));

        let Some((sec_mouse_pos, sec_mouse_state)) = sec_input else {
            if let Some(buffers) = world
                .resources
                .retained_ui
                .secondary_buffers
                .get_mut(&window_index)
            {
                buffers.z_sorted_nodes = sec_nodes;
            }
            continue;
        };

        let sec_just_pressed =
            sec_mouse_state.contains(crate::ecs::input::resources::MouseState::LEFT_JUST_PRESSED);
        let sec_just_released =
            sec_mouse_state.contains(crate::ecs::input::resources::MouseState::LEFT_JUST_RELEASED);
        let sec_mouse_down =
            sec_mouse_state.contains(crate::ecs::input::resources::MouseState::LEFT_CLICKED);
        let sec_right_released =
            sec_mouse_state.contains(crate::ecs::input::resources::MouseState::RIGHT_JUST_RELEASED);

        let sec_previous_active = world
            .resources
            .retained_ui
            .active_entity
            .filter(|a| sec_nodes.iter().any(|(e, _)| *e == *a));

        let sec_hit = hit_test_sorted_nodes(world, &sec_nodes, sec_mouse_pos);

        if let Some(entity) = sec_hit.entity
            && !sec_hit.disabled
            && let Some(interaction) = world.ui.get_ui_node_interaction_mut(entity)
        {
            interaction.hovered = true;
        }

        if sec_just_pressed
            && let Some(entity) = sec_hit.entity
            && !sec_hit.disabled
        {
            world.resources.retained_ui.active_entity = Some(entity);
            world.resources.retained_ui.focused_entity = Some(entity);
            if let Some(interaction) = world.ui.get_ui_node_interaction_mut(entity) {
                interaction.pressed = true;
                interaction.drag_start = Some(sec_mouse_pos);
            }
        }

        if sec_mouse_down
            && let Some(active) = sec_previous_active
            && let Some(interaction) = world.ui.get_ui_node_interaction_mut(active)
        {
            interaction.pressed = true;
            if interaction.drag_start.is_some() {
                interaction.dragging = true;
            }
        }

        if sec_just_released {
            if let Some(active) = world.resources.retained_ui.active_entity
                && sec_hit.entity == Some(active)
                && let Some(interaction) = world.ui.get_ui_node_interaction_mut(active)
            {
                interaction.clicked = true;
            }
            let is_sec_active = world
                .resources
                .retained_ui
                .active_entity
                .is_some_and(|a| sec_nodes.iter().any(|(e, _)| *e == a));
            if is_sec_active {
                if let Some(active) = world.resources.retained_ui.active_entity
                    && let Some(interaction) = world.ui.get_ui_node_interaction_mut(active)
                {
                    interaction.drag_start = None;
                    interaction.dragging = false;
                }
                world.resources.retained_ui.active_entity = None;
            }
        }

        if sec_right_released
            && !sec_hit.disabled
            && let Some(entity) = sec_hit.entity
            && let Some(interaction) = world.ui.get_ui_node_interaction_mut(entity)
        {
            interaction.right_clicked = true;
        }

        if let Some(buffers) = world
            .resources
            .retained_ui
            .secondary_buffers
            .get_mut(&window_index)
        {
            buffers.z_sorted_nodes = sec_nodes;
        }
    }
}