//! Optional overlays over an ordinary [`Stack`].
//!
//! An overlay host owns independent single and stackable controllers. Opening
//! any overlay removes the active single on each host the request traverses;
//! stackables remain in opening order, and the host's optional single is
//! always its final child. Controller lookup follows the active parent map
//! maintained by `EventManager`, so callers need only the requesting widget,
//! not access to `UiRenderState`.

use crate::prelude::*;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum OverlayKind {
    Single,
    Stackable,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CancelBehavior {
    Consume,
    PassThrough,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct OverlayOptions {
    pub kind: OverlayKind,
    pub cancel: CancelBehavior,
    pub blocks_prior_input: bool,
}

impl OverlayOptions {
    pub const fn single(cancel: CancelBehavior) -> Self {
        Self {
            kind: OverlayKind::Single,
            cancel,
            blocks_prior_input: false,
        }
    }

    pub const fn stackable(cancel: CancelBehavior) -> Self {
        Self {
            kind: OverlayKind::Stackable,
            cancel,
            blocks_prior_input: false,
        }
    }

    pub const fn blocking(mut self) -> Self {
        self.blocks_prior_input = true;
        self
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct OverlayHostOptions {
    pub single: bool,
    pub stackable: bool,
}

impl OverlayHostOptions {
    pub const BOTH: Self = Self {
        single: true,
        stackable: true,
    };
    pub const SINGLE: Self = Self {
        single: true,
        stackable: false,
    };
    pub const STACKABLE: Self = Self {
        single: false,
        stackable: true,
    };
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct OverlayHandle {
    host: WidgetId,
    kind: OverlayKind,
    serial: u64,
}

impl OverlayHandle {
    pub fn close<Rsc: HasEvents>(self, rsc: &mut Rsc) -> bool {
        match self.kind {
            OverlayKind::Single => {
                let Some(id) = rsc
                    .events()
                    .controllers
                    .id::<SingleOverlayController>(self.host)
                else {
                    return false;
                };
                rsc.with_controller(id, |controller: &mut SingleOverlayController, rsc| {
                    controller.close(id, self.serial, rsc)
                })
                .unwrap_or(false)
            }
            OverlayKind::Stackable => {
                let Some(id) = rsc
                    .events()
                    .controllers
                    .id::<StackableOverlayController>(self.host)
                else {
                    return false;
                };
                rsc.with_controller(id, |controller: &mut StackableOverlayController, rsc| {
                    controller.close(id, self.serial, rsc)
                })
                .unwrap_or(false)
            }
        }
    }
}

struct OverlayEntry {
    serial: u64,
    widget: WidgetId,
    previous_target: Option<ControllerId>,
    cancel: CancelBehavior,
    blocks_prior_input: bool,
}

pub struct SingleOverlayController {
    stack: WeakWidget<Stack>,
    active: Option<OverlayEntry>,
    next_serial: u64,
}

impl SingleOverlayController {
    fn new(stack: WeakWidget<Stack>) -> Self {
        Self {
            stack,
            active: None,
            next_serial: 1,
        }
    }

    fn open<Rsc: HasEvents>(
        &mut self,
        id: ControllerId,
        popup: StrongWidget,
        cancel: CancelBehavior,
        blocks_prior_input: bool,
        rsc: &mut Rsc,
    ) -> OverlayHandle {
        self.dismiss(id, rsc, true);
        let serial = self.take_serial();
        let widget = popup.id();
        let previous_target = rsc.events().controllers.command_target();
        (self.stack)(rsc).children.push(popup);
        self.active = Some(OverlayEntry {
            serial,
            widget,
            previous_target,
            cancel,
            blocks_prior_input,
        });
        rsc.set_command_target(Some(id));
        OverlayHandle {
            host: id.host(),
            kind: OverlayKind::Single,
            serial,
        }
    }

    fn close<Rsc: HasEvents>(&mut self, id: ControllerId, serial: u64, rsc: &mut Rsc) -> bool {
        if self.active.as_ref().map(|entry| entry.serial) != Some(serial) {
            return false;
        }
        self.dismiss(id, rsc, false);
        true
    }

    fn dismiss<Rsc: HasEvents>(
        &mut self,
        id: ControllerId,
        rsc: &mut Rsc,
        replacement: bool,
    ) -> bool {
        let Some(entry) = self.active.take() else {
            return false;
        };
        let current_target = rsc.events().controllers.command_target();
        let was_target = current_target == Some(id);
        remove_child(self.stack, entry.widget, rsc);
        if was_target {
            rsc.set_command_target(entry.previous_target);
        }
        if replacement && entry.cancel == CancelBehavior::PassThrough {
            rsc.set_command_target(entry.previous_target);
            let _ = rsc.run_command_before(Command::Escape, self.stack);
            if !was_target {
                rsc.set_command_target(current_target);
            }
        }
        true
    }

    fn take_serial(&mut self) -> u64 {
        let serial = self.next_serial;
        self.next_serial = self.next_serial.wrapping_add(1).max(1);
        serial
    }
}

impl<Rsc: HasEvents> Controller<Rsc> for SingleOverlayController {
    fn command(&mut self, command: Command, rsc: &mut Rsc) -> CommandResult {
        if command != Command::Escape {
            return CommandResult::Unused;
        }
        let Some(entry) = self.active.take() else {
            return CommandResult::Unused;
        };
        remove_child(self.stack, entry.widget, rsc);
        rsc.set_command_target(entry.previous_target);
        if entry.cancel == CancelBehavior::PassThrough {
            let _ = rsc.run_command(command);
        }
        CommandResult::Used
    }

    fn blocks_prior_input(&self) -> bool {
        self.active
            .as_ref()
            .is_some_and(|entry| entry.blocks_prior_input)
    }
}

pub struct StackableOverlayController {
    stack: WeakWidget<Stack>,
    entries: Vec<OverlayEntry>,
    next_serial: u64,
}

impl StackableOverlayController {
    fn new(stack: WeakWidget<Stack>) -> Self {
        Self {
            stack,
            entries: Vec::new(),
            next_serial: 1,
        }
    }

    fn open<Rsc: HasEvents>(
        &mut self,
        id: ControllerId,
        overlay: StrongWidget,
        cancel: CancelBehavior,
        blocks_prior_input: bool,
        rsc: &mut Rsc,
    ) -> OverlayHandle {
        let serial = self.take_serial();
        let widget = overlay.id();
        let previous_target = rsc.events().controllers.command_target();
        (self.stack)(rsc).children.push(overlay);
        self.entries.push(OverlayEntry {
            serial,
            widget,
            previous_target,
            cancel,
            blocks_prior_input,
        });
        rsc.set_command_target(Some(id));
        OverlayHandle {
            host: id.host(),
            kind: OverlayKind::Stackable,
            serial,
        }
    }

    fn close<Rsc: HasEvents>(&mut self, id: ControllerId, serial: u64, rsc: &mut Rsc) -> bool {
        let Some(index) = self.entries.iter().position(|entry| entry.serial == serial) else {
            return false;
        };
        close_single_at(id.host(), rsc, false);
        let previous_target = self.entries[index].previous_target;
        for entry in self.entries.drain(index..) {
            remove_child(self.stack, entry.widget, rsc);
        }
        if rsc.events().controllers.command_target() == Some(id) {
            rsc.set_command_target(previous_target);
        }
        true
    }

    fn take_serial(&mut self) -> u64 {
        let serial = self.next_serial;
        self.next_serial = self.next_serial.wrapping_add(1).max(1);
        serial
    }
}

impl<Rsc: HasEvents> Controller<Rsc> for StackableOverlayController {
    fn command(&mut self, command: Command, rsc: &mut Rsc) -> CommandResult {
        if command != Command::Escape {
            return CommandResult::Unused;
        }
        let mut used = false;
        loop {
            let Some(entry) = self.entries.pop() else {
                return if used {
                    CommandResult::Used
                } else {
                    CommandResult::Unused
                };
            };
            used = true;
            remove_child(self.stack, entry.widget, rsc);
            rsc.set_command_target(entry.previous_target);
            if entry.cancel == CancelBehavior::Consume {
                return CommandResult::Used;
            }
            let same_controller = entry.previous_target.is_some_and(|previous| {
                previous.host() == self.stack.id() && previous.is::<StackableOverlayController>()
            });
            if !same_controller {
                let _ = rsc.run_command(command);
                return CommandResult::Used;
            }
        }
    }

    fn blocks_prior_input(&self) -> bool {
        self.entries
            .last()
            .is_some_and(|entry| entry.blocks_prior_input)
    }
}

fn remove_child<Rsc: UiRsc>(stack: WeakWidget<Stack>, widget: WidgetId, rsc: &mut Rsc) {
    if let Some(stack) = rsc.widgets_mut().get_mut(&stack)
        && let Some(index) = stack.children.iter().position(|child| child.id() == widget)
    {
        stack.children.remove(index);
    }
}

fn close_single_at<Rsc: HasEvents>(host: WidgetId, rsc: &mut Rsc, replacement: bool) {
    let Some(id) = rsc.events().controllers.id::<SingleOverlayController>(host) else {
        return;
    };
    rsc.with_controller(id, |single: &mut SingleOverlayController, rsc| {
        single.dismiss(id, rsc, replacement);
    });
}

fn add_host<Rsc: HasEvents>(
    rsc: &mut Rsc,
    base: StrongWidget,
    options: OverlayHostOptions,
) -> WeakWidget<Stack> {
    let stack = Stack {
        children: vec![base],
        size: StackSize::Default,
    }
    .add(rsc);
    if options.single {
        rsc.register_controller(stack, SingleOverlayController::new(stack));
    }
    if options.stackable {
        rsc.register_controller(stack, StackableOverlayController::new(stack));
    }
    stack
}

pub fn overlay_host<Rsc, W, Tag>(
    content: W,
    options: OverlayHostOptions,
) -> impl WidgetIdFn<Rsc, Stack>
where
    Rsc: HasEvents,
    W: WidgetLike<Rsc, Tag>,
{
    move |rsc: &mut Rsc| {
        let content = content.add_strong(rsc);
        add_host(rsc, content, options)
    }
}

pub(crate) fn default_overlay_root<Rsc: HasEvents>(
    rsc: &mut Rsc,
    root: StrongWidget,
) -> StrongWidget {
    add_host(rsc, root, OverlayHostOptions::BOTH)
        .upgrade(rsc)
        .any()
}

pub trait OverlayRscExt: HasEvents {
    fn open_popup<W, Tag>(&mut self, origin: impl IdLike, popup: W) -> Option<OverlayHandle>
    where
        W: WidgetLike<Self, Tag>,
    {
        self.open_overlay(
            origin,
            popup,
            OverlayOptions::single(CancelBehavior::PassThrough),
        )
    }

    fn open_modal<W, Tag>(&mut self, origin: impl IdLike, modal: W) -> Option<OverlayHandle>
    where
        W: WidgetLike<Self, Tag>,
    {
        let render_handle = self.ui().render_state();
        let (target, path) = self
            .events()
            .controllers
            .path_to::<StackableOverlayController>(origin.id(), &render_handle.get())?;
        clear_singles(&path, self);

        let content = modal.add(self);
        let content_strong = content.upgrade(self).any();
        let backdrop = rect(Srgba8::new(0, 0, 0, 96)).add(self);
        let backdrop_strong = backdrop.upgrade(self).any();
        let centered = Aligned {
            inner: content_strong,
            align: Align::CENTER.into(),
        }
        .add_strong(self)
        .any();
        let layer = Stack {
            children: vec![backdrop_strong, centered],
            size: StackSize::Default,
        }
        .add_strong(self)
        .any();
        let handle = open_stackable_at(target, layer, CancelBehavior::Consume, true, self)?;

        let blocking = modal_pointer_senses();
        self.register_event(content, blocking.clone(), |_, _| {});
        self.register_event(backdrop, blocking, move |ctx, rsc| {
            if matches!(ctx.data.sense, CursorSense::PressStart(_)) {
                handle.close(rsc);
            }
        });
        Some(handle)
    }

    fn open_overlay<W, Tag>(
        &mut self,
        origin: impl IdLike,
        overlay: W,
        options: OverlayOptions,
    ) -> Option<OverlayHandle>
    where
        W: WidgetLike<Self, Tag>,
    {
        match options.kind {
            OverlayKind::Single => {
                let render_handle = self.ui().render_state();
                let target = self
                    .events()
                    .controllers
                    .nearest_id::<SingleOverlayController>(origin.id(), &render_handle.get())?;
                let overlay = overlay.add_strong(self).any();
                self.with_controller(target, |single: &mut SingleOverlayController, rsc| {
                    single.open(
                        target,
                        overlay,
                        options.cancel,
                        options.blocks_prior_input,
                        rsc,
                    )
                })
            }
            OverlayKind::Stackable => {
                let render_handle = self.ui().render_state();
                let (target, path) = self
                    .events()
                    .controllers
                    .path_to::<StackableOverlayController>(origin.id(), &render_handle.get())?;
                clear_singles(&path, self);
                let overlay = overlay.add_strong(self).any();
                open_stackable_at(
                    target,
                    overlay,
                    options.cancel,
                    options.blocks_prior_input,
                    self,
                )
            }
        }
    }
}

impl<Rsc: HasEvents> OverlayRscExt for Rsc {}

fn modal_pointer_senses() -> CursorSenses {
    CursorSense::drag_senses()
        | CursorSense::PressStart(CursorButton::Right)
        | CursorSense::Pressing(CursorButton::Right)
        | CursorSense::PressEnd(CursorButton::Right)
        | CursorSense::PressStart(CursorButton::Middle)
        | CursorSense::Pressing(CursorButton::Middle)
        | CursorSense::PressEnd(CursorButton::Middle)
        | CursorSense::Scroll(Axis::X)
        | CursorSense::Scroll(Axis::Y)
}

fn clear_singles<Rsc: HasEvents>(path: &[WidgetId], rsc: &mut Rsc) {
    let controllers: Vec<ControllerId> = path
        .iter()
        .filter_map(|&host| rsc.events().controllers.id::<SingleOverlayController>(host))
        .collect();
    for id in controllers {
        rsc.with_controller(id, |single: &mut SingleOverlayController, rsc| {
            single.dismiss(id, rsc, true);
        });
    }
}

fn open_stackable_at<Rsc: HasEvents>(
    target: ControllerId,
    overlay: StrongWidget,
    cancel: CancelBehavior,
    blocks_prior_input: bool,
    rsc: &mut Rsc,
) -> Option<OverlayHandle> {
    rsc.with_controller(target, |stackable: &mut StackableOverlayController, rsc| {
        stackable.open(target, overlay, cancel, blocks_prior_input, rsc)
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::harness::{Harness, HarnessState, TouchAction};
    use std::{cell::Cell, rc::Rc};

    struct EscapeCounter(Rc<Cell<usize>>);

    impl Controller<StdRsc<HarnessState>> for EscapeCounter {
        fn command(&mut self, command: Command, _rsc: &mut StdRsc<HarnessState>) -> CommandResult {
            if command == Command::Escape {
                self.0.set(self.0.get() + 1);
                CommandResult::Used
            } else {
                CommandResult::Unused
            }
        }
    }

    fn stack_len(harness: &Harness, handle: OverlayHandle) -> usize {
        harness
            .rsc
            .widgets()
            .get_dyn(handle.host)
            .unwrap()
            .as_any()
            .downcast_ref::<Stack>()
            .unwrap()
            .children
            .len()
    }

    fn rooted_harness() -> (Harness, WeakWidget<Rect>, Rc<Cell<usize>>) {
        let mut harness = Harness::new(Vec2::new(200.0, 300.0), 1.0);
        let base = rect(PaintId::WHITE).add(&mut harness.rsc);
        let count = Rc::new(Cell::new(0));
        harness
            .rsc
            .register_controller(base, EscapeCounter(count.clone()));
        let target = harness
            .rsc
            .events()
            .controllers
            .id::<EscapeCounter>(base.id())
            .unwrap();
        harness.rsc.set_command_target(Some(target));
        let root = base.upgrade(&mut harness.rsc).any();
        harness.state.set_root(&mut harness.rsc, root);
        harness.frame(0);
        (harness, base, count)
    }

    #[test]
    fn popup_is_single_and_passes_escape_through() {
        let (mut harness, base, escapes) = rooted_harness();
        let first = harness.rsc.open_popup(base, rect(PaintId::RED)).unwrap();
        assert_eq!(stack_len(&harness, first), 2);

        let second = harness.rsc.open_popup(base, rect(PaintId::BLUE)).unwrap();
        assert_eq!(escapes.get(), 1, "replacement reaches the old owner");
        assert!(!first.close(&mut harness.rsc), "a replaced handle is stale");
        assert_eq!(stack_len(&harness, second), 2);

        assert_eq!(
            harness.rsc.run_command(Command::Escape),
            CommandResult::Used
        );
        assert_eq!(
            escapes.get(),
            1,
            "the replacement had no prior command target"
        );
        assert_eq!(stack_len(&harness, second), 1);
    }

    #[test]
    fn stackables_survive_singles_and_close_from_the_top() {
        let (mut harness, base, escapes) = rooted_harness();
        let options = OverlayOptions::stackable(CancelBehavior::Consume);
        let first = harness
            .rsc
            .open_overlay(base, rect(PaintId::RED), options)
            .unwrap();
        let single = harness.rsc.open_popup(base, rect(PaintId::GREEN)).unwrap();
        let second = harness
            .rsc
            .open_overlay(base, rect(PaintId::BLUE), options)
            .unwrap();

        assert!(!single.close(&mut harness.rsc));
        assert_eq!(stack_len(&harness, second), 3);
        assert_eq!(
            harness.rsc.run_command(Command::Escape),
            CommandResult::Used
        );
        assert_eq!(stack_len(&harness, first), 2);
        assert_eq!(
            harness.rsc.run_command(Command::Escape),
            CommandResult::Used
        );
        assert_eq!(stack_len(&harness, first), 1);
        assert_eq!(escapes.get(), 0, "stackable cancellation is consumed");
        assert_eq!(
            harness.rsc.run_command(Command::Escape),
            CommandResult::Used
        );
        assert_eq!(escapes.get(), 1);
    }

    #[test]
    fn an_explicit_single_can_consume_cancellation() {
        let (mut harness, base, escapes) = rooted_harness();
        let overlay = harness
            .rsc
            .open_overlay(
                base,
                rect(PaintId::RED),
                OverlayOptions::single(CancelBehavior::Consume),
            )
            .unwrap();

        assert_eq!(
            harness.rsc.run_command(Command::Escape),
            CommandResult::Used
        );
        assert_eq!(escapes.get(), 0);
        assert_eq!(stack_len(&harness, overlay), 1);
    }

    #[test]
    fn a_stackable_can_pass_cancellation_to_the_entry_under_it() {
        let (mut harness, base, escapes) = rooted_harness();
        let first = harness
            .rsc
            .open_overlay(
                base,
                rect(PaintId::RED),
                OverlayOptions::stackable(CancelBehavior::Consume),
            )
            .unwrap();
        harness
            .rsc
            .open_overlay(
                base,
                rect(PaintId::BLUE),
                OverlayOptions::stackable(CancelBehavior::PassThrough),
            )
            .unwrap();

        assert_eq!(
            harness.rsc.run_command(Command::Escape),
            CommandResult::Used
        );
        assert_eq!(stack_len(&harness, first), 1);
        assert_eq!(escapes.get(), 0);
    }

    #[test]
    fn closing_a_stackable_handle_closes_everything_above_it() {
        let (mut harness, base, _) = rooted_harness();
        let options = OverlayOptions::stackable(CancelBehavior::Consume);
        let first = harness
            .rsc
            .open_overlay(base, rect(PaintId::RED), options)
            .unwrap();
        let second = harness
            .rsc
            .open_overlay(base, rect(PaintId::BLUE), options)
            .unwrap();
        let single = harness.rsc.open_popup(base, rect(PaintId::GREEN)).unwrap();

        assert!(first.close(&mut harness.rsc));
        assert!(!second.close(&mut harness.rsc));
        assert!(!single.close(&mut harness.rsc));
        assert_eq!(stack_len(&harness, first), 1);
    }

    #[test]
    fn backdrop_dismisses_modal_without_clicking_the_base() {
        let mut harness = Harness::new(Vec2::new(200.0, 300.0), 1.0);
        let clicks = Rc::new(Cell::new(0));
        let base = rect(PaintId::WHITE).add(&mut harness.rsc);
        let counted = clicks.clone();
        harness
            .rsc
            .register_event(base, CursorSense::click(), move |_, _| {
                counted.set(counted.get() + 1);
            });
        let root = base.upgrade(&mut harness.rsc).any();
        harness.state.set_root(&mut harness.rsc, root);
        harness.frame(0);

        let modal = harness
            .rsc
            .open_modal(base, rect(PaintId::RED).sized((80, 80)))
            .unwrap();
        harness.frame(1);
        assert_eq!(stack_len(&harness, modal), 2);
        assert!(
            harness
                .rsc
                .events()
                .controllers
                .command_target_blocks_input()
        );
        harness.touch(TouchAction::Down, Vec2::new(100.0, 150.0), 2);
        harness.touch(TouchAction::Up, Vec2::new(100.0, 150.0), 3);
        assert_eq!(stack_len(&harness, modal), 2);
        harness.touch(TouchAction::Down, Vec2::new(5.0, 5.0), 4);

        assert_eq!(clicks.get(), 0);
        assert_eq!(stack_len(&harness, modal), 1);
        assert!(
            !harness
                .rsc
                .events()
                .controllers
                .command_target_blocks_input()
        );
    }

    #[test]
    fn local_single_host_passes_a_modal_to_the_root() {
        let mut harness = Harness::new(Vec2::new(200.0, 300.0), 1.0);
        let leaf = rect(PaintId::WHITE).add(&mut harness.rsc);
        let local = overlay_host(leaf, OverlayHostOptions::SINGLE).add(&mut harness.rsc);
        let root = local.upgrade(&mut harness.rsc).any();
        harness.state.set_root(&mut harness.rsc, root);
        harness.frame(0);

        let popup = harness.rsc.open_popup(leaf, rect(PaintId::RED)).unwrap();
        assert_eq!(popup.host, local.id());
        let modal = harness
            .rsc
            .open_modal(leaf, rect(PaintId::BLUE).sized((80, 80)))
            .unwrap();

        assert!(!popup.close(&mut harness.rsc));
        assert_ne!(modal.host, local.id());
    }

    #[test]
    fn replacing_a_local_single_does_not_cancel_an_ancestor_overlay() {
        let mut harness = Harness::new(Vec2::new(200.0, 300.0), 1.0);
        let leaf = rect(PaintId::WHITE).add(&mut harness.rsc);
        let local = overlay_host(leaf, OverlayHostOptions::SINGLE).add(&mut harness.rsc);
        let root = local.upgrade(&mut harness.rsc).any();
        harness.state.set_root(&mut harness.rsc, root);
        harness.frame(0);

        let first_modal = harness
            .rsc
            .open_modal(leaf, rect(PaintId::RED).sized((80, 80)))
            .unwrap();
        harness.frame(1);
        let popup = harness.rsc.open_popup(leaf, rect(PaintId::GREEN)).unwrap();
        let second_modal = harness
            .rsc
            .open_modal(leaf, rect(PaintId::BLUE).sized((80, 80)))
            .unwrap();

        assert!(!popup.close(&mut harness.rsc));
        assert_eq!(stack_len(&harness, first_modal), 3);
        assert!(second_modal.close(&mut harness.rsc));
        assert_eq!(stack_len(&harness, first_modal), 2);
    }

    #[test]
    fn inactive_origins_have_no_overlay_route() {
        let (mut harness, _, _) = rooted_harness();
        let inactive = rect(PaintId::WHITE).add(&mut harness.rsc);
        assert!(
            harness
                .rsc
                .open_popup(inactive, rect(PaintId::RED))
                .is_none()
        );
    }

    #[test]
    fn replacing_the_root_forgets_the_old_routing_tree() {
        let (mut harness, old, _) = rooted_harness();
        let new = rect(PaintId::WHITE).add(&mut harness.rsc);
        let root = new.upgrade(&mut harness.rsc).any();
        harness.state.set_root(&mut harness.rsc, root);
        harness.frame(1);

        assert!(harness.rsc.open_popup(old, rect(PaintId::RED)).is_none());
        assert!(harness.rsc.open_popup(new, rect(PaintId::RED)).is_some());
    }
}