//! xdg-session-management-v1 server side: sessions and named toplevels whose
//! window-management state is persisted by `Windows`' `SessionStore` and
//! reapplied on the first commit of a restored toplevel.

pub mod generated;

use std::collections::HashMap;

use generated::server::{
    xdg_session_manager_v1::{self, XdgSessionManagerV1},
    xdg_session_v1::{self, XdgSessionV1},
    xdg_toplevel_session_v1::{self, XdgToplevelSessionV1},
};
use smithay::reexports::wayland_protocols::xdg::shell::server::xdg_toplevel::XdgToplevel;
use smithay::reexports::wayland_server::backend::ClientId;
use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
use smithay::reexports::wayland_server::{
    Client, DataInit, Dispatch, DisplayHandle, GlobalDispatch, New, Resource, WEnum,
    Weak as WlWeak,
};
use smithay::wayland::shell::xdg::ToplevelSurface;
use tracing::debug;

use crate::session_store::{SavedToplevel, SessionStore};

const MANAGER_VERSION: u32 = 1;
const MAX_ID_LEN: usize = 1024;

pub(crate) fn valid_session_id(s: &str) -> bool {
    !s.is_empty()
        && s.len() <= MAX_ID_LEN
        && !s.chars().any(|c| c.is_control() || c == '/')
}

pub(crate) fn valid_toplevel_name(s: &str) -> bool {
    valid_session_id(s)
}

/// User data of an `xdg_session_v1`.
pub struct SessionData {
    pub session_id: String,
}

/// User data of an `xdg_toplevel_session_v1`.
pub struct ToplevelSessionData {
    pub session_id: String,
    /// `None` when the toplevel was unknown to forge; the object stays inert.
    pub surface: Option<WlSurface>,
}

/// Binding carried on a `Window`.
#[derive(Debug)]
pub struct WindowSession {
    pub session_id: String,
    pub name: String,
    pub handle: WlWeak<XdgToplevelSessionV1>,
    pub pending_restore: Option<SavedToplevel>,
}

#[derive(Debug)]
pub enum AttachError {
    NoWindow,
    AlreadyMapped,
}

struct ActiveSession {
    resource: WlWeak<XdgSessionV1>,
    client: ClientId,
    #[allow(dead_code)]
    reason: xdg_session_manager_v1::Reason,
    toplevels: HashMap<String, WlSurface>,
}

#[derive(Default)]
pub struct SessionManagerState {
    active: HashMap<String, ActiveSession>,
    bound: HashMap<WlSurface, String>,
}

impl SessionManagerState {
    pub fn new<D>(display: &DisplayHandle) -> Self
    where
        D: GlobalDispatch<XdgSessionManagerV1, ()>,
        D: Dispatch<XdgSessionManagerV1, ()>,
        D: Dispatch<XdgSessionV1, SessionData>,
        D: Dispatch<XdgToplevelSessionV1, ToplevelSessionData>,
        D: SessionManagementHandler,
        D: 'static,
    {
        display.create_global::<D, XdgSessionManagerV1, _>(MANAGER_VERSION, ());
        Self::default()
    }

    /// Free the name binding of a destroyed toplevel; persisted data stays.
    pub fn toplevel_destroyed(&mut self, surface: &WlSurface) {
        if let Some(sid) = self.bound.remove(surface) {
            if let Some(s) = self.active.get_mut(&sid) {
                s.toplevels.retain(|_, wl| wl != surface);
            }
        }
    }

    fn session_live(&self, sid: &str, resource: &XdgSessionV1) -> bool {
        self.active
            .get(sid)
            .and_then(|s| s.resource.upgrade().ok())
            .is_some_and(|r| &r == resource)
    }
}

pub trait SessionManagementHandler {
    fn session_management_state(&mut self) -> &mut SessionManagerState;
    fn session_store(&mut self) -> &mut SessionStore;
    fn session_toplevel(&mut self, toplevel: &XdgToplevel) -> Option<ToplevelSurface>;
    fn session_attach(
        &mut self,
        toplevel: &ToplevelSurface,
        ws: WindowSession,
    ) -> Result<(), AttachError>;
    fn session_detach(&mut self, surface: &WlSurface);
    fn session_detach_all(&mut self, session_id: &str);
    fn session_rename(&mut self, surface: &WlSurface, name: &str);
}

impl<D> GlobalDispatch<XdgSessionManagerV1, (), D> for SessionManagerState
where
    D: GlobalDispatch<XdgSessionManagerV1, ()>,
    D: Dispatch<XdgSessionManagerV1, ()>,
    D: Dispatch<XdgSessionV1, SessionData>,
    D: Dispatch<XdgToplevelSessionV1, ToplevelSessionData>,
    D: SessionManagementHandler,
    D: 'static,
{
    fn bind(
        _state: &mut D,
        _display: &DisplayHandle,
        _client: &Client,
        manager: New<XdgSessionManagerV1>,
        _data: &(),
        data_init: &mut DataInit<'_, D>,
    ) {
        data_init.init(manager, ());
        debug!("Session manager bound");
    }
}

impl<D> Dispatch<XdgSessionManagerV1, (), D> for SessionManagerState
where
    D: Dispatch<XdgSessionManagerV1, ()>,
    D: Dispatch<XdgSessionV1, SessionData>,
    D: Dispatch<XdgToplevelSessionV1, ToplevelSessionData>,
    D: SessionManagementHandler,
    D: 'static,
{
    fn request(
        state: &mut D,
        client: &Client,
        manager: &XdgSessionManagerV1,
        request: xdg_session_manager_v1::Request,
        _data: &(),
        _display: &DisplayHandle,
        data_init: &mut DataInit<'_, D>,
    ) {
        match request {
            xdg_session_manager_v1::Request::GetSession { id, reason, session_id } => {
                let reason = match reason {
                    WEnum::Value(r) => r,
                    WEnum::Unknown(_) => {
                        data_init.init(id, SessionData { session_id: String::new() });
                        manager.post_error(
                            xdg_session_manager_v1::Error::InvalidReason,
                            "invalid reason",
                        );
                        return;
                    }
                };
                if let Some(requested) = session_id.as_deref() {
                    if !valid_session_id(requested) {
                        data_init.init(id, SessionData { session_id: String::new() });
                        manager.post_error(
                            xdg_session_manager_v1::Error::InvalidSessionId,
                            "invalid session id",
                        );
                        return;
                    }
                }

                let restoring = session_id
                    .as_deref()
                    .is_some_and(|sid| state.session_store().has_session(sid));
                let sid = if restoring {
                    session_id.unwrap()
                } else {
                    state.session_store().generate_id()
                };

                let previous = state
                    .session_management_state()
                    .active
                    .get(&sid)
                    .map(|s| (s.client.clone(), s.resource.clone()));
                if let Some((owner, old)) = previous {
                    if owner == client.id() {
                        data_init.init(id, SessionData { session_id: String::new() });
                        manager.post_error(
                            xdg_session_manager_v1::Error::InUse,
                            "session already in use by this client",
                        );
                        return;
                    }
                    if let Ok(old) = old.upgrade() {
                        old.replaced();
                    }
                    state.session_detach_all(&sid);
                    let sm = state.session_management_state();
                    sm.active.remove(&sid);
                    sm.bound.retain(|_, s| s != &sid);
                }

                let session = data_init.init(id, SessionData { session_id: sid.clone() });
                state.session_management_state().active.insert(
                    sid.clone(),
                    ActiveSession {
                        resource: session.downgrade(),
                        client: client.id(),
                        reason,
                        toplevels: HashMap::new(),
                    },
                );
                state.session_store().touch(&sid);
                if restoring {
                    session.restored();
                } else {
                    session.created(sid);
                }
            }
            xdg_session_manager_v1::Request::Destroy => {}
        }
    }
}

fn bind_toplevel<D>(
    state: &mut D,
    session: &XdgSessionV1,
    sid: &str,
    id: New<XdgToplevelSessionV1>,
    toplevel: &XdgToplevel,
    name: String,
    restore: bool,
    data_init: &mut DataInit<'_, D>,
) where
    D: Dispatch<XdgToplevelSessionV1, ToplevelSessionData>,
    D: SessionManagementHandler,
    D: 'static,
{
    let Some(surface_tl) = state.session_toplevel(toplevel) else {
        data_init.init(id, ToplevelSessionData { session_id: sid.to_string(), surface: None });
        return;
    };
    let wl = surface_tl.wl_surface().clone();
    let handle = data_init.init(
        id,
        ToplevelSessionData { session_id: sid.to_string(), surface: Some(wl.clone()) },
    );

    if !state.session_management_state().session_live(sid, session) {
        return;
    }
    if !valid_toplevel_name(&name) {
        session.post_error(xdg_session_v1::Error::InvalidName, "invalid toplevel name");
        return;
    }
    {
        let sm = state.session_management_state();
        if sm.bound.contains_key(&wl) {
            session.post_error(xdg_session_v1::Error::AlreadyAdded, "toplevel already added");
            return;
        }
        if sm.active.get(sid).is_some_and(|s| s.toplevels.contains_key(&name)) {
            session.post_error(xdg_session_v1::Error::NameInUse, "toplevel name in use");
            return;
        }
    }

    let pending_restore = if restore { state.session_store().lookup(sid, &name) } else { None };
    let ws = WindowSession {
        session_id: sid.to_string(),
        name: name.clone(),
        handle: handle.downgrade(),
        pending_restore,
    };
    match state.session_attach(&surface_tl, ws) {
        Ok(()) => {
            let sm = state.session_management_state();
            sm.bound.insert(wl.clone(), sid.to_string());
            if let Some(s) = sm.active.get_mut(sid) {
                s.toplevels.insert(name, wl);
            }
        }
        Err(AttachError::AlreadyMapped) => {
            session.post_error(xdg_session_v1::Error::AlreadyMapped, "toplevel already mapped");
        }
        Err(AttachError::NoWindow) => {}
    }
}

impl<D> Dispatch<XdgSessionV1, SessionData, D> for SessionManagerState
where
    D: Dispatch<XdgSessionV1, SessionData>,
    D: Dispatch<XdgToplevelSessionV1, ToplevelSessionData>,
    D: SessionManagementHandler,
    D: 'static,
{
    fn request(
        state: &mut D,
        _client: &Client,
        session: &XdgSessionV1,
        request: xdg_session_v1::Request,
        data: &SessionData,
        _display: &DisplayHandle,
        data_init: &mut DataInit<'_, D>,
    ) {
        let sid = data.session_id.as_str();
        match request {
            xdg_session_v1::Request::AddToplevel { id, toplevel, name } => {
                bind_toplevel(state, session, sid, id, &toplevel, name, false, data_init);
            }
            xdg_session_v1::Request::RestoreToplevel { id, toplevel, name } => {
                bind_toplevel(state, session, sid, id, &toplevel, name, true, data_init);
            }
            xdg_session_v1::Request::RemoveToplevel { name } => {
                if !state.session_management_state().session_live(sid, session) {
                    return;
                }
                state.session_store().remove_toplevel(sid, &name);
                let live = state
                    .session_management_state()
                    .active
                    .get_mut(sid)
                    .and_then(|s| s.toplevels.remove(&name));
                if let Some(wl) = live {
                    state.session_management_state().bound.remove(&wl);
                    state.session_detach(&wl);
                }
            }
            xdg_session_v1::Request::Remove => {
                if state.session_management_state().session_live(sid, session) {
                    state.session_store().remove_session(sid);
                }
            }
            xdg_session_v1::Request::Destroy => {}
        }
    }

    fn destroyed(
        state: &mut D,
        _client: ClientId,
        session: &XdgSessionV1,
        data: &SessionData,
    ) {
        let sid = data.session_id.as_str();
        if sid.is_empty() {
            return;
        }
        let ours = state
            .session_management_state()
            .active
            .get(sid)
            .is_some_and(|s| s.resource.id() == session.id());
        if !ours {
            return;
        }
        state.session_detach_all(sid);
        let sm = state.session_management_state();
        sm.active.remove(sid);
        sm.bound.retain(|_, s| s != sid);
    }
}

impl<D> Dispatch<XdgToplevelSessionV1, ToplevelSessionData, D> for SessionManagerState
where
    D: Dispatch<XdgToplevelSessionV1, ToplevelSessionData>,
    D: SessionManagementHandler,
    D: 'static,
{
    fn request(
        state: &mut D,
        _client: &Client,
        _handle: &XdgToplevelSessionV1,
        request: xdg_toplevel_session_v1::Request,
        data: &ToplevelSessionData,
        _display: &DisplayHandle,
        _data_init: &mut DataInit<'_, D>,
    ) {
        match request {
            xdg_toplevel_session_v1::Request::Rename { name } => {
                let sid = data.session_id.as_str();
                let Some(surface) = data.surface.as_ref() else { return };
                let (session, old_name) = {
                    let sm = state.session_management_state();
                    let Some(active) = sm.active.get(sid) else { return };
                    let Ok(session) = active.resource.upgrade() else { return };
                    let old = active
                        .toplevels
                        .iter()
                        .find(|(_, wl)| *wl == surface)
                        .map(|(n, _)| n.clone());
                    (session, old)
                };
                let Some(old_name) = old_name else { return };
                if !valid_toplevel_name(&name) {
                    session.post_error(xdg_session_v1::Error::InvalidName, "invalid toplevel name");
                    return;
                }
                if old_name == name {
                    return;
                }
                {
                    let sm = state.session_management_state();
                    let Some(active) = sm.active.get_mut(sid) else { return };
                    if active.toplevels.contains_key(&name) {
                        session.post_error(xdg_session_v1::Error::NameInUse, "toplevel name in use");
                        return;
                    }
                    let wl = active.toplevels.remove(&old_name).unwrap();
                    active.toplevels.insert(name.clone(), wl);
                }
                state.session_store().rename_toplevel(sid, &old_name, &name);
                state.session_rename(surface, &name);
            }
            xdg_toplevel_session_v1::Request::Destroy => {}
        }
    }
}

#[macro_export]
macro_rules! delegate_session_management {
    ($(@<$( $lt:tt $( : $clt:tt $(+ $dlt:tt )* )? ),+>)? $ty: ty) => {
        smithay::reexports::wayland_server::delegate_global_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::protocols::session_management::generated::server::xdg_session_manager_v1::XdgSessionManagerV1: ()
        ] => $crate::protocols::session_management::SessionManagerState);

        smithay::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::protocols::session_management::generated::server::xdg_session_manager_v1::XdgSessionManagerV1: ()
        ] => $crate::protocols::session_management::SessionManagerState);

        smithay::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::protocols::session_management::generated::server::xdg_session_v1::XdgSessionV1: $crate::protocols::session_management::SessionData
        ] => $crate::protocols::session_management::SessionManagerState);

        smithay::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::protocols::session_management::generated::server::xdg_toplevel_session_v1::XdgToplevelSessionV1: $crate::protocols::session_management::ToplevelSessionData
        ] => $crate::protocols::session_management::SessionManagerState);
    };
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn ids_and_names_are_validated() {
        assert!(valid_session_id("abc123"));
        assert!(valid_session_id("ünïcode-ok"));
        assert!(!valid_session_id(""));
        assert!(!valid_session_id("a/b"));
        assert!(!valid_session_id("a\nb"));
        assert!(!valid_session_id(&"x".repeat(MAX_ID_LEN + 1)));
        assert!(valid_toplevel_name("main"));
        assert!(!valid_toplevel_name("\u{7}"));
    }
}
