//! Live runtime state structs for active sessions.

use std::{
    collections::HashMap,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicU64},
    },
    time::Instant,
};

use dashmap::DashMap;

use super::{
    audit::{AuditLog, MAX_SIGNALED_ENTITIES},
    signal_state::SessionSignalState,
    types::{AgentId, SessionId},
};
use crate::wal::format::session::SessionSeqNo;

// ── SessionState ──────────────────────────────────────────────────────────────

/// Runtime state for an active session. Stored in `TidalDb::sessions`.
pub struct SessionState {
    pub id: SessionId,
    pub user_id: u64,
    pub agent_id: AgentId,
    pub policy_name: String,
    /// Wall-clock instant when the session was started.
    pub started_at: Instant,
    /// `started_at` expressed as nanoseconds since Unix epoch (for archival).
    pub started_at_ns: u64,
    /// Caller-supplied metadata attached to the session.
    pub metadata: HashMap<String, String>,
    /// Per-signal-type combined decay + window state (keyed by signal type name).
    pub signals: DashMap<String, SessionSignalState>,
    /// Entity IDs that have received any signal in this session.
    ///
    /// Bounded at [`MAX_SIGNALED_ENTITIES`]; insert exclusively via
    /// [`SessionState::track_signaled_entity`] so the cap is always enforced.
    pub signaled_entities: DashMap<u64, ()>,
    /// Timestamped free-text annotations (e.g., preference hints).
    /// Each entry is `(timestamp_ns, hint_text)`. Capped at 100 entries.
    pub annotations: Mutex<Vec<(u64, String)>>,
    /// Total signals accepted (for audit and policy count cap).
    pub signals_written: AtomicU64,
    /// Total signals rejected by policy.
    pub signals_rejected: AtomicU64,
    /// Policy audit log with bounded eviction.
    pub audit_log: Mutex<AuditLog>,
    /// `true` once `close_session` has consumed the handle.
    pub closed: Arc<AtomicBool>,
}

impl SessionState {
    /// Construct a fresh `SessionState` with all five empty-state fields
    /// (`signals`, `signaled_entities`, `annotations`, the two `AtomicU64`
    /// counters, and `audit_log`) initialized to their canonical empty values.
    ///
    /// Both production construction sites — `start_session` (live) and
    /// `restore_session_wal_events` (crash restore) — route through this
    /// constructor so the 14-field literal is never hand-duplicated and an
    /// atomic's initial value can never silently diverge between the two paths.
    ///
    /// The caller supplies the durable identity fields plus the in-process
    /// `closed` flag (`start_session` shares it with the returned
    /// `SessionHandle`; restore allocates a fresh one).
    ///
    /// # Durable age invariant
    ///
    /// `started_at` is a monotonic [`Instant`] for **live metrics only**. The
    /// authoritative session age is `started_at_ns` (nanoseconds since the Unix
    /// epoch). On restore the caller must **back-date** `started_at` so that
    /// `started_at.elapsed()` and `now.duration_since(started_at)` both reflect
    /// the true wall-clock age — the monotonic clock is not durable, so seeding
    /// it with a fresh `Instant::now()` would reset every restored session's age
    /// to ~zero and silently void TTL / `max_session_duration` enforcement.
    #[must_use]
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        id: SessionId,
        user_id: u64,
        agent_id: AgentId,
        policy_name: String,
        started_at: Instant,
        started_at_ns: u64,
        metadata: HashMap<String, String>,
        closed: Arc<AtomicBool>,
    ) -> Self {
        Self {
            id,
            user_id,
            agent_id,
            policy_name,
            started_at,
            started_at_ns,
            metadata,
            signals: DashMap::new(),
            signaled_entities: DashMap::new(),
            annotations: Mutex::new(Vec::new()),
            signals_written: AtomicU64::new(0),
            signals_rejected: AtomicU64::new(0),
            audit_log: Mutex::new(AuditLog::new()),
            closed,
        }
    }

    /// Record that `entity_id` received a signal in this session, enforcing the
    /// [`MAX_SIGNALED_ENTITIES`] cap.
    ///
    /// Returns `true` if the entity is now tracked (already present or newly
    /// inserted), `false` if it was dropped because the set is at capacity. The
    /// FOR SESSION entity boost degrades gracefully: a session that touches more
    /// than `MAX_SIGNALED_ENTITIES` distinct entities boosts the first that many.
    ///
    /// This is the only sanctioned write path into `signaled_entities`; all
    /// callers (live signal write and WAL restore) must route through it so the
    /// bound cannot be bypassed.
    pub fn track_signaled_entity(&self, entity_id: u64) -> bool {
        // Already tracked: cheap fast path, no capacity check needed.
        if self.signaled_entities.contains_key(&entity_id) {
            return true;
        }
        // New entity. Check the cap BEFORE taking an entry lock: DashMap::len
        // locks every shard, so checking it while holding a per-shard entry lock
        // would self-deadlock. The cap is a structural safety bound, so a benign
        // race (briefly exceeding by the number of concurrent writers) is fine —
        // it converges back under the bound and never grows without limit.
        if self.signaled_entities.len() >= MAX_SIGNALED_ENTITIES {
            return false;
        }
        self.signaled_entities.insert(entity_id, ());
        true
    }
}

// ── SessionHandle ─────────────────────────────────────────────────────────────

/// Move-only handle to an active session.
///
/// Ownership is consumed by `TidalDb::close_session` to prevent use-after-close
/// at the type level. The `closed` `Arc<AtomicBool>` provides runtime
/// defense-in-depth for any clone held separately.
#[derive(Debug)]
pub struct SessionHandle {
    pub id: SessionId,
    pub user_id: u64,
    pub agent_id: AgentId,
    pub policy_name: String,
    pub started_at: Instant,
    /// Shared with `SessionState`; set to `true` by `close_session`.
    pub closed: Arc<AtomicBool>,
}

// ── SessionSeqNoTracker ──────────────────────────────────────────────────────

/// Per-session monotonic write counter for idempotent replay.
///
/// Tracks the highest seqno applied locally. Writes with seqno <= hwm
/// are silently dropped (idempotent replay; state is already reflected).
#[derive(Debug, Default)]
pub struct SessionSeqNoTracker {
    hwm: DashMap<u64, SessionSeqNo>, // keyed by session_id.as_u64()
}

impl SessionSeqNoTracker {
    #[must_use]
    pub fn new() -> Self {
        Self {
            hwm: DashMap::new(),
        }
    }

    /// Returns `true` if this write should be applied (seqno > hwm).
    /// Returns `false` if the write is a duplicate -- skip.
    /// Updates the HWM on accept.
    #[must_use]
    pub fn should_apply(&self, session_id: u64, seqno: SessionSeqNo) -> bool {
        let mut entry = self.hwm.entry(session_id).or_insert(SessionSeqNo::ZERO);
        if seqno > *entry {
            *entry = seqno;
            true
        } else {
            false
        }
    }

    /// Current HWM for a session (ZERO if unknown).
    #[must_use]
    pub fn hwm(&self, session_id: u64) -> SessionSeqNo {
        self.hwm.get(&session_id).map_or(SessionSeqNo::ZERO, |v| *v)
    }

    /// Set HWM directly (used on follower startup to restore from WAL).
    pub fn set_hwm(&self, session_id: u64, seqno: SessionSeqNo) {
        self.hwm.insert(session_id, seqno);
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use std::sync::atomic::Ordering;

    use super::*;

    fn empty_state() -> SessionState {
        SessionState::new(
            SessionId(1),
            1,
            AgentId::new("agent-a").unwrap(),
            "default".to_string(),
            Instant::now(),
            1,
            HashMap::new(),
            Arc::new(AtomicBool::new(false)),
        )
    }

    /// Regression (WARNING DRY): `SessionState::new` must initialize every
    /// empty-state field to its canonical empty value so the two production
    /// construction sites cannot silently diverge on a counter's initial value
    /// or a missing reset.
    #[test]
    fn new_initializes_empty_state_fields() {
        let state = SessionState::new(
            SessionId(9),
            42,
            AgentId::new("agent-z").unwrap(),
            "p".to_string(),
            Instant::now(),
            123,
            HashMap::new(),
            Arc::new(AtomicBool::new(false)),
        );
        assert_eq!(state.id, SessionId(9));
        assert_eq!(state.user_id, 42);
        assert_eq!(state.started_at_ns, 123);
        assert_eq!(state.signals.len(), 0);
        assert_eq!(state.signaled_entities.len(), 0);
        assert_eq!(state.annotations.lock().unwrap().len(), 0);
        assert_eq!(state.signals_written.load(Ordering::Relaxed), 0);
        assert_eq!(state.signals_rejected.load(Ordering::Relaxed), 0);
        assert!(!state.closed.load(Ordering::Acquire));
    }

    #[test]
    fn track_signaled_entity_is_idempotent() {
        let state = empty_state();
        assert!(state.track_signaled_entity(7));
        assert!(state.track_signaled_entity(7)); // re-track same id
        assert_eq!(state.signaled_entities.len(), 1);
    }

    #[test]
    fn track_signaled_entity_enforces_cap() {
        let state = empty_state();
        // Fill exactly to the cap with distinct ids.
        for i in 0..MAX_SIGNALED_ENTITIES as u64 {
            assert!(state.track_signaled_entity(i));
        }
        assert_eq!(state.signaled_entities.len(), MAX_SIGNALED_ENTITIES);

        // A new distinct id past the cap is dropped...
        assert!(!state.track_signaled_entity(MAX_SIGNALED_ENTITIES as u64));
        assert_eq!(state.signaled_entities.len(), MAX_SIGNALED_ENTITIES);

        // ...but an already-tracked id still reports tracked (no growth).
        assert!(state.track_signaled_entity(0));
        assert_eq!(state.signaled_entities.len(), MAX_SIGNALED_ENTITIES);
    }

    #[test]
    fn seqno_tracker_accepts_monotone() {
        let t = SessionSeqNoTracker::new();
        assert!(t.should_apply(1, SessionSeqNo(1)));
        assert!(t.should_apply(1, SessionSeqNo(2)));
        assert!(t.should_apply(1, SessionSeqNo(5)));
        assert_eq!(t.hwm(1), SessionSeqNo(5));
    }

    #[test]
    fn seqno_tracker_rejects_duplicate() {
        let t = SessionSeqNoTracker::new();
        assert!(t.should_apply(1, SessionSeqNo(3)));
        assert!(!t.should_apply(1, SessionSeqNo(3))); // duplicate
        assert!(!t.should_apply(1, SessionSeqNo(2))); // older
        assert!(t.should_apply(1, SessionSeqNo(4))); // resumes
    }

    #[test]
    fn seqno_tracker_independent_sessions() {
        let t = SessionSeqNoTracker::new();
        assert!(t.should_apply(1, SessionSeqNo(10)));
        assert!(t.should_apply(2, SessionSeqNo(1))); // separate session starts fresh
        assert_eq!(t.hwm(1), SessionSeqNo(10));
        assert_eq!(t.hwm(2), SessionSeqNo(1));
    }

    #[test]
    fn seqno_tracker_legacy_none_seqno_always_applies() {
        // Legacy events have no seqno; the caller skips the should_apply check
        // for None. Verify tracker treats unknown session as HWM=ZERO.
        let t = SessionSeqNoTracker::new();
        assert_eq!(t.hwm(999), SessionSeqNo::ZERO);
    }
}
