//! Session lifecycle operations on `TidalDb`.

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

use crate::replication::idempotency::IdempotencyKey;
use crate::schema::{EntityId, TidalError, Timestamp};
use crate::session::{
    self as session_mod, AgentId, AuditEntry, SessionHandle, SessionId, SessionInfo,
    SessionSnapshot, SessionState, SessionSummary,
};
use crate::storage::{Tag, encode_key};
use crate::wal::format::session::SessionSeqNo;

use super::TidalDb;

impl TidalDb {
    /// Start a new agent session.
    ///
    /// Creates a session-scoped signal context for the given agent. The
    /// session is identified by its `SessionId` and keyed to `user_id` and
    /// `agent_id`. The `policy_name` must match a policy declared via
    /// `SchemaBuilder::session_policy()`.
    ///
    /// # Errors
    ///
    /// - `TidalError::Schema` if `policy_name` is not found in the schema.
    /// - `TidalError::Internal` if no schema was provided at open time.
    #[tracing::instrument(skip_all, fields(user_id, agent_id, policy_name))]
    pub fn start_session(
        &self,
        user_id: u64,
        agent_id: &str,
        policy_name: &str,
        metadata: HashMap<String, String>,
    ) -> crate::Result<SessionHandle> {
        self.require_writeable("start_session")?;
        // Validate policy exists in schema.
        let schema = self.schema_def.as_ref().ok_or_else(|| {
            TidalError::internal("start_session", "no schema: open with with_schema()")
        })?;
        if schema.session_policy(policy_name).is_none() {
            return Err(TidalError::internal(
                "start_session",
                format!("policy '{policy_name}' not found in schema"),
            ));
        }

        let parsed_agent_id = AgentId::new(agent_id)
            .map_err(|e| TidalError::internal("start_session", format!("invalid agent_id: {e}")))?;

        let session_id = SessionId::from_raw(self.next_session_id.fetch_add(1, Ordering::Relaxed));

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

        // Capture started_at once -- shared between SessionState and SessionHandle.
        let started_at = std::time::Instant::now();
        let started_at_ns = Timestamp::now().as_nanos();

        let state = Arc::new(SessionState {
            id: session_id,
            user_id,
            agent_id: parsed_agent_id.clone(),
            policy_name: policy_name.to_owned(),
            started_at,
            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: Arc::clone(&closed),
        });

        self.sessions.insert(session_id, Arc::clone(&state));

        // Persist the start record in non-ephemeral storage so that a crash
        // during an active session can be detected at next startup.
        if let Some(storage) = self.storage.as_ref() {
            let start_bytes = session_mod::serialize_start_record(&state);
            let start_key = encode_key(EntityId::new(session_id.as_u64()), Tag::Session, b"start");
            if let Err(e) = storage.items_engine().put(&start_key, &start_bytes) {
                tracing::warn!(error = %e, session_id = %session_id, "failed to persist session start record");
            }
        }

        // Write session start event to the session journal (fire-and-forget).
        if let Ok(guard) = self.wal.lock()
            && let Some(wal) = guard.as_ref()
            && let Err(e) = wal.session_start(
                session_id.as_u64(),
                user_id,
                started_at_ns,
                parsed_agent_id.as_str(),
                policy_name,
            )
        {
            tracing::warn!(error = %e, "session_start WAL write failed");
        }

        #[cfg(feature = "metrics")]
        self.metrics.active_sessions.fetch_add(1, Ordering::Relaxed);

        Ok(SessionHandle {
            id: session_id,
            user_id,
            agent_id: parsed_agent_id,
            policy_name: policy_name.to_owned(),
            started_at,
            closed,
        })
    }

    /// Close a session and return a summary.
    ///
    /// Takes ownership of the `SessionHandle` to prevent use-after-close at
    /// compile time. The session snapshot is archived to `closed_sessions`.
    ///
    /// # Errors
    ///
    /// - `TidalError::Internal` if the session was already removed (double-close).
    #[allow(clippy::needless_pass_by_value)] // Intentional: move semantics prevent use-after-close at the type level.
    pub fn close_session(&self, handle: SessionHandle) -> crate::Result<SessionSummary> {
        // Mark the handle as closed (runtime defense-in-depth).
        handle.closed.store(true, Ordering::Release);
        self.close_session_internal(handle.id, false)
    }

    /// Internal: close a session by ID without requiring a `SessionHandle`.
    ///
    /// Used by the TTL sweeper and shutdown cleanup. Sets the `closed`
    /// `AtomicBool` on the session state to prevent further signal writes
    /// from any handle that may still be held (defense-in-depth).
    ///
    /// # Errors
    ///
    /// - `TidalError::Internal` if the session is not found (already closed).
    pub fn close_session_internal(
        &self,
        session_id: SessionId,
        auto_closed: bool,
    ) -> crate::Result<SessionSummary> {
        let (_id, state) = self.sessions.remove(&session_id).ok_or_else(|| {
            TidalError::internal(
                "close_session",
                format!("session {session_id} not found (already closed?)"),
            )
        })?;

        // Mark as closed to prevent further signal writes from any
        // outstanding SessionHandle references.
        state.closed.store(true, Ordering::Release);

        let duration_ms = state.started_at.elapsed().as_millis() as u64;
        let signals_written = state.signals_written.load(Ordering::Relaxed);
        let rejections = state.signals_rejected.load(Ordering::Relaxed);

        // Build and archive the frozen snapshot.
        let snapshot = session_mod::build_frozen_snapshot(&state, duration_ms);

        // Persist the frozen snapshot and remove the start record atomically.
        if let Some(storage) = self.storage.as_ref() {
            let snapshot_key = encode_key(
                EntityId::new(session_id.as_u64()),
                Tag::Session,
                b"snapshot",
            );
            let start_key = encode_key(EntityId::new(session_id.as_u64()), Tag::Session, b"start");
            let snapshot_bytes = session_mod::serialize_snapshot(&snapshot);
            let mut batch = crate::storage::WriteBatch::new();
            batch.put(snapshot_key, snapshot_bytes);
            batch.delete(start_key);
            if let Err(e) = storage.items_engine().write_batch(batch) {
                tracing::warn!(error = %e, session_id = %session_id, "failed to persist session snapshot");
            }
        }

        // Write session close event to the session journal (fire-and-forget).
        if let Ok(guard) = self.wal.lock()
            && let Some(wal) = guard.as_ref()
        {
            let _ = wal.session_close(session_id.as_u64());
        }

        // Evict oldest closed sessions if the cap is exceeded.
        if self.closed_sessions.len() >= session_mod::MAX_CLOSED_SESSIONS {
            let mut entries: Vec<_> = self
                .closed_sessions
                .iter()
                .map(|r| (*r.key(), r.value().closed_at_ns))
                .collect();
            entries.sort_unstable_by_key(|(_, ts)| *ts);
            for (id, _) in entries.iter().take(session_mod::EVICT_BATCH_SIZE) {
                self.closed_sessions.remove(id);
            }
        }

        // Resolve user_id before state is dropped.
        let user_id = state.user_id;

        // M7p2: Clean up rate limiter bucket for this session.
        self.rate_limiter
            .remove(state.agent_id.as_str(), session_id.as_u64());

        // M6p4: Cross-session preference aggregation.
        // For each entity that received a session signal, blend its stored item
        // embedding into the user's global preference vector with adaptive learning rate.
        self.apply_session_preference_update(user_id, &snapshot);

        self.closed_sessions.insert(session_id, snapshot);

        #[cfg(feature = "metrics")]
        {
            self.metrics.active_sessions.fetch_sub(1, Ordering::Relaxed);
            self.metrics
                .closed_sessions_total
                .fetch_add(1, Ordering::Relaxed);
            if auto_closed {
                self.metrics
                    .session_auto_closed_total
                    .fetch_add(1, Ordering::Relaxed);
            }
        }

        tracing::debug!(
            session_id = %session_id,
            auto_closed,
            signals_written,
            rejections,
            duration_ms,
            "session closed"
        );

        Ok(SessionSummary {
            id: session_id,
            duration_ms,
            signals_written,
            rejections,
            auto_closed,
        })
    }

    /// List all currently active sessions.
    #[must_use]
    pub fn active_sessions(&self) -> Vec<SessionInfo> {
        self.sessions
            .iter()
            .map(|entry| {
                let s = entry.value();
                SessionInfo {
                    id: s.id,
                    user_id: s.user_id,
                    agent_id: s.agent_id.as_str().to_owned(),
                    started_at_ns: s.started_at_ns,
                    signals_written: s.signals_written.load(Ordering::Relaxed),
                }
            })
            .collect()
    }

    /// Write a session-scoped signal for an entity.
    ///
    /// Session signals are tracked in the session's in-memory `SessionHotState`
    /// with aggressive decay (5-minute half-life by default). They do **not**
    /// propagate to the global `SignalLedger`; they exist only within this session
    /// and are archived on `close_session`.
    ///
    /// If the session has a policy, it is evaluated before the write; rejected
    /// signals are counted and logged to the audit trail.
    ///
    /// # Errors
    ///
    /// - `TidalError::Internal` if `signal_type` is not in the schema, the session
    ///   is closed, or not found.
    /// - `TidalError::PolicyViolation` if the policy rejects the signal.
    /// - `TidalError::SessionExpired` if the session's duration limit is exceeded.
    #[allow(clippy::significant_drop_tightening)] // state_ref must live for the duration of the method.
    #[allow(clippy::too_many_lines)]
    pub fn session_signal(
        &self,
        handle: &SessionHandle,
        signal_type: &str,
        entity_id: EntityId,
        weight: f64,
        ts: Timestamp,
        annotation: Option<String>,
    ) -> crate::Result<()> {
        // Runtime guard: check the closed flag.
        if handle.closed.load(Ordering::Acquire) {
            return Err(TidalError::internal(
                "session_signal",
                format!("session {} is closed", handle.id),
            ));
        }

        // M7p2: Rate limit check (per-agent, per-session).
        if let Err((retry_after_ms, limit)) = self
            .rate_limiter
            .check(handle.agent_id.as_str(), handle.id.as_u64())
        {
            #[cfg(feature = "metrics")]
            self.metrics
                .rate_limited_total
                .fetch_add(1, Ordering::Relaxed);
            return Err(TidalError::RateLimited {
                agent_id: handle.agent_id.as_str().to_owned(),
                limit,
                retry_after_ms,
            });
        }

        // Validate signal_type exists in the schema.
        if let Some(ledger) = self.ledger.as_ref()
            && ledger.resolve_signal_type(signal_type).is_err()
        {
            return Err(TidalError::internal(
                "session_signal",
                format!("unknown signal type: '{signal_type}'"),
            ));
        }

        let state_ref = self.sessions.get(&handle.id).ok_or_else(|| {
            TidalError::internal("session_signal", format!("session {} not found", handle.id))
        })?;
        let state = state_ref.value();

        // Policy evaluation.
        if let Some(schema) = &self.schema_def
            && let Some(policy) = schema.session_policy(&state.policy_name)
        {
            let evaluator = session_mod::PolicyEvaluator::new(policy, &state.policy_name);
            match evaluator.check(signal_type, state, std::time::Instant::now()) {
                Ok(()) => {
                    // Record accepted entry via bounded AuditLog.
                    let entry = AuditEntry {
                        timestamp_ns: ts.as_nanos(),
                        signal_type: signal_type.to_owned(),
                        accepted: true,
                        reason: None,
                    };
                    if let Ok(mut log) = state.audit_log.lock() {
                        log.push(entry);
                    }
                }
                Err(violation) => {
                    // Record rejected entry via bounded AuditLog.
                    let entry = AuditEntry {
                        timestamp_ns: ts.as_nanos(),
                        signal_type: signal_type.to_owned(),
                        accepted: false,
                        reason: Some(violation.reason.clone()),
                    };
                    if let Ok(mut log) = state.audit_log.lock() {
                        log.push(entry);
                    }
                    state.signals_rejected.fetch_add(1, Ordering::Relaxed);

                    // Dispatch on the typed violation kind (no string parsing).
                    return match violation.kind {
                        session_mod::PolicyViolationKind::Expired => {
                            Err(TidalError::SessionExpired {
                                session_id: handle.id.as_u64(),
                                max_duration_secs: policy.max_session_duration.as_secs_f64(),
                            })
                        }
                        _ => Err(TidalError::PolicyViolation {
                            signal_type: violation.signal_type,
                            policy_name: violation.policy_name,
                            reason: violation.reason,
                        }),
                    };
                }
            }
        }

        // Resolve per-signal-type lambda from schema; fall back to 5-min default.
        let lambda = self
            .schema_def
            .as_ref()
            .and_then(|s| s.signal(signal_type))
            .and_then(|def| def.decay().lambda())
            .unwrap_or(session_mod::DEFAULT_SESSION_LAMBDA);

        // Write to session signal state (CAS decay + windowed count).
        let ts_ns = ts.as_nanos();
        let sig_entry = state
            .signals
            .entry(signal_type.to_owned())
            .or_insert_with(|| session_mod::SessionSignalState::new(ts_ns, lambda));
        sig_entry.on_signal(weight, ts_ns);
        drop(sig_entry);

        // Track signaled entity.
        state.signaled_entities.insert(entity_id.as_u64(), ());

        // Store annotation with timestamp (capped at MAX_ANNOTATIONS).
        let ann_for_wal = annotation.clone();
        if let Some(ann) = annotation
            && let Ok(mut anns) = state.annotations.lock()
            && anns.len() < session_mod::MAX_ANNOTATIONS
        {
            anns.push((ts_ns, ann));
        }

        // Increment and capture the write seqno for idempotent replay.
        // fetch_add returns the OLD value; seqno is old + 1.
        let seqno_raw = state.signals_written.fetch_add(1, Ordering::Relaxed) + 1;

        // Derive a BLAKE3 idempotency key from session context + operation payload.
        let ikey = {
            let mut op_bytes = Vec::with_capacity(8 + signal_type.len());
            op_bytes.extend_from_slice(&entity_id.as_u64().to_le_bytes());
            op_bytes.extend_from_slice(signal_type.as_bytes());
            IdempotencyKey::derive(handle.id.as_u64(), SessionSeqNo(seqno_raw), &op_bytes)
        };

        // Write session signal event to the session journal (fire-and-forget).
        #[allow(clippy::cast_possible_truncation)]
        if let Ok(guard) = self.wal.lock()
            && let Some(wal) = guard.as_ref()
        {
            let _ = wal.session_signal(
                handle.id.as_u64(),
                entity_id.as_u64(),
                weight as f32,
                ts_ns,
                signal_type,
                ann_for_wal.as_deref(),
                Some(seqno_raw),
                Some(ikey.0),
            );
        }

        Ok(())
    }

    /// Retrieve a snapshot of an active or archived session.
    ///
    /// For active sessions the decay scores are computed at the current
    /// wall-clock time. For archived sessions the scores are frozen at the
    /// moment `close_session` was called.
    ///
    /// # Errors
    ///
    /// - `TidalError::Internal` if the session is not found.
    pub fn session_snapshot(&self, session_id: SessionId) -> crate::Result<SessionSnapshot> {
        // Try active sessions first.
        if let Some(state_ref) = self.sessions.get(&session_id) {
            let state = state_ref.value();
            let now_ns = Timestamp::now().as_nanos();
            return Ok(session_mod::build_snapshot(state, now_ns));
        }

        // Fall back to archived sessions (in-memory cache).
        if let Some(snap_ref) = self.closed_sessions.get(&session_id) {
            return Ok(snap_ref.value().clone());
        }

        // Last resort: check persistent storage for archived snapshots.
        if let Some(storage) = self.storage.as_ref() {
            let snapshot_key = encode_key(
                EntityId::new(session_id.as_u64()),
                Tag::Session,
                b"snapshot",
            );
            if let Ok(Some(bytes)) = storage.items_engine().get(&snapshot_key)
                && let Some(snapshot) = session_mod::deserialize_snapshot(&bytes)
            {
                // Warm the in-memory cache to avoid repeated storage reads.
                self.closed_sessions.insert(session_id, snapshot.clone());
                return Ok(snapshot);
            }
        }

        Err(TidalError::internal(
            "session_snapshot",
            format!("session {session_id} not found"),
        ))
    }

    /// Apply cross-session preference aggregation.
    ///
    /// For each entity that received a session signal, read its stored item
    /// embedding and blend it into the user's global preference vector.
    /// Uses adaptive learning rate starting at 0.1, decaying as
    /// `base / (1 + ln(count + 1))` so that early sessions have high
    /// influence and later sessions refine the preference more gently.
    /// Entities without stored embeddings are skipped.
    fn apply_session_preference_update(&self, user_id: u64, snapshot: &SessionSnapshot) {
        let Some(storage) = self.storage.as_ref() else {
            return; // ephemeral mode with no storage
        };

        let prefs = self.preference_vectors();

        for &entity_id in &snapshot.signaled_entities {
            // Read the item embedding from storage.
            // Embeddings are stored at Tag::Meta with suffix "EMB:content".
            let key = encode_key(
                EntityId::new(entity_id),
                crate::storage::Tag::Meta,
                b"EMB:content",
            );
            if let Ok(Some(bytes)) = storage.items_engine().get(&key)
                && let Ok(emb) =
                    crate::storage::vector::lifecycle::serde::deserialize_embedding(&bytes)
            {
                let _ = prefs.update(user_id, &emb);
            }
        }
    }

    /// Retrieve the policy audit log for a session.
    ///
    /// Returns all accept/reject decisions recorded by the policy evaluator
    /// for the given session.
    ///
    /// # Errors
    ///
    /// - `TidalError::Internal` if the session is not found or the audit log
    ///   mutex is poisoned.
    pub fn session_audit(&self, session_id: SessionId) -> crate::Result<Vec<AuditEntry>> {
        // Try active sessions first.
        if let Some(state_ref) = self.sessions.get(&session_id) {
            let state = state_ref.value();
            let log = state
                .audit_log
                .lock()
                .map_err(|_| TidalError::internal("session_audit", "audit_log mutex poisoned"))?;
            return Ok(log.entries().to_vec());
        }

        // For archived sessions, return the audit log captured at close time.
        if let Some(snap_ref) = self.closed_sessions.get(&session_id) {
            return Ok(snap_ref.value().audit_log.clone());
        }

        Err(TidalError::internal(
            "session_audit",
            format!("session {session_id} not found"),
        ))
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
#[path = "session_tests.rs"]
mod tests;
