//! Session restore from persistent storage and WAL replay.

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};

use crate::schema::{EntityId, Timestamp};
use crate::session::{self as session_mod, AgentId, SessionId, SessionState};
use crate::storage::{Tag, encode_key, parse_key};
use crate::wal::format::session::SessionSeqNo;

use super::TidalDb;

impl TidalDb {
    /// Scan persistent storage for previously archived session snapshots.
    ///
    /// Called once at startup from `TidalDbBuilder::open()`. Walks every key
    /// in the items storage engine, filters for `Tag::Session` keys, and:
    /// - `b"snapshot"` suffix -> deserializes into `closed_sessions`
    /// - `b"start"` suffix   -> logs a warning (session was active at shutdown)
    ///
    /// In ephemeral mode the in-memory backend starts empty each time so this
    /// scan always returns nothing -- which is correct.
    pub(crate) fn restore_sessions(&self) {
        let Some(storage) = self.storage.as_ref() else {
            return;
        };
        let engine = storage.items_engine();
        let mut restored = 0usize;
        let mut orphaned = 0usize;

        for item in engine.scan_prefix(&[]) {
            let (key, value) = match item {
                Ok(kv) => kv,
                Err(e) => {
                    tracing::warn!(error = %e, "storage scan error during session restore");
                    continue;
                }
            };
            let Some((_entity_id, tag, suffix)) = parse_key(&key) else {
                continue;
            };
            if tag != Tag::Session {
                continue;
            }
            match suffix {
                b"snapshot" => {
                    if let Some(snapshot) = session_mod::deserialize_snapshot(&value) {
                        self.closed_sessions.insert(snapshot.id, snapshot);
                        restored += 1;
                    } else {
                        tracing::warn!("corrupt session snapshot key {:?}; skipping", key);
                    }
                }
                b"start" => {
                    if let Some((session_id, _user_id, _started_at_ns, _metadata)) =
                        session_mod::deserialize_start_record(&value)
                    {
                        tracing::warn!(
                            session_id = %session_id,
                            "session was active at last shutdown; in-memory signal state is lost"
                        );
                    }
                    orphaned += 1;
                }
                _ => {}
            }
        }

        if restored > 0 || orphaned > 0 {
            tracing::info!(restored, orphaned, "session restore complete");
        }
    }

    /// Replay session journal events to restore active sessions from a crash.
    ///
    /// Scans the event list for Start/Signal/Close triples. Sessions that have
    /// a `Start` but no matching `Close` are restored as active sessions with
    /// their signals replayed into `SessionSignalState`.
    ///
    /// Called from `from_parts()` during startup. The storage-backed
    /// `restore_sessions()` is additive and handles archived snapshots; this
    /// method handles in-flight sessions that were active at crash time.
    #[allow(clippy::too_many_lines, clippy::cast_possible_truncation)]
    pub(super) fn restore_session_wal_events(
        &self,
        events: &[crate::wal::format::SessionWalEvent],
    ) {
        let (open_sessions, session_signals) = Self::partition_session_events(events);

        if open_sessions.is_empty() {
            return;
        }

        let Some(schema) = self.schema_def.as_ref() else {
            tracing::warn!("no schema available; cannot restore sessions from WAL");
            return;
        };

        let now_ns = Timestamp::now().as_nanos();

        for (session_id, (user_id, started_at_ns, agent_id_str, policy_name)) in open_sessions {
            if schema.session_policy(&policy_name).is_none() {
                tracing::warn!(
                    session_id,
                    policy_name,
                    "restored session has unknown policy, skipping"
                );
                continue;
            }

            let agent_id = AgentId::new(&agent_id_str).unwrap_or_else(|_| {
                tracing::warn!(
                    session_id,
                    agent_id = agent_id_str,
                    "invalid agent_id in WAL, using fallback 'restored'"
                );
                // Fallback: AgentId::new("restored") is always valid (lowercase ASCII).
                AgentId::new("restored").expect("'restored' is a valid AgentId")
            });

            let closed = Arc::new(AtomicBool::new(false));

            // Rehydrate metadata from the persisted start record in storage.
            // The WAL Start event does not carry metadata; the storage start record does.
            let metadata = self
                .storage
                .as_ref()
                .and_then(|storage| {
                    let key = encode_key(EntityId::new(session_id), Tag::Session, b"start");
                    storage.items_engine().get(&key).ok().flatten()
                })
                .and_then(|bytes| session_mod::deserialize_start_record(&bytes))
                .map(|(_, _, _, meta)| meta)
                .unwrap_or_default();

            let state = Arc::new(SessionState {
                id: SessionId::from_raw(session_id),
                user_id,
                agent_id,
                policy_name: policy_name.clone(),
                // We lost the exact monotonic Instant -- approximate with "now".
                started_at: std::time::Instant::now(),
                started_at_ns,
                metadata,
                signals: dashmap::DashMap::new(),
                signaled_entities: dashmap::DashMap::new(),
                annotations: std::sync::Mutex::new(Vec::new()),
                signals_written: AtomicU64::new(0),
                signals_rejected: AtomicU64::new(0),
                audit_log: std::sync::Mutex::new(session_mod::AuditLog::new()),
                closed,
            });

            // Replay signals and annotations into the restored session state.
            // Advance the seqno HWM so incoming replication events are correctly deduped.
            let mut max_seqno: Option<u64> = None;
            if let Some(signals) = session_signals.get(&session_id) {
                for (entity_id, weight, ts_ns, signal_name, annotation, seqno_raw) in signals {
                    let lambda = schema
                        .signal(signal_name)
                        .and_then(|def| def.decay().lambda())
                        .unwrap_or(session_mod::DEFAULT_SESSION_LAMBDA);

                    let ss = state
                        .signals
                        .entry(signal_name.clone())
                        .or_insert_with(|| session_mod::SessionSignalState::new(now_ns, lambda));
                    ss.on_signal(f64::from(*weight), *ts_ns);
                    drop(ss);

                    state.signaled_entities.insert(*entity_id, ());
                    state.signals_written.fetch_add(1, Ordering::Relaxed);

                    // Track the highest seqno seen so we can advance the HWM after replay.
                    if let Some(s) = *seqno_raw {
                        max_seqno = Some(max_seqno.map_or(s, |m| m.max(s)));
                    }

                    // Replay annotation if present.
                    if let Some(ann) = annotation
                        && let Ok(mut anns) = state.annotations.lock()
                        && anns.len() < session_mod::MAX_ANNOTATIONS
                    {
                        anns.push((*ts_ns, ann.clone()));
                    }
                }
            }

            // Advance the seqno tracker HWM so the bridge deduplicates any
            // replication events that repeat signals already replayed from the WAL.
            if let Some(hwm) = max_seqno {
                self.session_bridge
                    .seqno_tracker()
                    .set_hwm(session_id, SessionSeqNo(hwm));
            }

            let sid = SessionId::from_raw(session_id);
            self.sessions.insert(sid, state);

            // Advance next_session_id past any restored session IDs.
            loop {
                let current = self.next_session_id.load(Ordering::Acquire);
                if session_id < current {
                    break;
                }
                if self
                    .next_session_id
                    .compare_exchange(
                        current,
                        session_id + 1,
                        Ordering::Release,
                        Ordering::Relaxed,
                    )
                    .is_ok()
                {
                    break;
                }
            }

            tracing::info!(session_id, user_id, "restored active session from WAL");
        }
    }

    /// Partition session WAL events into open sessions and their signals.
    ///
    /// Returns `(open_sessions, session_signals)` where `open_sessions` maps
    /// `session_id` to `(user_id, started_at_ns, agent_id, policy_name)` for
    /// sessions that have a Start but no Close.
    ///
    /// Each signal tuple is `(entity_id, weight, ts_ns, signal_name, annotation, seqno)`.
    /// `seqno` is `Some(n)` for events written after the M8p4 seqno mechanism was added,
    /// `None` for legacy events.
    #[allow(clippy::type_complexity)]
    pub(super) fn partition_session_events(
        events: &[crate::wal::format::SessionWalEvent],
    ) -> (
        HashMap<u64, (u64, u64, String, String)>,
        HashMap<u64, Vec<(u64, f32, u64, String, Option<String>, Option<u64>)>>,
    ) {
        use crate::wal::format::SessionWalEvent;

        let mut open_sessions: HashMap<u64, (u64, u64, String, String)> = HashMap::new();
        let mut session_signals: HashMap<
            u64,
            Vec<(u64, f32, u64, String, Option<String>, Option<u64>)>,
        > = HashMap::new();

        for event in events {
            match event {
                SessionWalEvent::Start {
                    session_id,
                    user_id,
                    started_at_ns,
                    agent_id,
                    policy_name,
                } => {
                    open_sessions.insert(
                        *session_id,
                        (
                            *user_id,
                            *started_at_ns,
                            agent_id.clone(),
                            policy_name.clone(),
                        ),
                    );
                }
                SessionWalEvent::Signal {
                    session_id,
                    entity_id,
                    weight,
                    ts_ns,
                    signal_name,
                    annotation,
                    session_seqno,
                    ..
                } => {
                    session_signals.entry(*session_id).or_default().push((
                        *entity_id,
                        *weight,
                        *ts_ns,
                        signal_name.clone(),
                        annotation.clone(),
                        session_seqno.map(|s| s.0),
                    ));
                }
                SessionWalEvent::Close { session_id } => {
                    open_sessions.remove(session_id);
                    session_signals.remove(session_id);
                }
            }
        }

        (open_sessions, session_signals)
    }
}
