bevy-expedition-demo/src/inventory/ui.rs

use std::f32::consts::FRAC_PI_2;

use bevy::{ecs::relationship::RelatedSpawner, prelude::*, ui_widgets::{ControlOrientation, CoreScrollbarThumb, Scrollbar}};

use crate::{inventory::{ActiveInventory, Inventory, item::Item}, ui::{UiRoot, UiRotateEvent}};

const UI_SLOT_ASSET_PATH: &'static str = "sprites/ui/inventory_slot.png";
const TEMP_ITEM_PATH: &'static str = "sprites/items/choco_bar.png";

#[derive(Component, Reflect)]
#[require(Node)]
pub struct UiInventoryManager;

#[derive(Component, Reflect)]
#[require(Node)]
pub struct UiInventory(pub Entity);

#[derive(Component, Reflect)]
#[require(Node, ImageNode)]
pub struct UiInventorySlot(pub UVec2);

#[derive(Component, Reflect)]
#[require(Node, ImageNode)]
pub struct UiItem(pub Entity);

#[derive(Component, Reflect)]
pub struct HoveredItem;

#[derive(Component, Reflect)]
pub struct HoveredSlot;

#[derive(Component, Reflect)]
pub struct DraggedItem(pub Item, pub UVec2);

fn ui_item_node_data(item: &Item) -> (Val, Val, Val, Val, UiTransform) {
    match item.rotated {
        true => (
            percent(50. * (item.size.x as f32 - item.size.y as f32)),
            percent(50. * (item.size.y as f32 - item.size.x as f32)),
            percent(100. * item.size.y as f32),
            percent(100. * item.size.x as f32),
            UiTransform::from_rotation(Rot2::radians(-FRAC_PI_2)),
        ),
        false => (
            auto(),
            auto(),
            percent(100. * item.size.x as f32),
            percent(100. * item.size.y as f32),
            UiTransform::default(),
        ),
    }
}

fn update_ui_node(item: &Item, mut node: Mut<'_, Node>, mut ui_transform: Mut<'_, UiTransform>) {
    let (left, top, min_width, min_height, new_ui_transform) = ui_item_node_data(item);
    node.left = left;
    node.top = top;
    node.min_width = min_width;
    node.min_height = min_height;
    ui_transform.rotation = new_ui_transform.rotation;
}

fn reset_slots_colors(query: Query<&mut ImageNode, With<UiInventorySlot>>) {
    for mut image in query {
        image.color = Color::WHITE;
    }
}

fn on_slot_over(
    e: On<Pointer<Over>>,
    mut commands: Commands,
    mut image_query: Query<&mut ImageNode, With<UiInventorySlot>>,
    hovered_slots: Query<Entity, With<HoveredSlot>>,
    has_dragged_item: Option<Single<(), With<DraggedItem>>>,
) {
    if let Ok(mut image) = image_query.get_mut(e.event_target()) {
        image.color = Color::WHITE.darker(0.3);
    };
    for slot_id in hovered_slots {
        commands.entity(slot_id).remove::<HoveredSlot>();
    }
    if has_dragged_item.is_none() {
        commands.entity(e.event_target()).insert(HoveredSlot);
    }
}

fn on_slot_out(e: On<Pointer<Out>>, mut query: Query<&mut ImageNode, With<UiInventorySlot>>) {
    if let Ok(mut image) = query.get_mut(e.event_target()) {
        image.color = Color::WHITE;
    }
}

fn on_item_over(
    e: On<Pointer<Over>>,
    mut commands: Commands,
    query: Query<(), With<UiItem>>,
    has_hovered_item: Option<Single<(), With<HoveredItem>>>,
) {
    if has_hovered_item.is_some() {
        return;
    }
    if let Ok(_) = query.get(e.event_target()) {
        commands.entity(e.event_target()).insert(HoveredItem);
    }
}

fn on_item_out(
    e: On<Pointer<Out>>,
    mut commands: Commands,
    query: Query<(), (With<UiItem>, With<HoveredItem>)>,
) {
    if let Ok(_) = query.get(e.event_target()) {
        commands.entity(e.event_target()).remove::<HoveredItem>();
    }
}

fn on_item_drag_start(
    e: On<Pointer<DragStart>>,
    mut commands: Commands,
    ui_query: Query<&UiItem, With<HoveredItem>>,
    item_query: Query<&Item>,
    hovered_slot: Single<&UiInventorySlot, With<HoveredSlot>>
) {
    if let Ok(ui_item) = ui_query.get(e.event_target()) {
        let Ok(item) = item_query.get(ui_item.0) else {
            error!("UiItem {} is pointing to non-existing Item", e.event_target());
            return;
        };
        let Some(item_position) = item.position else {
            return;
        };
        let slot_position = hovered_slot.0;
        let diff = slot_position - item_position;
        commands.entity(e.event_target()).insert(DraggedItem(item.clone(), diff));
    }
}

fn on_item_drag(e: On<Pointer<Drag>>, mut query: Query<&mut UiTransform, With<UiItem>>) {
    if let Ok(mut transform) = query.get_mut(e.event_target()) {
        transform.translation = Val2::px(e.distance.x, e.distance.y);
    }
}

fn on_item_drag_end(
    e: On<Pointer<DragEnd>>,
    mut commands: Commands,
    mut query: Query<&mut UiTransform, With<UiItem>>,
) {
    if let Ok(mut transform) = query.get_mut(e.event_target()) {
        transform.translation = Val2::ZERO;
        commands.entity(e.event_target()).remove::<DraggedItem>();
    }
}

fn on_item_drag_drop(
    event: On<Pointer<DragDrop>>,
    mut commands: Commands,
    mut ui_item_query: Query<(Entity, &UiItem, &mut DraggedItem, &mut Node, &mut UiTransform)>,
    ui_inventory_query: Query<&UiInventory>,
    mut item_query: Query<&mut Item>,
    slot_query: Query<(&ChildOf, &UiInventorySlot, Option<&Children>)>,
    slot_id_query: Query<(Entity, &UiInventorySlot, &ChildOf), With<UiInventorySlot>>,
    inventory_query: Query<(&Inventory, Option<&Children>)>,
) {
    let Ok((ui_item_entity, UiItem(item_entity), mut dragged_item, node, ui_transform)) = ui_item_query.get_mut(event.dropped) else {
        return;
    };
    let Ok((slot_parent, UiInventorySlot(new_position), slot_children)) = slot_query.get(event.event_target()) else {
        return;
    };
    if slot_children.is_some() {
        return;
    }
    let Ok(UiInventory(inventory_id)) = ui_inventory_query.get(slot_parent.0) else {
        return;
    };
    let Ok((inventory, inventory_children)) = inventory_query.get(*inventory_id) else {
        return;
    };

    let items = match inventory_children {
        Some(children) => &children[..],
        None => &[],
    };
    
    let actual_position = new_position.as_ivec2() - dragged_item.1.as_ivec2();
    info!("{actual_position:?}");
    let temp_item = &mut dragged_item.0;
    if actual_position.is_negative_bitmask() == 0 {
        temp_item.position = Some(actual_position.as_uvec2());
    }

    let Some((slot_id, _, _)) = slot_id_query.iter().find(|(_, slot_pos, slot_parent_id)| slot_pos.0 == temp_item.position.unwrap() && slot_parent_id == &slot_parent) else {
        return;
    };

    if inventory.can_replace(item_query.as_readonly(), items, *item_entity, temp_item) {
        let mut item = item_query.get_mut(*item_entity).unwrap();
        item.position = temp_item.position;
        item.size = temp_item.size;
        item.rotated = temp_item.rotated;
        commands.entity(ui_item_entity).insert(ChildOf(slot_id));
        commands.entity(*item_entity).insert(ChildOf(*inventory_id));
        commands.entity(slot_id).insert(HoveredSlot);
        update_ui_node(item.as_ref(), node, ui_transform);
    } else {
        if let Ok(item) = item_query.get(*item_entity) {
            update_ui_node(item, node, ui_transform);
        }
    }
    commands.run_system_cached(reset_slots_colors);
}

pub fn on_ui_rotate(
    _: On<UiRotateEvent>,
    mut item_query: Query<&mut Item>,
    parent_query: Query<&ChildOf>,
    inventory_query: Query<(&Inventory, Option<&Children>)>,
    ui_item_query: Query<(&UiItem, &mut UiTransform, &mut Node, Option<&mut DraggedItem>), With<HoveredItem>>,
) {
    for (ui_item, ui_transform, node, maybe_dragged) in ui_item_query {
        let Ok(item_parent) = parent_query.get(ui_item.0) else {
            continue;
        };
        let Ok((inventory, children)) = inventory_query.get(item_parent.0) else {
            continue;
        };

        let children = match children {
            Some(children) => &children[..],
            None => &[],
        };
        let (was_rotated, item) = match maybe_dragged {
            Some(mut temp_item) => {
                temp_item.0.rotate();
                temp_item.1 = uvec2(temp_item.1.y, temp_item.1.x);
                (true, Some(temp_item.0.clone()))
            },
            None => {
                let result = inventory.can_rotate(item_query.as_readonly(), children, ui_item.0);
                let mut out_item = None;
                if result {
                    if let Ok(mut item) = item_query.get_mut(ui_item.0) {
                        item.rotate();
                        out_item = Some(item.clone());
                    }
                }
                (result, out_item)
            },
        };
        if was_rotated {
            update_ui_node(&item.unwrap(), node, ui_transform);
        }
    }
}

fn ui_manager_bundle(children: Vec<Entity>) -> impl Bundle {
    (
        UiInventoryManager,
        Node {
            position_type: PositionType::Absolute,
            width: percent(50.),
            height: percent(100.),
            scrollbar_width: 8.,
            display: Display::Grid,
            grid_template_columns: vec![RepeatedGridTrack::flex(1, 1.), RepeatedGridTrack::auto(1)],
            grid_template_rows: vec![RepeatedGridTrack::flex(1, 1.), RepeatedGridTrack::auto(1)],
            ..default()
        },
        Pickable::IGNORE,
        GlobalZIndex::default(),
        Children::spawn(SpawnWith(move |parent: &mut RelatedSpawner<ChildOf>| {
            let scroll_area_id = parent.spawn((
                Node {
                    width: percent(100.),
                    height: percent(100.),
                    display: Display::Flex,
                    flex_direction: FlexDirection::Column,
                    align_items: AlignItems::Center,
                    justify_content: JustifyContent::SpaceAround,
                    overflow: Overflow::scroll_y(),
                    ..default()
                },
                ScrollPosition(Vec2::ZERO),
            )).add_children(children.as_slice()).id();
            parent.spawn((
                Node {
                    min_width: px(8),
                    grid_row: GridPlacement::start(1),
                    grid_column: GridPlacement::start(2),
                    ..default()
                },
                Scrollbar {
                    orientation: ControlOrientation::Vertical,
                    target: scroll_area_id,
                    min_thumb_length: 8.0,
                },
                BackgroundColor(Color::hsl(0., 0., 0.5)),
                children![(
                    Node {
                        position_type: PositionType::Absolute,
                        border_radius: BorderRadius::all(px(4.)),
                        ..default()
                    },
                    BackgroundColor(Color::hsl(0., 0., 0.3)),
                    CoreScrollbarThumb,
                )],
            ));
        })),
    )
}

fn ui_inventory_bundle(
    inventory: &Inventory,
    inventory_entity: Entity,
) -> impl Bundle {
    (
        UiInventory(inventory_entity),
        Node {
            align_self: AlignSelf::Stretch,
            align_content: AlignContent::Center,
            display: Display::Grid,
            width: percent(100.),
            aspect_ratio: Some(inventory.size.x as f32 / inventory.size.y as f32),
            grid_auto_columns: vec![GridTrack::percent(100. / inventory.size.x as f32)],
            grid_auto_rows: vec![GridTrack::percent(100. / inventory.size.y as f32)],
            ..default()
        },
        Pickable::IGNORE,
        GlobalZIndex::default(),
        Name::new(format!("UiInventory ({}x{})", inventory.size.x, inventory.size.y)),
    )
}

fn inventory_slot_bundle(x: u32, y: u32, image: Handle<Image>) -> impl Bundle {
    (
        UiInventorySlot(UVec2::new(x, y)),
        ImageNode {
            color: Color::WHITE,
            image,
            image_mode: NodeImageMode::Stretch,
            ..default()
        },
        Node {
            width: percent(100.),
            height: percent(100.),
            grid_column: GridPlacement::start(x as i16 + 1),
            grid_row: GridPlacement::start(y as i16 + 1),
            aspect_ratio: Some(1.),
            ..default()
        },
        Pickable {
            should_block_lower: true,
            is_hoverable: true,
        },
        BackgroundColor::DEFAULT,
        Name::new(format!("UiInventorySlot({x},{y})")),
    )
}

fn ui_item_bundle(item: &Item, item_entity: Entity, image: Handle<Image>) -> impl Bundle {
    let (left, top, min_width, min_height, ui_transform) = ui_item_node_data(item);
    (
        UiItem(item_entity),
        ImageNode {
            image,
            image_mode: NodeImageMode::Stretch,
            ..default()
        },
        Node {
            left,
            top,
            min_width,
            min_height,
            ..default()
        },
        BackgroundColor(Color::hsla(0., 0., 0., 0.5)),
        ui_transform,
        GlobalZIndex(1),
        Pickable {
            should_block_lower: false,
            is_hoverable: true,
        },
        Name::new(format!("UiItem ({},{})",
            item.position.unwrap_or_default().x,
            item.position.unwrap_or_default().y,
        )),
    )
}

pub fn setup_ui_inventory(
    mut commands: Commands,
    asset_server: Res<AssetServer>,
    inventory_query: Query<(Entity, &Inventory, Option<&Children>), With<ActiveInventory>>,
    item_query: Query<&Item>,
    root_query: Query<Entity, With<UiRoot>>,
) {
    let Ok(root) = root_query.single() else {
        error!("Query contains more than one UiRoot");
        return;
    };
    let ui_slot_image: Handle<Image> = asset_server.load(UI_SLOT_ASSET_PATH);
    let temp_item_image: Handle<Image> = asset_server.load(TEMP_ITEM_PATH);
    let mut inventory_ids = Vec::new();
    for (inventory_entity, inventory, children) in inventory_query.iter().sort::<Entity>() {
        let items = match children {
            Some(children) => {
                children.iter().filter_map(|item_entity| {
                    match item_query.get(item_entity) {
                        Ok(item) => Some((item, item_entity)),
                        Err(err) => {
                            warn!("Error querying item {item_entity}: {err}");
                            None
                        },
                    }
                }).collect::<Vec<(&Item, Entity)>>()
            }
            None => Vec::new(),
        };
        let inventory_id = commands.spawn(ui_inventory_bundle(inventory, inventory_entity))
            .with_children(|commands| {
                for x in 0..inventory.size.x { for y in 0..inventory.size.y {
                    let mut slot_commands = commands.spawn(inventory_slot_bundle(x, y, ui_slot_image.clone()));
                    slot_commands.observe(on_slot_over)
                        .observe(on_slot_out)
                        .observe(on_item_drag_drop);
                    if let Some((item, entity)) = items.iter()
                        .find(|(i, _)| i.position.unwrap_or_default() == UVec2::new(x, y)) {
                        slot_commands.with_children(|commands| {
                            commands.spawn(ui_item_bundle(item, *entity, temp_item_image.clone()))
                                .observe(on_item_over)
                                .observe(on_item_out)
                                .observe(on_item_drag_start)
                                .observe(on_item_drag)
                                .observe(on_item_drag_end);
                        });
                    }
                } }
            }).id();
        inventory_ids.push(inventory_id);
    }
    let inventory_manager = commands.spawn(ui_manager_bundle(inventory_ids)).id();
    commands.entity(root).add_child(inventory_manager);
}

pub fn clear_ui_inventory(
    mut commands: Commands,
    inventory_query: Query<Entity, With<UiInventoryManager>>,
) {
    for entity in inventory_query {
        commands.entity(entity).despawn();
    }
}