//! Example: spatial-query flock on a 2D top-down map.
//!
//! Built on the same camera + grid skeleton as `overhead_map`, but
//! tile-free so the flock entities aren't competing for visual space with
//! map imagery:
//!
//!   - 60 entities random-walk in 2D across the X-Z world plane.
//!   - One entity is "chosen" (white, larger). Every other entity is
//!     recoloured every frame from its cell relationship to the chosen one:
//!       * red    — same cell as the chosen entity
//!       * yellow — vertex-adjacent (`bevy_a5::query::vertex_neighbors`)
//!       * green  — anywhere else
//!   - A cyan gizmo line connects the chosen entity to its closest
//!     neighbour; a magenta line connects it to the farthest. Each line is
//!     labelled at its midpoint with the live distance in metres.
//!   - The chosen cell and every vertex-adjacent neighbour are outlined on
//!     the grid so you can see the cell membership classification visually
//!     even when no entities sit in those cells.
//!
//! Why `vertex_neighbors` and not `neighbors`? `neighbors` returns cells
//! that share an *edge* with the centre — usually 5 cells for an A5
//! pentagon. Around the icosahedral vertices of the A5 tiling, that can
//! miss the diagonally-touching cells; `vertex_neighbors` picks up every
//! cell sharing at least one vertex (typically 10 cells), which is what
//! "neighbouring" usually means for proximity classification. Both are
//! exposed in `bevy_a5::query`; pick the one that matches your semantics.
//!
//! Controls:
//!   - WASD or click-and-drag to pan
//!   - Mouse wheel or Q/E to zoom
//!   - 1..9 / 0 / `[` `]` change the A5 grid resolution (or use the dropdown)
//!
//! Run with:
//!   cargo run --example flock

use bevy::asset::RenderAssetUsages;
use bevy::camera::ScalingMode;
use bevy::input::mouse::AccumulatedMouseScroll;
use bevy::mesh::{Indices, Mesh, PrimitiveTopology};
use bevy::prelude::*;
use bevy::window::PrimaryWindow;
use bevy_a5::prelude::*;
use bevy_a5::query::{spherical_cap, vertex_neighbors};
use rand::prelude::*;
use std::collections::HashSet;

/// Map centre: Paris.
const CENTER_LAT: f64 = 48.8566;
const CENTER_LON: f64 = 2.3522;
/// Anchors the world coordinate system: 1 reference-tile = `TILE_SIZE`
/// world units across `360°/2^REFERENCE_ZOOM_INT` of longitude. Used by
/// `lonlat_scale` to keep world units interpretable as metres even though
/// no tiles are actually fetched in this example.
const REFERENCE_ZOOM_INT: i32 = 8;
const DEFAULT_RESOLUTION: i32 = 6;
const TILE_SIZE: f32 = 256.0;
const VIEWPORT_HEIGHT_MIN: f32 = 80.0;
const VIEWPORT_HEIGHT_MAX: f32 = 4096.0;

/// Resolution clamp range for the dropdown / keyboard shortcuts.
const MIN_RESOLUTION: i32 = 1;
const MAX_RESOLUTION: i32 = 15;

/// Flock parameters.
const FLOCK_SIZE: usize = 60;
const ENTITY_RADIUS: f32 = 4.0;
const CHOSEN_RADIUS: f32 = 7.0;
const SPAWN_HALF_EXTENT: f32 = 250.0; // world units around (0, 0)
const SURFACE_Y: f32 = 2.0; // entities draw on top of the grid (y=1)
const KICK_STRENGTH: f32 = 60.0; // world units / sqrt(s)
const VELOCITY_DAMPING: f32 = 0.6; // per second
const SPAWN_SEED: u64 = 0xDEAD_BEEF;

fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .add_plugins(BevyA5Plugin)
        .insert_resource(PlanetSettings::earth())
        .insert_resource(GridResolution(DEFAULT_RESOLUTION))
        .insert_resource(GridMeshState::default())
        .insert_resource(DropdownState::default())
        .insert_resource(DistanceTargets::default())
        .add_systems(Startup, setup)
        .add_systems(
            Update,
            (
                pan_camera_keyboard,
                pan_camera_drag,
                zoom_camera,
                resolution_toggle_button,
                resolution_option_button,
                keyboard_resolution_shortcut,
                update_dropdown_visibility,
                update_resolution_dropdown_label,
                rebuild_grid_mesh,
                wander_flock,
                compute_distance_targets,
                colour_flock,
                draw_distance_lines,
                highlight_chosen_cell,
                update_distance_labels,
                update_flock_hud,
            )
                .chain(),
        )
        .run();
}

#[derive(Resource)]
struct GridResolution(i32);

#[derive(Resource, Default)]
struct DropdownState {
    open: bool,
}

#[derive(Resource, Default)]
struct GridMeshState {
    built_resolution: Option<i32>,
    built_center_world: Option<Vec2>,
    built_radius_m: f64,
}

#[derive(Resource, Default)]
struct DragState {
    last_cursor: Option<Vec2>,
    dragging: bool,
}

#[derive(Resource)]
struct FlockMaterials {
    same_cell: Handle<StandardMaterial>,
    neighbour_cell: Handle<StandardMaterial>,
    elsewhere: Handle<StandardMaterial>,
    chosen: Handle<StandardMaterial>,
}

#[derive(Resource, Default)]
struct DistanceTargets {
    /// World-space position (`y = SURFACE_Y`) of the closest non-chosen
    /// entity, plus the great-circle-equivalent distance in metres.
    closest: Option<(Vec3, f64)>,
    farthest: Option<(Vec3, f64)>,
}

#[derive(Component)]
struct MapCamera;
#[derive(Component)]
struct GridMesh;
#[derive(Component)]
struct ResolutionToggleButton;
#[derive(Component)]
struct ResolutionToggleLabel;
#[derive(Component)]
struct ResolutionOptionsPanel;
#[derive(Component)]
struct ResolutionOption(i32);
#[derive(Component)]
struct FlockMember {
    velocity: Vec2,
}
#[derive(Component)]
struct Chosen;
#[derive(Component)]
struct FlockHud;
#[derive(Component)]
struct ClosestLabel;
#[derive(Component)]
struct FarthestLabel;

const TOGGLE_BG_IDLE: Color = Color::linear_rgba(0.10, 0.10, 0.10, 0.85);
const TOGGLE_BG_HOVER: Color = Color::linear_rgba(0.20, 0.20, 0.20, 0.90);
const TOGGLE_BG_PRESSED: Color = Color::linear_rgba(0.25, 0.25, 0.25, 0.95);
const PANEL_BG: Color = Color::linear_rgba(0.08, 0.08, 0.08, 0.95);

fn option_bg(option_res: i32, active_res: i32, hovered: bool) -> Color {
    if option_res == active_res {
        if hovered {
            Color::linear_rgba(0.55, 0.40, 0.10, 0.95)
        } else {
            Color::linear_rgba(0.45, 0.30, 0.05, 0.85)
        }
    } else if hovered {
        Color::linear_rgba(0.30, 0.30, 0.30, 0.95)
    } else {
        Color::linear_rgba(0.0, 0.0, 0.0, 0.0)
    }
}

fn setup(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
) {
    commands.insert_resource(DragState::default());

    // Solid dark background — no map imagery, just the grid + flock.
    commands.insert_resource(ClearColor(Color::linear_rgb(0.04, 0.05, 0.08)));

    // Orthographic top-down camera.
    commands.spawn((
        MapCamera,
        Camera3d::default(),
        Projection::from(OrthographicProjection {
            scaling_mode: ScalingMode::FixedVertical {
                viewport_height: TILE_SIZE * 5.0,
            },
            ..OrthographicProjection::default_3d()
        }),
        Transform::from_translation(Vec3::new(0.0, 500.0, 0.0))
            .looking_at(Vec3::ZERO, Vec3::NEG_Z),
    ));

    commands.spawn((
        DirectionalLight {
            illuminance: 20_000.0,
            shadows_enabled: false,
            ..default()
        },
        Transform::from_rotation(Quat::from_euler(
            EulerRot::XYZ,
            -std::f32::consts::FRAC_PI_2,
            0.0,
            0.0,
        )),
    ));

    // Floating origin anchored at Paris (used by the plugin to give the
    // grid mesh a stable reference; the flock entities don't reference it).
    let origin_cell = GeoCell::from_lon_lat(CENTER_LON, CENTER_LAT, DEFAULT_RESOLUTION)
        .expect("centre coords valid");
    commands.spawn((FloatingOrigin::default(), origin_cell, Transform::default()));

    // Top-left: dropdown menu for picking the A5 resolution.
    commands
        .spawn(Node {
            position_type: PositionType::Absolute,
            top: Val::Px(8.0),
            left: Val::Px(12.0),
            flex_direction: FlexDirection::Column,
            row_gap: Val::Px(2.0),
            ..default()
        })
        .with_children(|parent| {
            parent
                .spawn((
                    ResolutionToggleButton,
                    Button,
                    Node {
                        padding: UiRect::axes(Val::Px(10.0), Val::Px(6.0)),
                        min_width: Val::Px(160.0),
                        ..default()
                    },
                    BackgroundColor(TOGGLE_BG_IDLE),
                ))
                .with_children(|btn| {
                    btn.spawn((
                        ResolutionToggleLabel,
                        Text::new(format!("Resolution: {}  ▾", DEFAULT_RESOLUTION)),
                        TextFont {
                            font_size: 16.0,
                            ..default()
                        },
                        TextColor(Color::WHITE),
                    ));
                });

            parent
                .spawn((
                    ResolutionOptionsPanel,
                    Node {
                        display: Display::None,
                        flex_direction: FlexDirection::Column,
                        min_width: Val::Px(160.0),
                        ..default()
                    },
                    BackgroundColor(PANEL_BG),
                ))
                .with_children(|panel| {
                    for res in MIN_RESOLUTION..=MAX_RESOLUTION {
                        panel
                            .spawn((
                                ResolutionOption(res),
                                Button,
                                Node {
                                    padding: UiRect::axes(Val::Px(10.0), Val::Px(4.0)),
                                    ..default()
                                },
                                BackgroundColor(option_bg(res, DEFAULT_RESOLUTION, false)),
                            ))
                            .with_children(|opt| {
                                opt.spawn((
                                    Text::new(format!("Resolution {}", res)),
                                    TextFont {
                                        font_size: 15.0,
                                        ..default()
                                    },
                                    TextColor(Color::WHITE),
                                ));
                            });
                    }
                });
        });

    // Top-right: flock HUD (chosen cell, colour bucket counts, distances).
    commands.spawn((
        FlockHud,
        Text::new(""),
        TextFont {
            font_size: 14.0,
            ..default()
        },
        TextColor(Color::WHITE),
        Node {
            position_type: PositionType::Absolute,
            top: Val::Px(8.0),
            right: Val::Px(12.0),
            padding: UiRect::axes(Val::Px(10.0), Val::Px(6.0)),
            ..default()
        },
        BackgroundColor(Color::linear_rgba(0.0, 0.0, 0.0, 0.65)),
    ));

    // World-anchored UI labels for the closest / farthest line midpoints.
    commands.spawn((
        ClosestLabel,
        Text::new(""),
        TextFont {
            font_size: 13.0,
            ..default()
        },
        TextColor(Color::linear_rgb(0.30, 1.0, 1.0)),
        Node {
            position_type: PositionType::Absolute,
            padding: UiRect::axes(Val::Px(4.0), Val::Px(2.0)),
            ..default()
        },
        BackgroundColor(Color::linear_rgba(0.0, 0.0, 0.0, 0.65)),
    ));
    commands.spawn((
        FarthestLabel,
        Text::new(""),
        TextFont {
            font_size: 13.0,
            ..default()
        },
        TextColor(Color::linear_rgb(1.0, 0.45, 1.0)),
        Node {
            position_type: PositionType::Absolute,
            padding: UiRect::axes(Val::Px(4.0), Val::Px(2.0)),
            ..default()
        },
        BackgroundColor(Color::linear_rgba(0.0, 0.0, 0.0, 0.65)),
    ));

    // Pre-cache the four flock materials.
    let flock_materials = FlockMaterials {
        same_cell: materials.add(StandardMaterial {
            base_color: Color::linear_rgb(1.00, 0.20, 0.20),
            emissive: LinearRgba::new(1.20, 0.20, 0.20, 1.0),
            unlit: true,
            ..default()
        }),
        neighbour_cell: materials.add(StandardMaterial {
            base_color: Color::linear_rgb(1.00, 0.85, 0.10),
            emissive: LinearRgba::new(1.20, 0.95, 0.15, 1.0),
            unlit: true,
            ..default()
        }),
        elsewhere: materials.add(StandardMaterial {
            base_color: Color::linear_rgb(0.20, 0.95, 0.40),
            emissive: LinearRgba::new(0.20, 1.10, 0.40, 1.0),
            unlit: true,
            ..default()
        }),
        chosen: materials.add(StandardMaterial {
            base_color: Color::linear_rgb(1.00, 1.00, 1.00),
            emissive: LinearRgba::new(2.50, 2.50, 2.50, 1.0),
            unlit: true,
            ..default()
        }),
    };

    // Spawn the flock at random world (X, Z) positions around the origin.
    let entity_mesh = meshes.add(Sphere::new(ENTITY_RADIUS));
    let chosen_mesh = meshes.add(Sphere::new(CHOSEN_RADIUS));
    let mut rng = StdRng::seed_from_u64(SPAWN_SEED);
    for i in 0..FLOCK_SIZE {
        let x = rng.gen_range(-SPAWN_HALF_EXTENT..SPAWN_HALF_EXTENT);
        let z = rng.gen_range(-SPAWN_HALF_EXTENT..SPAWN_HALF_EXTENT);
        let is_chosen = i == 0;
        let mut entity = commands.spawn((
            FlockMember {
                velocity: Vec2::ZERO,
            },
            Transform::from_translation(Vec3::new(x, SURFACE_Y, z)),
            Mesh3d(if is_chosen {
                chosen_mesh.clone()
            } else {
                entity_mesh.clone()
            }),
            MeshMaterial3d(if is_chosen {
                flock_materials.chosen.clone()
            } else {
                flock_materials.elsewhere.clone()
            }),
        ));
        if is_chosen {
            entity.insert(Chosen);
        }
    }
    commands.insert_resource(flock_materials);

    info!("flock: {} entities on overhead map", FLOCK_SIZE);
    info!(
        "Controls: WASD or drag to pan, mouse-wheel or Q/E to zoom, dropdown / 1..0 / [ ] for A5 resolution"
    );
}

// ---------------------------------------------------------------------------
// Camera / map systems (lifted from overhead_map without changes)
// ---------------------------------------------------------------------------

fn pan_camera_keyboard(
    time: Res<Time>,
    keys: Res<ButtonInput<KeyCode>>,
    mut query: Query<(&mut Transform, &Projection), With<MapCamera>>,
) {
    let Ok((mut transform, projection)) = query.single_mut() else {
        return;
    };
    let viewport_h = match projection {
        Projection::Orthographic(o) => match o.scaling_mode {
            ScalingMode::FixedVertical { viewport_height } => viewport_height,
            _ => 512.0,
        },
        _ => 512.0,
    };
    let speed = viewport_h * 0.6 * time.delta_secs();
    let mut delta = Vec3::ZERO;
    if keys.pressed(KeyCode::KeyW) {
        delta.z -= speed;
    }
    if keys.pressed(KeyCode::KeyS) {
        delta.z += speed;
    }
    if keys.pressed(KeyCode::KeyA) {
        delta.x -= speed;
    }
    if keys.pressed(KeyCode::KeyD) {
        delta.x += speed;
    }
    transform.translation += delta;
}

fn pan_camera_drag(
    mouse_buttons: Res<ButtonInput<MouseButton>>,
    windows: Query<&Window, With<PrimaryWindow>>,
    cameras: Query<(&Camera, &Projection, &GlobalTransform), With<MapCamera>>,
    mut camera_transforms: Query<&mut Transform, With<MapCamera>>,
    button_interactions: Query<&Interaction, With<Button>>,
    mut drag: ResMut<DragState>,
) {
    let Ok(window) = windows.single() else {
        return;
    };
    let Ok((camera, projection, _)) = cameras.single() else {
        return;
    };
    let Ok(mut transform) = camera_transforms.single_mut() else {
        return;
    };
    let cursor = window.cursor_position();
    let over_ui = button_interactions
        .iter()
        .any(|i| matches!(i, Interaction::Hovered | Interaction::Pressed));

    if mouse_buttons.just_pressed(MouseButton::Left) && !over_ui {
        drag.dragging = true;
        drag.last_cursor = cursor;
    }
    if mouse_buttons.just_released(MouseButton::Left) {
        drag.dragging = false;
        drag.last_cursor = None;
    }
    if !drag.dragging {
        drag.last_cursor = cursor;
        return;
    }
    let (Some(cur), Some(prev)) = (cursor, drag.last_cursor) else {
        drag.last_cursor = cursor;
        return;
    };
    if cur == prev {
        return;
    }
    let viewport_h = match projection {
        Projection::Orthographic(o) => match o.scaling_mode {
            ScalingMode::FixedVertical { viewport_height } => viewport_height,
            _ => 512.0,
        },
        _ => 512.0,
    };
    let Some(viewport_size) = camera.logical_viewport_size() else {
        return;
    };
    if viewport_size.y <= 0.0 {
        return;
    }
    let world_per_pixel = viewport_h / viewport_size.y;
    let pixel_delta = cur - prev;
    transform.translation.x -= pixel_delta.x * world_per_pixel;
    transform.translation.z -= pixel_delta.y * world_per_pixel;
    drag.last_cursor = cursor;
}

fn zoom_camera(
    time: Res<Time>,
    keys: Res<ButtonInput<KeyCode>>,
    scroll: Res<AccumulatedMouseScroll>,
    mut query: Query<&mut Projection, With<MapCamera>>,
) {
    let Ok(mut projection) = query.single_mut() else {
        return;
    };
    let Projection::Orthographic(ref mut ortho) = *projection else {
        return;
    };
    let ScalingMode::FixedVertical {
        ref mut viewport_height,
    } = ortho.scaling_mode
    else {
        return;
    };
    let mut zoom_steps: f32 = 0.0;
    if scroll.delta.y != 0.0 {
        zoom_steps -= scroll.delta.y;
    }
    if keys.pressed(KeyCode::KeyE) {
        zoom_steps -= 4.0 * time.delta_secs();
    }
    if keys.pressed(KeyCode::KeyQ) {
        zoom_steps += 4.0 * time.delta_secs();
    }
    if zoom_steps == 0.0 {
        return;
    }
    let factor = 1.15_f32.powf(zoom_steps);
    *viewport_height = (*viewport_height * factor).clamp(VIEWPORT_HEIGHT_MIN, VIEWPORT_HEIGHT_MAX);
}

fn rebuild_grid_mesh(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
    mut state: ResMut<GridMeshState>,
    resolution: Res<GridResolution>,
    camera_query: Query<(&Camera, &Projection, &GlobalTransform), With<MapCamera>>,
    grid_q: Query<Entity, With<GridMesh>>,
) {
    let Ok((_, _, camera_gt)) = camera_query.single() else {
        return;
    };
    let center_lon = CENTER_LON;
    let center_lat = CENTER_LAT;
    let scale = lonlat_scale(center_lat);
    let visible_radius_m = visible_radius_meters(&camera_query, scale);
    // Tight build radius (visible diagonal + 15 % margin for caching).
    // The previous 2× factor pushed `k_for_coverage` past `MAX_CELLS` at
    // moderate resolutions, which left empty corners in the viewport.
    let build_radius_m = visible_radius_m * 1.15;
    let cam_world = camera_gt.translation();
    let cam_xz = Vec2::new(cam_world.x, cam_world.z);
    let needs_resolution_rebuild =
        resolution.is_changed() || state.built_resolution != Some(resolution.0);
    let needs_pan_rebuild = match state.built_center_world {
        None => true,
        Some(prev) => {
            let drift_world = (cam_xz - prev).length() as f64;
            let metres_per_world = 110_574.0 / scale.1.max(1e-6);
            let drift_m = drift_world * metres_per_world;
            drift_m > (build_radius_m - visible_radius_m) * 0.5
                || visible_radius_m * 2.0 > state.built_radius_m
        }
    };
    if !needs_resolution_rebuild && !needs_pan_rebuild {
        return;
    }
    let (cam_lon, cam_lat) =
        map_pos_to_lonlat(cam_world.x, cam_world.z, center_lon, center_lat);
    let Some(cam_cell) = GeoCell::from_lon_lat(cam_lon, cam_lat, resolution.0) else {
        return;
    };
    // `spherical_cap` returns every cell whose centre lies within
    // `build_radius_m` of the camera — omnidirectional by construction, so
    // there are no off-axis gaps. `bevy_a5::query::spherical_cap` already
    // uncompacts to the input cell's resolution.
    //
    // No artificial cell-count cap: the count is bounded by
    // `π · build_radius² / cell_area`, which the user already controls via
    // their zoom (viewport height) and chosen resolution. Capping here just
    // re-introduces the silent-truncation failure mode the
    // `spherical_cap → AABB-cull` switch was meant to fix.
    let cells = spherical_cap(&cam_cell, build_radius_m).unwrap_or_default();

    // Viewport rectangle in world coords (15 % margin matching
    // `build_radius_m`) — used to skip cells whose AABB falls entirely
    // outside the visible area. At fine resolutions this drops a sizeable
    // fraction of the disc's cells (the corners we don't actually need).
    let viewport_margin = 1.15_f32;
    let (view_half_w_raw, view_half_h_raw) = view_half_extents(&camera_query);
    let view_half_w = view_half_w_raw * viewport_margin;
    let view_half_h = view_half_h_raw * viewport_margin;
    let viewport_min_x = cam_world.x - view_half_w;
    let viewport_max_x = cam_world.x + view_half_w;
    let viewport_min_z = cam_world.z - view_half_h;
    let viewport_max_z = cam_world.z + view_half_h;

    let mut positions: Vec<[f32; 3]> = Vec::new();
    let mut indices: Vec<u32> = Vec::new();
    let mut offset: u32 = 0;
    for cell in &cells {
        let Some(boundary) = cell.boundary() else {
            continue;
        };
        let verts: Vec<_> = if boundary.len() > 1 && boundary.first() == boundary.last() {
            boundary[..boundary.len() - 1].to_vec()
        } else {
            boundary
        };
        if verts.is_empty() {
            continue;
        }
        // Cull cells whose AABB is entirely outside the viewport rectangle.
        let mut min_x = f32::INFINITY;
        let mut max_x = f32::NEG_INFINITY;
        let mut min_z = f32::INFINITY;
        let mut max_z = f32::NEG_INFINITY;
        let projected: Vec<Vec3> = verts
            .iter()
            .map(|ll| {
                let p = lonlat_to_map_pos(ll.longitude(), ll.latitude(), center_lon, center_lat);
                if p.x < min_x {
                    min_x = p.x;
                }
                if p.x > max_x {
                    max_x = p.x;
                }
                if p.z < min_z {
                    min_z = p.z;
                }
                if p.z > max_z {
                    max_z = p.z;
                }
                p
            })
            .collect();
        if max_x < viewport_min_x
            || min_x > viewport_max_x
            || max_z < viewport_min_z
            || min_z > viewport_max_z
        {
            continue;
        }
        for p in &projected {
            positions.push([p.x, p.y, p.z]);
        }
        let n = projected.len() as u32;
        for i in 0..n {
            indices.push(offset + i);
            indices.push(offset + (i + 1) % n);
        }
        offset += n;
    }
    let mut new_mesh = Mesh::new(
        PrimitiveTopology::LineList,
        RenderAssetUsages::MAIN_WORLD | RenderAssetUsages::RENDER_WORLD,
    );
    new_mesh.insert_attribute(Mesh::ATTRIBUTE_POSITION, positions);
    new_mesh.insert_indices(Indices::U32(indices));
    for e in grid_q.iter() {
        commands.entity(e).despawn();
    }
    let (base_color, emissive) = grid_color_for_resolution(resolution.0);
    commands.spawn((
        GridMesh,
        Mesh3d(meshes.add(new_mesh)),
        MeshMaterial3d(materials.add(StandardMaterial {
            base_color,
            emissive,
            unlit: true,
            ..default()
        })),
        Transform::default(),
    ));
    state.built_resolution = Some(resolution.0);
    state.built_center_world = Some(cam_xz);
    state.built_radius_m = build_radius_m;
}

fn keyboard_resolution_shortcut(
    keys: Res<ButtonInput<KeyCode>>,
    mut resolution: ResMut<GridResolution>,
    mut grid_state: ResMut<GridMeshState>,
    mut origin_q: Query<&mut GeoCell, With<FloatingOrigin>>,
) {
    let direct = [
        (KeyCode::Digit1, 1),
        (KeyCode::Digit2, 2),
        (KeyCode::Digit3, 3),
        (KeyCode::Digit4, 4),
        (KeyCode::Digit5, 5),
        (KeyCode::Digit6, 6),
        (KeyCode::Digit7, 7),
        (KeyCode::Digit8, 8),
        (KeyCode::Digit9, 9),
        (KeyCode::Digit0, 10),
    ];
    let mut target: Option<i32> = None;
    for (key, value) in direct {
        if keys.just_pressed(key) {
            target = Some(value);
            break;
        }
    }
    if target.is_none() {
        if keys.just_pressed(KeyCode::BracketRight) {
            target = Some((resolution.0 + 1).min(MAX_RESOLUTION));
        } else if keys.just_pressed(KeyCode::BracketLeft) {
            target = Some((resolution.0 - 1).max(MIN_RESOLUTION));
        }
    }
    let Some(new) = target else {
        return;
    };
    if resolution.0 == new {
        return;
    }
    resolution.0 = new;
    grid_state.built_resolution = None;
    if let Ok(mut cell) = origin_q.single_mut() {
        if let Some(new_cell) = GeoCell::from_lon_lat(CENTER_LON, CENTER_LAT, new) {
            *cell = new_cell;
        }
    }
}

fn resolution_toggle_button(
    mut q: Query<
        (&Interaction, &mut BackgroundColor),
        (Changed<Interaction>, With<ResolutionToggleButton>),
    >,
    mut state: ResMut<DropdownState>,
) {
    for (interaction, mut bg) in q.iter_mut() {
        match *interaction {
            Interaction::Pressed => {
                state.open = !state.open;
                bg.0 = TOGGLE_BG_PRESSED;
            }
            Interaction::Hovered => bg.0 = TOGGLE_BG_HOVER,
            Interaction::None => bg.0 = TOGGLE_BG_IDLE,
        }
    }
}

fn resolution_option_button(
    interaction_q: Query<(&Interaction, &ResolutionOption), Changed<Interaction>>,
    mut style_q: Query<(&ResolutionOption, &Interaction, &mut BackgroundColor)>,
    mut resolution: ResMut<GridResolution>,
    mut grid_state: ResMut<GridMeshState>,
    mut origin_q: Query<&mut GeoCell, With<FloatingOrigin>>,
    mut dropdown: ResMut<DropdownState>,
) {
    for (interaction, opt) in interaction_q.iter() {
        if matches!(*interaction, Interaction::Pressed) && resolution.0 != opt.0 {
            resolution.0 = opt.0;
            grid_state.built_resolution = None;
            if let Ok(mut cell) = origin_q.single_mut() {
                if let Some(new_cell) = GeoCell::from_lon_lat(CENTER_LON, CENTER_LAT, opt.0) {
                    *cell = new_cell;
                }
            }
            dropdown.open = false;
        }
    }
    for (opt, interaction, mut bg) in style_q.iter_mut() {
        let hovered = matches!(*interaction, Interaction::Hovered | Interaction::Pressed);
        bg.0 = option_bg(opt.0, resolution.0, hovered);
    }
}

fn update_dropdown_visibility(
    state: Res<DropdownState>,
    mut panel_q: Query<&mut Node, With<ResolutionOptionsPanel>>,
) {
    if !state.is_changed() {
        return;
    }
    let Ok(mut node) = panel_q.single_mut() else {
        return;
    };
    node.display = if state.open {
        Display::Flex
    } else {
        Display::None
    };
}

fn update_resolution_dropdown_label(
    resolution: Res<GridResolution>,
    mut label_q: Query<&mut Text, With<ResolutionToggleLabel>>,
) {
    if !resolution.is_changed() {
        return;
    }
    let Ok(mut label) = label_q.single_mut() else {
        return;
    };
    label.0 = format!("Resolution: {}  ▾", resolution.0);
}

// ---------------------------------------------------------------------------
// Flock systems
// ---------------------------------------------------------------------------

/// 2D Brownian motion in the X-Z world plane.
fn wander_flock(time: Res<Time>, mut q: Query<(&mut FlockMember, &mut Transform)>) {
    let dt = time.delta_secs();
    if dt <= 0.0 {
        return;
    }
    let kick_scale = KICK_STRENGTH * dt.sqrt();
    let mut rng = thread_rng();
    for (mut flock, mut transform) in q.iter_mut() {
        let kick = Vec2::new(
            rng.gen_range(-1.0_f32..1.0),
            rng.gen_range(-1.0_f32..1.0),
        ) * kick_scale;
        flock.velocity += kick;
        flock.velocity *= 1.0 - VELOCITY_DAMPING * dt;
        transform.translation.x += flock.velocity.x * dt;
        transform.translation.z += flock.velocity.y * dt;
        transform.translation.y = SURFACE_Y;
    }
}

/// Find the closest and farthest non-chosen entities to the chosen one,
/// using planar Euclidean distance in the map's world coords converted to
/// metres via the lon/lat scale.
fn compute_distance_targets(
    chosen_q: Query<&Transform, With<Chosen>>,
    others_q: Query<&Transform, (With<FlockMember>, Without<Chosen>)>,
    mut targets: ResMut<DistanceTargets>,
) {
    *targets = DistanceTargets::default();
    let Ok(c_t) = chosen_q.single() else {
        return;
    };
    let metres_per_world = 110_574.0 / lonlat_scale(CENTER_LAT).1.max(1e-6);
    let chosen_xz = Vec2::new(c_t.translation.x, c_t.translation.z);
    for t in others_q.iter() {
        let other_xz = Vec2::new(t.translation.x, t.translation.z);
        let dist_world = (chosen_xz - other_xz).length() as f64;
        let dist_m = dist_world * metres_per_world;
        let pos = Vec3::new(t.translation.x, SURFACE_Y, t.translation.z);
        if targets.closest.is_none_or(|(_, d)| dist_m < d) {
            targets.closest = Some((pos, dist_m));
        }
        if targets.farthest.is_none_or(|(_, d)| dist_m > d) {
            targets.farthest = Some((pos, dist_m));
        }
    }
}

/// Recolour every non-chosen flock member based on its cell relationship
/// to the chosen entity's cell. Uses [`vertex_neighbors`] so cells sharing
/// only a vertex with the chosen cell still count as "neighbours" (yellow).
fn colour_flock(
    materials: Res<FlockMaterials>,
    resolution: Res<GridResolution>,
    chosen_q: Query<&Transform, With<Chosen>>,
    mut others_q: Query<
        (&Transform, &mut MeshMaterial3d<StandardMaterial>),
        (With<FlockMember>, Without<Chosen>),
    >,
) {
    let Ok(c_t) = chosen_q.single() else {
        return;
    };
    let (c_lon, c_lat) =
        map_pos_to_lonlat(c_t.translation.x, c_t.translation.z, CENTER_LON, CENTER_LAT);
    let Some(chosen_cell) = GeoCell::from_lon_lat(c_lon, c_lat, resolution.0) else {
        return;
    };
    let neighbour_cells: HashSet<u64> = vertex_neighbors(&chosen_cell)
        .unwrap_or_default()
        .into_iter()
        .map(|c| c.raw())
        .collect();
    let chosen_raw = chosen_cell.raw();

    for (t, mut material) in others_q.iter_mut() {
        let (lon, lat) =
            map_pos_to_lonlat(t.translation.x, t.translation.z, CENTER_LON, CENTER_LAT);
        let cell = GeoCell::from_lon_lat(lon, lat, resolution.0);
        material.0 = match cell {
            Some(c) if c.raw() == chosen_raw => materials.same_cell.clone(),
            Some(c) if neighbour_cells.contains(&c.raw()) => materials.neighbour_cell.clone(),
            _ => materials.elsewhere.clone(),
        };
    }
}

fn draw_distance_lines(
    mut gizmos: Gizmos,
    chosen_q: Query<&Transform, With<Chosen>>,
    targets: Res<DistanceTargets>,
) {
    let Ok(c_t) = chosen_q.single() else {
        return;
    };
    let chosen_pos = Vec3::new(c_t.translation.x, SURFACE_Y, c_t.translation.z);
    if let Some((near_pos, _)) = targets.closest {
        gizmos.line(chosen_pos, near_pos, Color::linear_rgb(0.30, 1.0, 1.0));
    }
    if let Some((far_pos, _)) = targets.farthest {
        gizmos.line(chosen_pos, far_pos, Color::linear_rgb(1.0, 0.45, 1.0));
    }
}

/// Outline the chosen entity's cell in white and every vertex-adjacent
/// neighbour cell in yellow, so cell membership is visible even when no
/// entities are inside those cells.
fn highlight_chosen_cell(
    mut gizmos: Gizmos,
    resolution: Res<GridResolution>,
    chosen_q: Query<&Transform, With<Chosen>>,
) {
    let Ok(c_t) = chosen_q.single() else {
        return;
    };
    let (c_lon, c_lat) =
        map_pos_to_lonlat(c_t.translation.x, c_t.translation.z, CENTER_LON, CENTER_LAT);
    let Some(chosen_cell) = GeoCell::from_lon_lat(c_lon, c_lat, resolution.0) else {
        return;
    };
    let neighbour_cells = vertex_neighbors(&chosen_cell).unwrap_or_default();

    for cell in &neighbour_cells {
        draw_cell_outline(
            &mut gizmos,
            cell,
            CENTER_LON,
            CENTER_LAT,
            Color::linear_rgb(1.0, 0.85, 0.10),
            1.5,
        );
    }
    draw_cell_outline(
        &mut gizmos,
        &chosen_cell,
        CENTER_LON,
        CENTER_LAT,
        Color::WHITE,
        2.0,
    );
}

fn draw_cell_outline(
    gizmos: &mut Gizmos,
    cell: &GeoCell,
    center_lon: f64,
    center_lat: f64,
    color: Color,
    y: f32,
) {
    let Some(boundary) = cell.boundary() else {
        return;
    };
    let verts: Vec<_> = if boundary.len() > 1 && boundary.first() == boundary.last() {
        boundary[..boundary.len() - 1].to_vec()
    } else {
        boundary
    };
    if verts.is_empty() {
        return;
    }
    let projected: Vec<Vec3> = verts
        .iter()
        .map(|ll| {
            let mut p = lonlat_to_map_pos(ll.longitude(), ll.latitude(), center_lon, center_lat);
            p.y = y;
            p
        })
        .collect();
    for i in 0..projected.len() {
        let next = (i + 1) % projected.len();
        gizmos.line(projected[i], projected[next], color);
    }
}

fn update_distance_labels(
    cameras: Query<(&Camera, &GlobalTransform), With<Camera3d>>,
    chosen_q: Query<&Transform, With<Chosen>>,
    targets: Res<DistanceTargets>,
    mut closest_q: Query<
        (&mut Text, &mut Node),
        (With<ClosestLabel>, Without<FarthestLabel>),
    >,
    mut farthest_q: Query<
        (&mut Text, &mut Node),
        (With<FarthestLabel>, Without<ClosestLabel>),
    >,
) {
    let Ok((camera, camera_gt)) = cameras.single() else {
        return;
    };
    let Ok(c_t) = chosen_q.single() else {
        return;
    };
    let chosen_pos = Vec3::new(c_t.translation.x, SURFACE_Y, c_t.translation.z);

    place_label(
        camera,
        camera_gt,
        chosen_pos,
        targets.closest,
        &mut closest_q.single_mut().ok(),
    );
    place_label(
        camera,
        camera_gt,
        chosen_pos,
        targets.farthest,
        &mut farthest_q.single_mut().ok(),
    );
}

fn place_label(
    camera: &Camera,
    camera_gt: &GlobalTransform,
    chosen_pos: Vec3,
    target: Option<(Vec3, f64)>,
    label: &mut Option<(Mut<Text>, Mut<Node>)>,
) {
    let Some((text, node)) = label.as_mut() else {
        return;
    };
    match target {
        Some((other_pos, dist)) => {
            let mid = (chosen_pos + other_pos) * 0.5;
            match camera.world_to_viewport(camera_gt, mid) {
                Ok(vp) => {
                    node.display = Display::Flex;
                    node.left = Val::Px(vp.x);
                    node.top = Val::Px(vp.y);
                    text.0 = format_distance(dist);
                }
                Err(_) => node.display = Display::None,
            }
        }
        None => node.display = Display::None,
    }
}

fn format_distance(metres: f64) -> String {
    if metres >= 1000.0 {
        format!("{:.2} km", metres / 1000.0)
    } else {
        format!("{metres:.1} m")
    }
}

fn update_flock_hud(
    resolution: Res<GridResolution>,
    chosen_q: Query<&Transform, With<Chosen>>,
    others_q: Query<&Transform, (With<FlockMember>, Without<Chosen>)>,
    targets: Res<DistanceTargets>,
    mut text_q: Query<&mut Text, With<FlockHud>>,
) {
    let Ok(c_t) = chosen_q.single() else {
        return;
    };
    let Ok(mut text) = text_q.single_mut() else {
        return;
    };
    let (c_lon, c_lat) =
        map_pos_to_lonlat(c_t.translation.x, c_t.translation.z, CENTER_LON, CENTER_LAT);
    let Some(chosen_cell) = GeoCell::from_lon_lat(c_lon, c_lat, resolution.0) else {
        return;
    };
    let neighbour_cells: HashSet<u64> = vertex_neighbors(&chosen_cell)
        .unwrap_or_default()
        .into_iter()
        .map(|c| c.raw())
        .collect();
    let chosen_raw = chosen_cell.raw();

    let mut same = 0;
    let mut neighbour = 0;
    let mut elsewhere = 0;
    for t in others_q.iter() {
        let (lon, lat) =
            map_pos_to_lonlat(t.translation.x, t.translation.z, CENTER_LON, CENTER_LAT);
        let Some(cell) = GeoCell::from_lon_lat(lon, lat, resolution.0) else {
            continue;
        };
        if cell.raw() == chosen_raw {
            same += 1;
        } else if neighbour_cells.contains(&cell.raw()) {
            neighbour += 1;
        } else {
            elsewhere += 1;
        }
    }

    let closest_str = targets
        .closest
        .map(|(_, d)| format_distance(d))
        .unwrap_or_else(|| "—".into());
    let farthest_str = targets
        .farthest
        .map(|(_, d)| format_distance(d))
        .unwrap_or_else(|| "—".into());

    text.0 = format!(
        "Flock spatial-query demo\n\
         \n\
         Chosen cell:        0x{:016x}\n\
         Chosen lon/lat:     {:>7.3}°, {:>6.3}°\n\
         Resolution:         {}\n\
         \n\
         Same cell (red):       {same:>3}\n\
         Neighbour (yellow):    {neighbour:>3}\n\
         Elsewhere (green):     {elsewhere:>3}\n\
         \n\
         Closest (cyan):     {closest_str}\n\
         Farthest (magenta): {farthest_str}",
        chosen_raw,
        c_lon,
        c_lat,
        resolution.0,
    );
}

// ---------------------------------------------------------------------------
// Map projection helpers
// ---------------------------------------------------------------------------

/// Half the visible viewport extents in world units, `(half_x, half_z)`.
fn view_half_extents(
    camera_query: &Query<(&Camera, &Projection, &GlobalTransform), With<MapCamera>>,
) -> (f32, f32) {
    let Ok((camera, projection, _)) = camera_query.single() else {
        return (256.0, 256.0);
    };
    let viewport_h = match projection {
        Projection::Orthographic(o) => match o.scaling_mode {
            ScalingMode::FixedVertical { viewport_height } => viewport_height,
            _ => 512.0,
        },
        _ => 512.0,
    };
    let aspect = camera
        .logical_viewport_size()
        .map(|s| if s.y > 0.0 { s.x / s.y } else { 1.0 })
        .unwrap_or(1.0);
    (viewport_h * 0.5 * aspect, viewport_h * 0.5)
}

fn visible_radius_meters(
    camera_query: &Query<(&Camera, &Projection, &GlobalTransform), With<MapCamera>>,
    scale: (f64, f64),
) -> f64 {
    let Ok((camera, projection, _)) = camera_query.single() else {
        return 200_000.0;
    };
    let viewport_h = match projection {
        Projection::Orthographic(o) => match o.scaling_mode {
            ScalingMode::FixedVertical { viewport_height } => viewport_height,
            _ => 512.0,
        },
        _ => 512.0,
    };
    let aspect = camera
        .logical_viewport_size()
        .map(|s| if s.y > 0.0 { s.x / s.y } else { 1.0 })
        .unwrap_or(1.0);
    let viewport_w = viewport_h * aspect;
    let half_diag_world =
        ((viewport_w * viewport_w + viewport_h * viewport_h).sqrt() * 0.5) as f64;
    let metres_per_world = 110_574.0 / scale.1.max(1e-6);
    half_diag_world * metres_per_world
}

fn grid_color_for_resolution(resolution: i32) -> (Color, LinearRgba) {
    let span = (MAX_RESOLUTION - MIN_RESOLUTION).max(1) as f32;
    let t = ((resolution - MIN_RESOLUTION).clamp(0, MAX_RESOLUTION - MIN_RESOLUTION)) as f32
        / span;
    let hue = t * 360.0;
    let base = Color::hsl(hue, 0.95, 0.55);
    let emissive_color = Color::hsl(hue, 0.95, 0.65);
    let emissive = LinearRgba::from(emissive_color) * 1.5;
    (base, emissive)
}

fn lonlat_to_map_pos(lon: f64, lat: f64, center_lon: f64, center_lat: f64) -> Vec3 {
    let scale = lonlat_scale(center_lat);
    let dx = (lon - center_lon) * scale.0;
    let dy = (lat - center_lat) * scale.1;
    Vec3::new(dx as f32, 1.0, -(dy) as f32)
}

fn map_pos_to_lonlat(x: f32, z: f32, center_lon: f64, center_lat: f64) -> (f64, f64) {
    let scale = lonlat_scale(center_lat);
    let lon = center_lon + (x as f64) / scale.0;
    let lat = center_lat + (-(z as f64)) / scale.1;
    (lon, lat)
}

fn lonlat_scale(center_lat: f64) -> (f64, f64) {
    let meters_per_deg_lon = 111_320.0 * center_lat.to_radians().cos();
    let meters_per_deg_lat = 110_574.0;
    let n = 2.0_f64.powi(REFERENCE_ZOOM_INT);
    let tile_deg_lon = 360.0 / n;
    let tile_meters = tile_deg_lon * meters_per_deg_lon;
    let meters_to_world = TILE_SIZE as f64 / tile_meters;
    (
        meters_per_deg_lon * meters_to_world,
        meters_per_deg_lat * meters_to_world,
    )
}
