//! Editor selection treatment: active/secondary AABB shells, x-ray, and corner brackets.

use avian3d::prelude::ColliderConstructor;
use bevy::asset::{load_internal_asset, uuid_handle};
use bevy::camera::visibility::RenderLayers;
use bevy::pbr::{Material, MaterialPlugin};
use bevy::prelude::*;
use bevy::render::render_resource::{
    AsBindGroup, BlendState, CompareFunction, Face, RenderPipelineDescriptor, ShaderType,
};
use bevy::shader::ShaderRef;
use shared::{
    ColliderDesc, ColliderShapeDesc, LevelObject, PhysicsBody, Primitive, RaytracingExcluded,
};

use crate::infra::EditorOnly;
use crate::state::scene_tools_active;
use crate::ui::UiState;
use crate::viewport::ViewportDisplayMode;

const SELECTION_OUTLINE_SHADER: Handle<Shader> =
    uuid_handle!("c4e8f1a2-9b3d-4e7c-a1f5-6d2b8e0c4a91");

const PRIMARY_VISIBLE_COLOR: Color = Color::srgba(1.0, 0.72, 0.28, 0.56);
const PRIMARY_OCCLUDED_COLOR: Color = Color::srgba(1.0, 0.66, 0.22, 0.28);
const SECONDARY_VISIBLE_COLOR: Color = Color::srgba(0.31, 0.75, 0.88, 0.42);
const SECONDARY_OCCLUDED_COLOR: Color = Color::srgba(0.31, 0.68, 0.82, 0.2);
const BOUNDS_PADDING: f32 = 1.06;

#[derive(Clone, Copy, ShaderType)]
struct OutlineUniform {
    color: LinearRgba,
    rim_power: f32,
    rim_mix: f32,
    pass_kind: u32,
}

#[derive(Asset, TypePath, AsBindGroup, Clone)]
pub struct SelectionOutlineVisibleMaterial {
    #[uniform(0)]
    uniform: OutlineUniform,
}

#[derive(Asset, TypePath, AsBindGroup, Clone)]
pub struct SelectionOutlineOccludedMaterial {
    #[uniform(0)]
    uniform: OutlineUniform,
}

fn outline_uniform(color: Color, pass_kind: u32) -> OutlineUniform {
    OutlineUniform {
        color: color.into(),
        rim_power: 5.5,
        rim_mix: 0.9,
        pass_kind,
    }
}

macro_rules! impl_outline_material {
    ($t:ty, $depth:expr) => {
        impl Material for $t {
            fn fragment_shader() -> ShaderRef {
                SELECTION_OUTLINE_SHADER.into()
            }

            fn alpha_mode(&self) -> AlphaMode {
                AlphaMode::Blend
            }

            fn enable_prepass() -> bool {
                false
            }

            fn enable_shadows() -> bool {
                false
            }

            fn specialize(
                _pipeline: &bevy::pbr::MaterialPipeline,
                descriptor: &mut RenderPipelineDescriptor,
                _layout: &bevy::mesh::MeshVertexBufferLayoutRef,
                _key: bevy::pbr::MaterialPipelineKey<Self>,
            ) -> Result<(), bevy::render::render_resource::SpecializedMeshPipelineError> {
                descriptor.primitive.cull_mode = Some(Face::Front);
                if let Some(depth) = &mut descriptor.depth_stencil {
                    depth.depth_write_enabled = Some(false);
                    depth.depth_compare = Some($depth);
                }
                if let Some(fragment) = &mut descriptor.fragment {
                    if let Some(target) = fragment.targets.first_mut().and_then(|t| t.as_mut()) {
                        target.blend = Some(BlendState::ALPHA_BLENDING);
                    }
                }
                Ok(())
            }
        }
    };
}

impl_outline_material!(SelectionOutlineVisibleMaterial, CompareFunction::LessEqual);
impl_outline_material!(SelectionOutlineOccludedMaterial, CompareFunction::Greater);

#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
enum SelectionOutlinePass {
    Visible,
    Occluded,
}

#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
enum SelectionOutlineRole {
    Primary,
    Secondary,
}

#[derive(Component)]
pub(crate) struct SelectionOutlineShell {
    root: Entity,
    pass: SelectionOutlinePass,
    role: SelectionOutlineRole,
}

#[derive(Resource)]
struct SelectionOutlineMaterials {
    primary_occluded: Handle<SelectionOutlineOccludedMaterial>,
    primary_visible: Handle<SelectionOutlineVisibleMaterial>,
    secondary_occluded: Handle<SelectionOutlineOccludedMaterial>,
    secondary_visible: Handle<SelectionOutlineVisibleMaterial>,
}

#[derive(Resource)]
struct SelectionOutlineUnitBox(Handle<Mesh>);

#[derive(Clone, Copy, Debug)]
struct SelectionBounds {
    center: Vec3,
    half_extents: Vec3,
}

pub struct SelectionOutlinePlugin;

impl Plugin for SelectionOutlinePlugin {
    fn build(&self, app: &mut App) {
        load_internal_asset!(
            app,
            SELECTION_OUTLINE_SHADER,
            "../../assets/shaders/selection_outline.wgsl",
            Shader::from_wgsl
        );

        app.add_plugins((
            MaterialPlugin::<SelectionOutlineVisibleMaterial>::default(),
            MaterialPlugin::<SelectionOutlineOccludedMaterial>::default(),
        ))
        .add_systems(
            Startup,
            (
                init_selection_outline_assets,
                init_selection_outline_materials,
            )
                .chain(),
        )
        .add_systems(
            Update,
            (
                sync_selection_outline_meshes,
                draw_selection_corner_brackets,
            )
                .chain()
                .run_if(scene_tools_active),
        );
    }
}

fn init_selection_outline_assets(mut commands: Commands, mut meshes: ResMut<Assets<Mesh>>) {
    commands.insert_resource(SelectionOutlineUnitBox(
        meshes.add(Mesh::from(Cuboid::new(1.0, 1.0, 1.0))),
    ));
}

fn init_selection_outline_materials(
    mut visible: ResMut<Assets<SelectionOutlineVisibleMaterial>>,
    mut occluded: ResMut<Assets<SelectionOutlineOccludedMaterial>>,
    mut commands: Commands,
) {
    commands.insert_resource(SelectionOutlineMaterials {
        primary_visible: visible.add(SelectionOutlineVisibleMaterial {
            uniform: outline_uniform(PRIMARY_VISIBLE_COLOR, 0),
        }),
        primary_occluded: occluded.add(SelectionOutlineOccludedMaterial {
            uniform: outline_uniform(PRIMARY_OCCLUDED_COLOR, 1),
        }),
        secondary_visible: visible.add(SelectionOutlineVisibleMaterial {
            uniform: outline_uniform(SECONDARY_VISIBLE_COLOR, 0),
        }),
        secondary_occluded: occluded.add(SelectionOutlineOccludedMaterial {
            uniform: outline_uniform(SECONDARY_OCCLUDED_COLOR, 1),
        }),
    });
}

#[allow(clippy::too_many_arguments)]
fn sync_selection_outline_meshes(
    ui_state: Res<UiState>,
    display: Res<ViewportDisplayMode>,
    materials: Res<SelectionOutlineMaterials>,
    unit_box: Res<SelectionOutlineUnitBox>,
    mut commands: Commands,
    mut shells: Query<(Entity, &SelectionOutlineShell, &mut Transform)>,
    children: Query<&Children>,
    mesh_entities: Query<(Entity, &Mesh3d), With<Mesh3d>>,
    globals: Query<&GlobalTransform>,
    mesh_assets: Res<Assets<Mesh>>,
    primitives: Query<&Primitive>,
    physics: Query<&PhysicsBody>,
    colliders: Query<&ColliderDesc>,
    level_objects: Query<(), (With<LevelObject>, Without<EditorOnly>)>,
) {
    let selected: Vec<Entity> = ui_state.selected_entities.iter().collect();

    let mut desired: std::collections::HashMap<
        (Entity, SelectionOutlinePass, SelectionOutlineRole),
        SelectionBounds,
    > = std::collections::HashMap::new();
    if !display.clean_game_view {
        for (index, &root) in selected.iter().enumerate() {
            if !level_objects.contains(root) {
                continue;
            }
            let Some(bounds) = selection_bounds(
                root,
                &children,
                &mesh_entities,
                &globals,
                &mesh_assets,
                &primitives,
                &physics,
                &colliders,
            ) else {
                continue;
            };
            if bounds.half_extents.max_element() <= f32::EPSILON {
                continue;
            }
            let role = if index == 0 {
                SelectionOutlineRole::Primary
            } else {
                SelectionOutlineRole::Secondary
            };
            desired.insert((root, SelectionOutlinePass::Occluded, role), bounds);
            desired.insert((root, SelectionOutlinePass::Visible, role), bounds);
        }
    }

    let mut existing = std::collections::HashSet::new();
    for (shell_entity, shell, mut transform) in &mut shells {
        let key = (shell.root, shell.pass, shell.role);
        if let Some(bounds) = desired.get(&key) {
            transform.translation = bounds.center;
            transform.scale = bounds.half_extents * 2.0;
            existing.insert(key);
        } else {
            commands.entity(shell_entity).despawn();
        }
    }

    for ((root, pass, role), bounds) in desired {
        if existing.contains(&(root, pass, role)) {
            continue;
        }

        let mut entity_commands = commands.spawn((
            EditorOnly,
            RaytracingExcluded,
            Name::new("Selection Outline"),
            SelectionOutlineShell { root, pass, role },
            Mesh3d(unit_box.0.clone()),
            Transform {
                translation: bounds.center,
                scale: bounds.half_extents * 2.0,
                ..default()
            },
            RenderLayers::default(),
            ChildOf(root),
        ));

        match (role, pass) {
            (SelectionOutlineRole::Primary, SelectionOutlinePass::Occluded) => {
                entity_commands.insert(MeshMaterial3d(materials.primary_occluded.clone()));
            }
            (SelectionOutlineRole::Primary, SelectionOutlinePass::Visible) => {
                entity_commands.insert(MeshMaterial3d(materials.primary_visible.clone()));
            }
            (SelectionOutlineRole::Secondary, SelectionOutlinePass::Occluded) => {
                entity_commands.insert(MeshMaterial3d(materials.secondary_occluded.clone()));
            }
            (SelectionOutlineRole::Secondary, SelectionOutlinePass::Visible) => {
                entity_commands.insert(MeshMaterial3d(materials.secondary_visible.clone()));
            }
        }
    }
}

fn draw_selection_corner_brackets(
    display: Res<ViewportDisplayMode>,
    mut gizmos: Gizmos,
    shells: Query<(&SelectionOutlineShell, &GlobalTransform)>,
) {
    if display.clean_game_view {
        return;
    }

    for (shell, global) in &shells {
        if shell.pass != SelectionOutlinePass::Visible {
            continue;
        }
        let (scale, rotation, center) = global.to_scale_rotation_translation();
        let half = scale.abs() * 0.5;
        if half.max_element() <= f32::EPSILON {
            continue;
        }
        let color = match shell.role {
            SelectionOutlineRole::Primary => Color::srgba(1.0, 0.76, 0.34, 0.98),
            SelectionOutlineRole::Secondary => Color::srgba(0.38, 0.78, 0.9, 0.78),
        };
        draw_corner_brackets(&mut gizmos, center, rotation, half, color);
    }
}

fn draw_corner_brackets(
    gizmos: &mut Gizmos,
    center: Vec3,
    rotation: Quat,
    half: Vec3,
    color: Color,
) {
    let lengths = (half * 0.28).clamp(Vec3::splat(0.06), Vec3::splat(0.55));
    for signs in [
        Vec3::new(-1.0, -1.0, -1.0),
        Vec3::new(1.0, -1.0, -1.0),
        Vec3::new(-1.0, 1.0, -1.0),
        Vec3::new(1.0, 1.0, -1.0),
        Vec3::new(-1.0, -1.0, 1.0),
        Vec3::new(1.0, -1.0, 1.0),
        Vec3::new(-1.0, 1.0, 1.0),
        Vec3::new(1.0, 1.0, 1.0),
    ] {
        let corner = center + rotation * (half * signs);
        for axis in [Vec3::X, Vec3::Y, Vec3::Z] {
            let length = lengths.dot(axis);
            let direction = rotation * (axis * -signs.dot(axis));
            gizmos.line(corner, corner + direction * length, color);
        }
    }
}

#[expect(
    clippy::too_many_arguments,
    reason = "selection bounds combine distinct authored geometry and physics sources"
)]
fn selection_bounds(
    root: Entity,
    children: &Query<&Children>,
    mesh_entities: &Query<(Entity, &Mesh3d), With<Mesh3d>>,
    globals: &Query<&GlobalTransform>,
    mesh_assets: &Assets<Mesh>,
    primitives: &Query<&Primitive>,
    physics: &Query<&PhysicsBody>,
    colliders: &Query<&ColliderDesc>,
) -> Option<SelectionBounds> {
    if let Ok(primitive) = primitives.get(root) {
        return Some(SelectionBounds {
            center: Vec3::ZERO,
            half_extents: primitive.size * 0.5 * BOUNDS_PADDING,
        });
    }

    if let Ok(body) = physics.get(root) {
        if let Some(bounds) = bounds_from_collider(&body.collider) {
            return Some(bounds);
        }
    }

    if let Ok(collider) = colliders.get(root) {
        if let Some(bounds) = bounds_from_collider_desc(collider) {
            return Some(bounds);
        }
    }

    union_mesh_bounds(root, children, mesh_entities, globals, mesh_assets)
}

fn bounds_from_collider(collider: &ColliderConstructor) -> Option<SelectionBounds> {
    match collider {
        ColliderConstructor::Cuboid {
            x_length,
            y_length,
            z_length,
        } => Some(SelectionBounds {
            center: Vec3::ZERO,
            half_extents: Vec3::new(*x_length, *y_length, *z_length) * 0.5 * BOUNDS_PADDING,
        }),
        ColliderConstructor::Sphere { radius } => Some(SelectionBounds {
            center: Vec3::ZERO,
            half_extents: Vec3::splat(*radius) * BOUNDS_PADDING,
        }),
        _ => None,
    }
}

fn bounds_from_collider_desc(collider: &ColliderDesc) -> Option<SelectionBounds> {
    if !collider.enabled {
        return None;
    }
    match &collider.shape {
        ColliderShapeDesc::Cuboid {
            x_length,
            y_length,
            z_length,
        } => Some(SelectionBounds {
            center: Vec3::ZERO,
            half_extents: Vec3::new(*x_length, *y_length, *z_length) * 0.5 * BOUNDS_PADDING,
        }),
        ColliderShapeDesc::Sphere { radius } => Some(SelectionBounds {
            center: Vec3::ZERO,
            half_extents: Vec3::splat(*radius) * BOUNDS_PADDING,
        }),
        ColliderShapeDesc::Capsule { radius, height } => Some(SelectionBounds {
            center: Vec3::ZERO,
            half_extents: Vec3::new(*radius, *height * 0.5, *radius) * BOUNDS_PADDING,
        }),
        ColliderShapeDesc::StaticMesh { .. } => None,
    }
}

fn union_mesh_bounds(
    root: Entity,
    children: &Query<&Children>,
    mesh_entities: &Query<(Entity, &Mesh3d), With<Mesh3d>>,
    globals: &Query<&GlobalTransform>,
    mesh_assets: &Assets<Mesh>,
) -> Option<SelectionBounds> {
    let root_global = globals.get(root).ok()?;
    let root_to_world = root_global.affine();
    let world_to_root = root_to_world.inverse();

    let mut min = Vec3::splat(f32::INFINITY);
    let mut max = Vec3::splat(f32::NEG_INFINITY);
    let mut any = false;

    for entity in collect_mesh_entities(root, children, mesh_entities) {
        let Ok((_, mesh3d)) = mesh_entities.get(entity) else {
            continue;
        };
        let Some(mesh) = mesh_assets.get(&mesh3d.0) else {
            continue;
        };
        let Some(aabb) = mesh.final_aabb else {
            continue;
        };
        let entity_global = globals.get(entity).ok()?;
        let entity_to_world = entity_global.affine();

        let mesh_min = Vec3::from(aabb.min);
        let mesh_max = Vec3::from(aabb.max);
        for corner in corner_iter(mesh_min, mesh_max) {
            let world = entity_to_world.transform_point3(corner);
            let local = world_to_root.transform_point3(world);
            min = min.min(local);
            max = max.max(local);
            any = true;
        }
    }

    if !any {
        return None;
    }

    let center = (min + max) * 0.5;
    let half_extents = (max - min) * 0.5 * BOUNDS_PADDING;
    Some(SelectionBounds {
        center,
        half_extents,
    })
}

fn corner_iter(min: Vec3, max: Vec3) -> impl Iterator<Item = Vec3> {
    [
        Vec3::new(min.x, min.y, min.z),
        Vec3::new(max.x, min.y, min.z),
        Vec3::new(min.x, max.y, min.z),
        Vec3::new(max.x, max.y, min.z),
        Vec3::new(min.x, min.y, max.z),
        Vec3::new(max.x, min.y, max.z),
        Vec3::new(min.x, max.y, max.z),
        Vec3::new(max.x, max.y, max.z),
    ]
    .into_iter()
}

fn collect_mesh_entities(
    root: Entity,
    children: &Query<&Children>,
    meshes: &Query<(Entity, &Mesh3d), With<Mesh3d>>,
) -> Vec<Entity> {
    let mut found = Vec::new();
    if meshes.get(root).is_ok() {
        found.push(root);
    }
    collect_mesh_descendants(root, children, meshes, &mut found);
    found
}

fn collect_mesh_descendants(
    entity: Entity,
    children: &Query<&Children>,
    meshes: &Query<(Entity, &Mesh3d), With<Mesh3d>>,
    found: &mut Vec<Entity>,
) {
    let Ok(kids) = children.get(entity) else {
        return;
    };
    for child in kids.iter() {
        if meshes.get(child).is_ok() {
            found.push(child);
        }
        collect_mesh_descendants(child, children, meshes, found);
    }
}