use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};

use spacesh_proto::ids::{GroupId, SurfaceId, WorkspaceId};
use spacesh_proto::status::SurfaceState;
use spacesh_proto::workspace::{Group, SurfaceSpec, SurfaceView, Workspace, WorkspaceView};

use crate::state_store::PersistState;
use crate::surface::SurfaceHandle;

/// Single-threaded owner of structure (workspaces/groups/trees + per-surface
/// specs) and the live actor map. Lives in the server router task.
#[derive(Default)]
pub struct Registry {
    counter: AtomicU64,
    groups: HashMap<GroupId, Group>,
    workspaces: HashMap<WorkspaceId, Workspace>,
    by_path: HashMap<String, WorkspaceId>,
    /// Live actors only. Absent id that exists in a workspace's `surfaces` = stopped.
    live: HashMap<SurfaceId, SurfaceHandle>,
    /// Ephemeral per-surface status. In-memory only (never persisted).
    states: HashMap<SurfaceId, SurfaceState>,
}

impl Registry {
    pub fn new() -> Self {
        Self::default()
    }

    fn next_id(&self, prefix: &str) -> String {
        let n = self.counter.fetch_add(1, Ordering::Relaxed);
        format!("{prefix}_{n:x}")
    }

    pub fn new_surface_id(&self) -> SurfaceId {
        SurfaceId(self.next_id("s"))
    }

    // ---- workspaces ----

    /// Idempotent by canonicalized path. Returns (workspace_id, created?).
    pub fn open_workspace(&mut self, path: PathBuf) -> (WorkspaceId, bool) {
        let canonical = path.canonicalize().unwrap_or(path);
        let key = canonical.to_string_lossy().to_string();
        if let Some(id) = self.by_path.get(&key) {
            return (id.clone(), false);
        }
        let id = WorkspaceId(self.next_id("w"));
        let name = canonical.file_name().map(|s| s.to_string_lossy().to_string()).unwrap_or_else(|| key.clone());
        let order = self.workspaces.len() as u32;
        self.workspaces.insert(id.clone(), Workspace {
            id: id.clone(), path: key.clone(), name, group_id: None, order,
            unread: false, layout: None, surfaces: HashMap::new(),
        });
        self.by_path.insert(key, id.clone());
        (id, true)
    }

    pub fn workspace(&self, id: &WorkspaceId) -> Option<&Workspace> {
        self.workspaces.get(id)
    }
    pub fn workspace_mut(&mut self, id: &WorkspaceId) -> Option<&mut Workspace> {
        self.workspaces.get_mut(id)
    }
    pub fn close_workspace(&mut self, id: &WorkspaceId) -> Vec<SurfaceId> {
        let Some(ws) = self.workspaces.remove(id) else { return vec![] };
        self.by_path.retain(|_, v| v != id);
        let ids: Vec<SurfaceId> = ws.surfaces.keys().cloned().collect();
        for sid in &ids {
            self.live.remove(sid);
        }
        ids
    }

    /// The workspace that owns a surface id, if any.
    pub fn workspace_of(&self, sid: &SurfaceId) -> Option<WorkspaceId> {
        self.workspaces.values().find(|w| w.surfaces.contains_key(sid)).map(|w| w.id.clone())
    }

    // ---- surfaces (structure) ----

    pub fn add_surface_spec(&mut self, ws: &WorkspaceId, sid: SurfaceId, spec: SurfaceSpec) {
        if let Some(w) = self.workspaces.get_mut(ws) {
            w.surfaces.insert(sid, spec);
        }
    }
    pub fn surface_spec(&self, sid: &SurfaceId) -> Option<SurfaceSpec> {
        self.workspaces.values().find_map(|w| w.surfaces.get(sid).cloned())
    }
    /// Remove a surface from its workspace's spec map and the tree.
    pub fn remove_surface(&mut self, sid: &SurfaceId) {
        self.live.remove(sid);
        if let Some(ws) = self.workspace_of(sid) {
            if let Some(w) = self.workspaces.get_mut(&ws) {
                w.surfaces.remove(sid);
                w.layout = w.layout.take().and_then(|l| spacesh_core::ops::remove_leaf(l, sid));
            }
        }
    }

    // ---- live actors ----

    pub fn set_live(&mut self, handle: SurfaceHandle) {
        self.live.insert(handle.id.clone(), handle);
    }
    pub fn live(&self, sid: &SurfaceId) -> Option<&SurfaceHandle> {
        self.live.get(sid)
    }
    pub fn mark_stopped(&mut self, sid: &SurfaceId) {
        self.live.remove(sid);
    }
    pub fn is_running(&self, sid: &SurfaceId) -> bool {
        self.live.contains_key(sid)
    }

    // ---- surface state ----

    pub fn set_state(&mut self, sid: &SurfaceId, state: SurfaceState) {
        self.states.insert(sid.clone(), state);
    }
    pub fn state(&self, sid: &SurfaceId) -> SurfaceState {
        self.states.get(sid).copied().unwrap_or_default()
    }
    pub fn drop_state(&mut self, sid: &SurfaceId) {
        self.states.remove(sid);
    }

    // ---- groups ----

    pub fn create_group(&mut self, name: String, color: String) -> GroupId {
        let id = GroupId(self.next_id("g"));
        let order = self.groups.len() as u32;
        self.groups.insert(id.clone(), Group { id: id.clone(), name, color, order });
        id
    }
    pub fn group_mut(&mut self, id: &GroupId) -> Option<&mut Group> {
        self.groups.get_mut(id)
    }
    pub fn delete_group(&mut self, id: &GroupId) {
        self.groups.remove(id);
        for w in self.workspaces.values_mut() {
            if w.group_id.as_ref() == Some(id) {
                w.group_id = None;
            }
        }
    }
    pub fn groups(&self) -> Vec<Group> {
        let mut g: Vec<Group> = self.groups.values().cloned().collect();
        g.sort_by_key(|x| x.order);
        g
    }

    // ---- views & persistence ----

    pub fn workspace_view(&self, id: &WorkspaceId) -> Option<WorkspaceView> {
        let w = self.workspaces.get(id)?;
        Some(self.to_view(w))
    }
    fn to_view(&self, w: &Workspace) -> WorkspaceView {
        let surfaces = w.surfaces.iter().map(|(sid, spec)| {
            (sid.clone(), SurfaceView {
                spec: spec.clone(),
                running: self.live.contains_key(sid),
                state: self.state(sid),
            })
        }).collect();
        WorkspaceView {
            id: w.id.clone(), path: w.path.clone(), name: w.name.clone(),
            group_id: w.group_id.clone(), order: w.order, unread: w.unread,
            layout: w.layout.clone(), surfaces,
        }
    }
    pub fn status(&self) -> (Vec<Group>, Vec<WorkspaceView>) {
        let mut ws: Vec<WorkspaceView> = self.workspaces.values().map(|w| self.to_view(w)).collect();
        ws.sort_by_key(|w| w.order);
        (self.groups(), ws)
    }
    pub fn persist_state(&self) -> PersistState {
        let mut workspaces: Vec<Workspace> = self.workspaces.values().cloned().collect();
        workspaces.sort_by_key(|w| w.order);
        PersistState { version: 1, groups: self.groups(), workspaces }
    }
    /// Replace all structure from a loaded snapshot (cold start). Clears live map.
    pub fn restore(&mut self, state: PersistState) {
        self.groups = state.groups.into_iter().map(|g| (g.id.clone(), g)).collect();
        self.workspaces.clear();
        self.by_path.clear();
        self.live.clear();
        self.states.clear();
        for w in state.workspaces {
            self.by_path.insert(w.path.clone(), w.id.clone());
            self.workspaces.insert(w.id.clone(), w);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use spacesh_proto::layout::{LayoutNode as LN, Orient};

    fn spec() -> SurfaceSpec {
        SurfaceSpec { command: "/bin/sh".into(), args: vec![], cwd: "/tmp".into(),
            agent_label: None, cols: 80, rows: 24, autostart: false }
    }

    #[test]
    fn open_is_idempotent() {
        let mut r = Registry::new();
        let (a, c1) = r.open_workspace(std::env::temp_dir());
        let (b, c2) = r.open_workspace(std::env::temp_dir());
        assert_eq!(a, b);
        assert!(c1 && !c2);
    }

    #[test]
    fn surface_running_then_stopped() {
        let mut r = Registry::new();
        let (ws, _) = r.open_workspace(std::env::temp_dir());
        let sid = r.new_surface_id();
        r.add_surface_spec(&ws, sid.clone(), spec());
        assert!(!r.is_running(&sid)); // spec present, no live actor = stopped
        let v = r.workspace_view(&ws).unwrap();
        assert_eq!(v.surfaces.get(&sid).unwrap().running, false);
    }

    #[test]
    fn remove_surface_updates_tree() {
        let mut r = Registry::new();
        let (ws, _) = r.open_workspace(std::env::temp_dir());
        let s1 = r.new_surface_id();
        let s2 = r.new_surface_id();
        r.add_surface_spec(&ws, s1.clone(), spec());
        r.add_surface_spec(&ws, s2.clone(), spec());
        r.workspace_mut(&ws).unwrap().layout = Some(LN::Split {
            orient: Orient::H, ratios: vec![0.5, 0.5],
            children: vec![LN::leaf(s1.clone()), LN::leaf(s2.clone())],
        });
        r.remove_surface(&s2);
        let w = r.workspace(&ws).unwrap();
        assert!(!w.surfaces.contains_key(&s2));
        assert_eq!(w.layout, Some(LN::leaf(s1))); // split collapsed
    }

    #[test]
    fn restore_round_trips_through_persist_state() {
        let mut r = Registry::new();
        let (ws, _) = r.open_workspace(std::env::temp_dir());
        r.add_surface_spec(&ws, r.new_surface_id(), spec());
        let state = r.persist_state();
        let mut r2 = Registry::new();
        r2.restore(state.clone());
        assert_eq!(r2.persist_state().workspaces.len(), state.workspaces.len());
    }

    #[test]
    fn delete_group_ungroups_members() {
        let mut r = Registry::new();
        let (ws, _) = r.open_workspace(std::env::temp_dir());
        let g = r.create_group("prod".into(), "#fff".into());
        r.workspace_mut(&ws).unwrap().group_id = Some(g.clone());
        r.delete_group(&g);
        assert!(r.workspace(&ws).unwrap().group_id.is_none());
    }

    #[test]
    fn state_defaults_idle_and_can_be_set() {
        let mut r = Registry::new();
        let (ws, _) = r.open_workspace(std::env::temp_dir());
        let sid = r.new_surface_id();
        r.add_surface_spec(&ws, sid.clone(), spec());
        assert_eq!(r.state(&sid), spacesh_proto::status::SurfaceState::Idle);
        r.set_state(&sid, spacesh_proto::status::SurfaceState::Work);
        assert_eq!(r.state(&sid), spacesh_proto::status::SurfaceState::Work);
        let v = r.workspace_view(&ws).unwrap();
        assert_eq!(v.surfaces.get(&sid).unwrap().state, spacesh_proto::status::SurfaceState::Work);
    }

    #[test]
    fn drop_state_resets_to_idle() {
        let mut r = Registry::new();
        let (ws, _) = r.open_workspace(std::env::temp_dir());
        let sid = r.new_surface_id();
        r.add_surface_spec(&ws, sid.clone(), spec());
        r.set_state(&sid, spacesh_proto::status::SurfaceState::Error);
        r.drop_state(&sid);
        assert_eq!(r.state(&sid), spacesh_proto::status::SurfaceState::Idle);
    }
}