// Per-session window-management state for xdg-session-management-v1.
//
// On-disk format, `$XDG_STATE_HOME/forge/sessions.json`:
//
//     { "<session_id>": { "last_used": 1723700000,
//                         "toplevels": { "<name>": { "x": .., "y": .., "w": .., "h": ..,
//                                                    "maximized": false, "output": "DP-2",
//                                                    "stack": 3 } } } }

use std::collections::{BTreeMap, HashMap};
use std::io::Read;
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH};

use serde::{Deserialize, Serialize};
use smithay::utils::{Logical, Rectangle};
use tracing::warn;

use crate::geometry_store::{write_json_atomic, OutputRef, SavedGeo};

const RETENTION_DAYS: u64 = 60;
const MAX_SESSIONS: usize = 256;

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SavedToplevel {
    #[serde(flatten)]
    pub geo: SavedGeo,
    #[serde(default)]
    pub stack: u32,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct SavedSession {
    pub last_used: u64,
    #[serde(default)]
    pub toplevels: BTreeMap<String, SavedToplevel>,
}

#[derive(Debug)]
pub struct SessionStore {
    path: PathBuf,
    sessions: HashMap<String, SavedSession>,
}

fn now_secs() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

impl SessionStore {
    pub fn load() -> Self {
        Self::with_path(Self::default_path())
    }

    /// Load from `path`; a missing or corrupt file yields an empty store.
    /// Sessions unused for `RETENTION_DAYS` are dropped, then the oldest
    /// beyond `MAX_SESSIONS`.
    pub fn with_path(path: PathBuf) -> Self {
        let sessions: HashMap<String, SavedSession> = match std::fs::read_to_string(&path) {
            Ok(s) => serde_json::from_str(&s).unwrap_or_else(|e| {
                warn!("session_store: ignoring corrupt {}: {e}", path.display());
                HashMap::new()
            }),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => HashMap::new(),
            Err(e) => {
                warn!("session_store: read {}: {e}", path.display());
                HashMap::new()
            }
        };
        let mut store = Self { path, sessions };
        if store.evict(now_secs()) {
            store.persist();
        }
        store
    }

    fn default_path() -> PathBuf {
        dirs::state_dir()
            .unwrap_or_else(|| dirs::data_dir().unwrap_or_default())
            .join("forge")
            .join("sessions.json")
    }

    fn evict(&mut self, now: u64) -> bool {
        let before = self.sessions.len();
        let cutoff = now.saturating_sub(RETENTION_DAYS * 86_400);
        self.sessions.retain(|_, s| s.last_used >= cutoff);
        if self.sessions.len() > MAX_SESSIONS {
            let mut by_age: Vec<(String, u64)> = self
                .sessions
                .iter()
                .map(|(k, s)| (k.clone(), s.last_used))
                .collect();
            by_age.sort_by_key(|(_, t)| *t);
            for (k, _) in by_age.into_iter().take(self.sessions.len() - MAX_SESSIONS) {
                self.sessions.remove(&k);
            }
        }
        self.sessions.len() != before
    }

    pub fn has_session(&self, id: &str) -> bool {
        self.sessions.contains_key(id)
    }

    /// Create the session if needed and bump its `last_used`.
    pub fn touch(&mut self, id: &str) {
        let now = now_secs();
        let s = self.sessions.entry(id.to_string()).or_default();
        s.last_used = now;
        self.persist();
    }

    pub fn lookup(&self, id: &str, name: &str) -> Option<SavedToplevel> {
        self.sessions
            .get(id)?
            .toplevels
            .get(name)
            .cloned()
            .filter(|t| t.geo.is_usable())
    }

    pub fn record(
        &mut self,
        id: &str,
        name: &str,
        rect: Rectangle<i32, Logical>,
        maximized: bool,
        output: OutputRef,
        stack: u32,
    ) {
        let geo = SavedGeo::from_rect(rect, maximized, output);
        if !geo.is_usable() {
            return;
        }
        let entry = SavedToplevel { geo, stack };
        let session = self.sessions.entry(id.to_string()).or_default();
        session.last_used = now_secs();
        let previous = session.toplevels.insert(name.to_string(), entry.clone());
        if previous.as_ref() == Some(&entry) {
            return;
        }
        self.persist();
    }

    pub fn remove_toplevel(&mut self, id: &str, name: &str) {
        let removed = self
            .sessions
            .get_mut(id)
            .map(|s| s.toplevels.remove(name).is_some())
            .unwrap_or(false);
        if removed {
            self.persist();
        }
    }

    pub fn remove_session(&mut self, id: &str) {
        if self.sessions.remove(id).is_some() {
            self.persist();
        }
    }

    pub fn rename_toplevel(&mut self, id: &str, from: &str, to: &str) {
        let Some(s) = self.sessions.get_mut(id) else { return };
        if let Some(entry) = s.toplevels.remove(from) {
            s.toplevels.insert(to.to_string(), entry);
            self.persist();
        }
    }

    /// A fresh 32-hex-char id not present in the store.
    pub fn generate_id(&self) -> String {
        loop {
            let id = random_hex_128();
            if !self.sessions.contains_key(&id) {
                return id;
            }
        }
    }

    fn persist(&self) {
        write_json_atomic(&self.path, &self.sessions, "session_store");
    }
}

fn random_hex_128() -> String {
    let mut buf = [0u8; 16];
    let read_ok = std::fs::File::open("/dev/urandom")
        .and_then(|mut f| f.read_exact(&mut buf))
        .is_ok();
    if !read_ok {
        use std::collections::hash_map::RandomState;
        use std::hash::{BuildHasher, Hash, Hasher};
        static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
        let seed = (
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .map(|d| d.as_nanos())
                .unwrap_or(0),
            std::process::id(),
            COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
        );
        for (i, chunk) in buf.chunks_mut(8).enumerate() {
            let mut h = RandomState::new().build_hasher();
            (seed, i).hash(&mut h);
            chunk.copy_from_slice(&h.finish().to_le_bytes());
        }
    }
    buf.iter().map(|b| format!("{b:02x}")).collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use smithay::utils::{Point, Size};

    fn rect(x: i32, y: i32, w: i32, h: i32) -> Rectangle<i32, Logical> {
        Rectangle::new(Point::from((x, y)), Size::from((w, h)))
    }

    fn temp_store() -> (tempfile::TempDir, SessionStore) {
        let dir = tempfile::tempdir().unwrap();
        let store = SessionStore::with_path(dir.path().join("sessions.json"));
        (dir, store)
    }

    #[test]
    fn touch_creates_session() {
        let (_d, mut s) = temp_store();
        assert!(!s.has_session("a"));
        s.touch("a");
        assert!(s.has_session("a"));
        assert!(s.path.exists());
    }

    #[test]
    fn record_and_lookup_round_trip() {
        let (_d, mut s) = temp_store();
        s.record("a", "main", rect(10, 20, 640, 480), true, OutputRef { connector: Some("DP-2".into()), identity: None }, 3);
        let t = s.lookup("a", "main").unwrap();
        assert_eq!((t.geo.x, t.geo.y, t.geo.w, t.geo.h), (10, 20, 640, 480));
        assert!(t.geo.maximized);
        assert_eq!(t.geo.output.as_deref(), Some("DP-2"));
        assert_eq!(t.stack, 3);

        let reloaded = SessionStore::with_path(s.path.clone());
        assert_eq!(reloaded.lookup("a", "main"), Some(t));
    }

    #[test]
    fn identical_record_does_not_rewrite() {
        let (_d, mut s) = temp_store();
        s.record("a", "main", rect(0, 0, 640, 480), false, OutputRef::default(), 0);
        let m1 = std::fs::metadata(&s.path).unwrap().modified().unwrap();
        std::thread::sleep(std::time::Duration::from_millis(20));
        s.record("a", "main", rect(0, 0, 640, 480), false, OutputRef::default(), 0);
        let m2 = std::fs::metadata(&s.path).unwrap().modified().unwrap();
        assert_eq!(m1, m2);
    }

    #[test]
    fn unusable_rect_is_not_stored() {
        let (_d, mut s) = temp_store();
        s.record("a", "main", rect(0, 0, 10, 10), false, OutputRef::default(), 0);
        assert!(s.lookup("a", "main").is_none());
    }

    #[test]
    fn remove_and_rename() {
        let (_d, mut s) = temp_store();
        s.record("a", "one", rect(0, 0, 640, 480), false, OutputRef::default(), 0);
        s.record("a", "two", rect(0, 0, 640, 480), false, OutputRef::default(), 0);
        s.rename_toplevel("a", "one", "uno");
        assert!(s.lookup("a", "one").is_none());
        assert!(s.lookup("a", "uno").is_some());
        s.remove_toplevel("a", "uno");
        assert!(s.lookup("a", "uno").is_none());
        assert!(s.lookup("a", "two").is_some());
        s.remove_session("a");
        assert!(!s.has_session("a"));
    }

    #[test]
    fn eviction_drops_stale_sessions() {
        let (_d, mut s) = temp_store();
        s.sessions.insert("old".into(), SavedSession { last_used: 1, ..Default::default() });
        s.sessions.insert("new".into(), SavedSession { last_used: now_secs(), ..Default::default() });
        s.persist();
        let reloaded = SessionStore::with_path(s.path.clone());
        assert!(!reloaded.has_session("old"));
        assert!(reloaded.has_session("new"));
    }

    #[test]
    fn eviction_caps_session_count() {
        let (_d, mut s) = temp_store();
        let now = now_secs();
        for i in 0..(MAX_SESSIONS + 5) {
            s.sessions.insert(format!("s{i}"), SavedSession { last_used: now - i as u64, ..Default::default() });
        }
        assert!(s.evict(now));
        assert_eq!(s.sessions.len(), MAX_SESSIONS);
        assert!(s.has_session("s0"));
        assert!(!s.has_session(&format!("s{}", MAX_SESSIONS + 4)));
    }

    #[test]
    fn generate_id_is_hex_and_unique() {
        let (_d, mut s) = temp_store();
        let id = s.generate_id();
        assert_eq!(id.len(), 32);
        assert!(id.chars().all(|c| c.is_ascii_hexdigit()));
        s.touch(&id);
        assert_ne!(s.generate_id(), id);
    }

    #[test]
    fn corrupt_file_loads_empty() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("sessions.json");
        std::fs::write(&path, "{ not json").unwrap();
        let s = SessionStore::with_path(path);
        assert!(s.sessions.is_empty());
    }
}
