use gpui::{
    AnyElement, App, Div, ElementId, Entity, FocusHandle, InteractiveElement, IntoElement,
    ParentElement, RenderOnce, StyleRefinement, Styled, Subscription, Window, actions, div,
};

use crate::theme::ActiveTheme;

actions!(ely, [FocusNext, FocusPrev]);

/// Focus-colored border while focused. Give the element a 1px border.
pub trait FocusRing: InteractiveElement + Sized {
    fn focus_ring(self, cx: &App) -> Self {
        let ring = cx.theme().colors.focus;
        self.focus(move |style| style.border_color(ring))
    }
}

impl<E: InteractiveElement> FocusRing for E {}

/// Owns Tab and Shift-Tab for its subtree. `trap` keeps focus inside.
#[derive(IntoElement)]
pub struct FocusScope {
    base: Div,
    handle: FocusHandle,
    trap: bool,
    root: bool,
}

impl FocusScope {
    pub fn new(handle: &FocusHandle) -> Self {
        Self {
            base: div().track_focus(handle),
            handle: handle.clone(),
            trap: false,
            root: false,
        }
    }

    pub fn trap(mut self) -> Self {
        self.trap = true;
        self
    }

    /// The window's outermost scope: it takes focus when the focused element leaves the tree.
    pub fn root(mut self) -> Self {
        self.root = true;
        self
    }
}

/// Keeps a root scope's claim on lost focus while it renders.
struct Root {
    _lost: Subscription,
}

impl Styled for FocusScope {
    fn style(&mut self) -> &mut StyleRefinement {
        self.base.style()
    }
}

impl ParentElement for FocusScope {
    fn extend(&mut self, elements: impl IntoIterator<Item = AnyElement>) {
        self.base.extend(elements);
    }
}

impl RenderOnce for FocusScope {
    fn render(self, window: &mut Window, cx: &mut App) -> impl IntoElement {
        if self.root {
            let handle = self.handle.clone();
            window.use_keyed_state("focus-root", cx, |window, cx| Root {
                _lost: cx.on_focus_lost(window, move |_: &mut Root, window, _| {
                    log::info!("focus: its element left the tree; the root takes it");
                    window.focus(&handle);
                }),
            });
        }
        let (next, prev) = (self.handle.clone(), self.handle);
        let trap = self.trap;
        self.base
            .on_action(move |_: &FocusNext, window, cx| step(&next, trap, true, window, cx))
            .on_action(move |_: &FocusPrev, window, cx| step(&prev, trap, false, window, cx))
    }
}

fn step(scope: &FocusHandle, trap: bool, forward: bool, window: &mut Window, cx: &mut App) {
    let advance = |window: &mut Window| {
        if forward {
            window.focus_next()
        } else {
            window.focus_prev()
        }
    };
    if !trap {
        advance(window);
        return;
    }
    let origin = window.focused(cx);
    let mut first = None;
    loop {
        advance(window);
        let focused = window.focused(cx);
        if scope.contains_focused(window, cx) && focused.as_ref() != Some(scope) {
            return;
        }
        if focused.is_none() || focused == origin || (first.is_some() && focused == first) {
            break;
        }
        if first.is_none() {
            first = focused;
        }
    }
    log::warn!("focus scope: trap holds no tab stop; focus stays");
    if let Some(origin) = origin {
        window.focus(&origin);
    }
}

/// Focus an overlay holds while open, and what it held before.
pub(crate) struct Takeover {
    pub focus: FocusHandle,
    previous: Option<FocusHandle>,
    taken: bool,
    returned: bool,
    _lost: Subscription,
}

/// Keyed focus for an overlay; focused on its first render, and again whenever the focused element leaves the tree.
pub(crate) fn take_focus(
    key: impl Into<gpui::ElementId>,
    window: &mut Window,
    cx: &mut App,
) -> Entity<Takeover> {
    let state = window.use_keyed_state(key, cx, |window, cx| Takeover {
        focus: cx.focus_handle(),
        previous: None,
        taken: false,
        returned: false,
        _lost: cx.on_focus_lost(window, |takeover: &mut Takeover, window, _| {
            if !takeover.returned {
                log::info!("focus: its element left the tree; the overlay takes it back");
                window.focus(&takeover.focus);
            }
        }),
    });
    if !state.read(cx).taken {
        let previous = window.focused(cx);
        window.focus(&state.read(cx).focus.clone());
        state.update(cx, |takeover, _| {
            takeover.previous = previous;
            takeover.taken = true;
        });
    }
    state
}

/// Returns focus to whatever held it before the overlay opened.
pub(crate) fn give_back(state: &Entity<Takeover>, window: &mut Window, cx: &mut App) {
    let previous = state.update(cx, |takeover, _| {
        takeover.returned = true;
        takeover.previous.clone()
    });
    match previous {
        Some(previous) => {
            window.focus(&previous);
            log::info!("focus: handed back after an overlay");
        }
        None => log::info!("focus: overlay closed, nothing was focused before"),
    }
}

/// A focus handle kept for `id`, a Tab stop while enabled.
pub(crate) fn tab_stop(
    id: ElementId,
    enabled: bool,
    window: &mut Window,
    cx: &mut App,
) -> FocusHandle {
    window
        .use_keyed_state(id, cx, |_, cx| cx.focus_handle())
        .read(cx)
        .clone()
        .tab_stop(enabled)
}