//! # Session lifecycle and persistence guarantees
//!
//! [`AgentSession`] drives a background turn-processing loop that persists
//! conversation history to an optional [`SessionStore`].  The following
//! invariants are maintained under concurrent `reset()`, `flush()`, `close()`,
//! and `interrupt()` calls:
//!
//! ## Epoch-based reset isolation
//!
//! Every `reset()` bumps an atomic epoch counter.  The background loop
//! captures the epoch at the start of each turn and re-checks it at every
//! mutation point (history push, store append, post-turn write-back).  If the
//! epoch has advanced, the turn is **silently abandoned** — no stale data is
//! written to history or the store.
//!
//! ## No duplicate user-message persistence
//!
//! The user message is appended to the store **once** before the turn runs.
//! After the turn completes, only assistant / tool messages produced during
//! the turn are appended; the user message is skipped via `previous_len`
//! bookkeeping.
//!
//! ## Flush / close during active turns
//!
//! `close()` sets a shared `closed` flag that prevents further message
//! processing, interrupts any in-flight turn, and performs a final flush.
//! Because `AgentSession` is `Clone` and the underlying `mpsc::Sender` is
//! reference-counted, calling `close()` on one handle does not drop the
//! input channel for other clones — the `closed` flag (shared via `Arc`)
//! is what guards against post-close activity.  Any turn that observes
//! the closed flag after returning from the LLM will skip its post-turn
//! persistence to avoid racing with the final flush.
//!
//! ## Interrupt during ask_user wait
//!
//! Calling `interrupt()` (directly or via `reset()` / `close()`) sends
//! `AgentOp::Interrupt` through the ops channel.  The agent loop honours
//! this even while blocked on an `ask_user` tool, returning an interrupted
//! result so the background loop can proceed to the next epoch check.
//!
//! ## Drop behaviour
//!
//! `Drop` performs a **best-effort** flush.  For strict durability, call
//! `flush()` or `close()` explicitly before dropping.

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

use tokio::sync::{mpsc, Mutex};
use tokio::task::JoinHandle;

use crate::loop_::BackpressurePolicy;
use crate::{AgentEvent, AgentLoop, AgentOp, LlmClient, Message, Result, SessionStore};

/// A managed agent session that processes user messages in a background task,
/// persists conversation history, and supports reset / interrupt / close
/// lifecycle operations.
///
/// See the [module-level documentation](self) for persistence guarantees.
#[derive(Clone)]
pub struct AgentSession {
    session_id: String,
    input_tx: mpsc::Sender<String>,
    current_ops_tx: Arc<Mutex<Option<mpsc::Sender<AgentOp>>>>,
    ops_backpressure: BackpressurePolicy,
    ops_capacity: usize,
    on_event: Arc<dyn Fn(AgentEvent) + Send + Sync>,
    ops_saturated: Arc<AtomicBool>,
    history: Arc<Mutex<Vec<Message>>>,
    store: Option<Arc<dyn SessionStore>>,
    epoch: Arc<AtomicU64>,
    closed: Arc<AtomicBool>,
}

impl AgentSession {
    pub fn new(
        agent_loop: AgentLoop,
        llm: Arc<dyn LlmClient>,
        on_event: impl Fn(AgentEvent) + Send + Sync + 'static,
    ) -> (Self, JoinHandle<()>) {
        Self::new_with_history(
            format!("session-{}", uuid_like()),
            None,
            Vec::new(),
            agent_loop,
            llm,
            on_event,
        )
    }

    pub fn new_with_store(
        session_id: impl Into<String>,
        store: Arc<dyn SessionStore>,
        agent_loop: AgentLoop,
        llm: Arc<dyn LlmClient>,
        on_event: impl Fn(AgentEvent) + Send + Sync + 'static,
    ) -> (Self, JoinHandle<()>) {
        Self::new_with_history(
            session_id.into(),
            Some(store),
            Vec::new(),
            agent_loop,
            llm,
            on_event,
        )
    }

    pub async fn resume(
        session_id: &str,
        store: Arc<dyn SessionStore>,
        agent_loop: AgentLoop,
        llm: Arc<dyn LlmClient>,
        on_event: impl Fn(AgentEvent) + Send + Sync + 'static,
    ) -> Result<(Self, JoinHandle<()>)> {
        let history = store.load(session_id).await?;
        if let Some(file_store) = store
            .as_any()
            .downcast_ref::<crate::session_store::FileSessionStore>()
        {
            if let Some(meta) = file_store.load_meta(session_id).await? {
                if meta.agent_loop_version != env!("CARGO_PKG_VERSION") {
                    eprintln!(
                        "warning: session '{}' saved by version {}, current version {}",
                        session_id,
                        meta.agent_loop_version,
                        env!("CARGO_PKG_VERSION")
                    );
                }
            }
        }

        Ok(Self::new_with_history(
            session_id.to_string(),
            Some(store),
            history,
            agent_loop,
            llm,
            on_event,
        ))
    }

    /// Send a user message into the session for processing.
    ///
    /// Returns immediately; the message is queued and processed by the
    /// background loop.  If the session has been closed, the message is
    /// silently dropped.
    pub async fn send(&self, msg: &str) {
        if self.closed.load(Ordering::SeqCst) {
            return;
        }
        let _ = self.input_tx.send(msg.to_string()).await;
    }

    /// Interrupt the currently running turn (if any).
    ///
    /// This sends `AgentOp::Interrupt` through the ops channel so the agent
    /// loop can break out of tool execution or ask_user wait promptly.
    pub async fn interrupt(&self) {
        self.send_op(AgentOp::Interrupt).await;
    }

    pub async fn inject(&self, context: &str) {
        self.send_op(AgentOp::InjectHint(context.to_string())).await;
    }

    pub async fn reply(&self, call_id: Option<&str>, answer: &str) {
        self.send_op(AgentOp::AskUserAnswer {
            call_id: call_id.map(ToString::to_string),
            answer: answer.to_string(),
        })
        .await;
    }

    /// Send an operation to the current turn's ops channel, respecting the
    /// configured [`BackpressurePolicy`].
    ///
    /// Returns `true` if the op was successfully sent, `false` if it was
    /// dropped or rejected (or no active turn exists).
    ///
    /// # Backpressure behavior
    ///
    /// - **Block**: Awaits until capacity is available (default).
    /// - **Reject**: Returns `false` immediately if the channel is full.
    /// - **DropOldest**: Drops the newest item (the one being sent) when full,
    ///   since tokio's mpsc does not expose a pop-oldest primitive. For true
    ///   ring-buffer semantics consider a dedicated data structure.
    pub async fn send_op(&self, op: AgentOp) -> bool {
        let tx = match self.current_ops_tx.lock().await.clone() {
            Some(tx) => tx,
            None => return false,
        };
        match self.ops_backpressure {
            BackpressurePolicy::Block => {
                let _ = tx.send(op).await;
                true
            }
            BackpressurePolicy::Reject => match tx.try_send(op) {
                Ok(()) => true,
                Err(mpsc::error::TrySendError::Full(rejected_op)) => {
                    // Emit OpsSaturated once on first saturation.
                    if !self.ops_saturated.swap(true, Ordering::SeqCst) {
                        (self.on_event)(AgentEvent::OpsSaturated {
                            capacity: self.ops_capacity,
                        });
                    }
                    (self.on_event)(AgentEvent::OpRejected {
                        reason: format!("{rejected_op:?}"),
                    });
                    false
                }
                Err(mpsc::error::TrySendError::Closed(_)) => false,
            },
            BackpressurePolicy::DropOldest => match tx.try_send(op) {
                Ok(()) => true,
                Err(mpsc::error::TrySendError::Full(dropped_op)) => {
                    // Emit OpsSaturated once on first saturation.
                    if !self.ops_saturated.swap(true, Ordering::SeqCst) {
                        (self.on_event)(AgentEvent::OpsSaturated {
                            capacity: self.ops_capacity,
                        });
                    }
                    (self.on_event)(AgentEvent::OpDropped {
                        reason: format!("{dropped_op:?}"),
                    });
                    false
                }
                Err(mpsc::error::TrySendError::Closed(_)) => false,
            },
        }
    }

    /// Reset the session: bump epoch, interrupt in-flight turn, clear history
    /// and delete persisted state.
    ///
    /// Any turn that was in progress will observe the epoch change and
    /// silently discard its results rather than writing stale data.
    pub async fn reset(&self) {
        self.epoch.fetch_add(1, Ordering::SeqCst);
        self.interrupt().await;
        {
            let mut history = self.history.lock().await;
            history.clear();
        }
        if let Some(store) = &self.store {
            let _ = store.delete(&self.session_id).await;
        }
    }

    pub async fn history(&self) -> Vec<Message> {
        self.history.lock().await.clone()
    }

    /// Force-flush persisted session state to durable storage.
    ///
    /// Use this before process exit when strict durability is required.
    /// This avoids relying on `Drop`, which is intentionally best-effort.
    pub async fn flush(&self) -> Result<()> {
        if let Some(store) = &self.store {
            store.flush(&self.session_id).await?;
        }
        Ok(())
    }

    /// Flush and gracefully close this session.
    ///
    /// After close:
    /// - New `send()` calls are silently dropped.
    /// - The background loop's current turn (if any) is interrupted.
    /// - A final flush is issued to persist any buffered state.
    ///
    /// Prefer calling this in shutdown paths when durability is important.
    /// Note: because `AgentSession` is `Clone`, other handles may still
    /// exist; the `closed` flag is shared so they will also observe it.
    pub async fn close(self) -> Result<()> {
        self.closed.store(true, Ordering::SeqCst);
        self.interrupt().await;
        self.flush().await
    }

    /// Returns `true` if [`close()`](Self::close) has been called.
    pub fn is_closed(&self) -> bool {
        self.closed.load(Ordering::SeqCst)
    }

    fn new_with_history(
        session_id: String,
        store: Option<Arc<dyn SessionStore>>,
        initial_history: Vec<Message>,
        agent_loop: AgentLoop,
        llm: Arc<dyn LlmClient>,
        on_event: impl Fn(AgentEvent) + Send + Sync + 'static,
    ) -> (Self, JoinHandle<()>) {
        let ch_cfg = agent_loop.channel_config().clone();
        let input_capacity = ch_cfg.input_capacity;
        let ops_capacity = ch_cfg.ops_capacity;
        let ops_backpressure = ch_cfg.ops_backpressure;

        let (input_tx, mut input_rx) = mpsc::channel::<String>(input_capacity);
        let current_ops_tx = Arc::new(Mutex::new(None));
        let ops_saturated = Arc::new(AtomicBool::new(false));
        let history = Arc::new(Mutex::new(initial_history));
        let epoch = Arc::new(AtomicU64::new(0));
        let closed = Arc::new(AtomicBool::new(false));

        let history_bg = history.clone();
        let current_ops_tx_bg = current_ops_tx.clone();
        let ops_saturated_bg = ops_saturated.clone();
        let store_bg = store.clone();
        let session_id_bg = session_id.clone();
        let epoch_bg = epoch.clone();
        let closed_bg = closed.clone();
        let on_event_arc: Arc<dyn Fn(AgentEvent) + Send + Sync> = Arc::new(on_event);
        let on_event_session = on_event_arc.clone();

        let handle = tokio::spawn(async move {
            while let Some(user_message) = input_rx.recv().await {
                // Capture epoch at turn start for staleness detection.
                let turn_epoch = epoch_bg.load(Ordering::SeqCst);

                // If the session is closed, drain remaining messages without
                // processing them.
                if closed_bg.load(Ordering::SeqCst) {
                    continue;
                }

                // --- Phase 1: append user message to in-memory history ---
                let previous_len = {
                    let mut history_lock = history_bg.lock().await;
                    if epoch_bg.load(Ordering::SeqCst) != turn_epoch {
                        continue;
                    }
                    history_lock.push(Message::user(&user_message));
                    history_lock.len() - 1
                };

                // --- Phase 2: persist user message to store (once) ---
                if epoch_bg.load(Ordering::SeqCst) != turn_epoch {
                    continue;
                }
                if let Some(store) = &store_bg {
                    let _ = store
                        .append(&session_id_bg, &Message::user(&user_message))
                        .await;
                }

                // --- Phase 3: run the agent turn ---
                if epoch_bg.load(Ordering::SeqCst) != turn_epoch {
                    continue;
                }
                let turn_input = { history_bg.lock().await.clone() };
                let (turn_ops_tx, turn_ops_rx) = mpsc::channel::<AgentOp>(ops_capacity);
                // Reset saturation flag so each turn can independently observe
                // its first saturation event.
                ops_saturated_bg.store(false, Ordering::SeqCst);
                {
                    let mut slot = current_ops_tx_bg.lock().await;
                    *slot = Some(turn_ops_tx);
                }

                let event_handler = on_event_arc.clone();
                let turn_result = agent_loop
                    .run_with_ops(llm.as_ref(), turn_input, Some(turn_ops_rx), move |event| {
                        event_handler(event)
                    })
                    .await;

                {
                    let mut slot = current_ops_tx_bg.lock().await;
                    *slot = None;
                }

                // --- Phase 4: write-back results ---
                // Guard: skip write-back if epoch changed (reset) or session
                // is closed.
                if epoch_bg.load(Ordering::SeqCst) != turn_epoch {
                    continue;
                }
                if closed_bg.load(Ordering::SeqCst) {
                    continue;
                }
                match turn_result {
                    Ok(result) => {
                        // Skip persisting the user message that was already
                        // appended before running this turn — prevents
                        // duplicate user message in the store.
                        let start_idx = previous_len.saturating_add(1);
                        let new_messages = if result.messages.len() > start_idx {
                            result.messages[start_idx..].to_vec()
                        } else {
                            Vec::new()
                        };

                        {
                            let mut history_lock = history_bg.lock().await;
                            // Final epoch check while holding the lock to
                            // prevent a concurrent reset from being
                            // overwritten.
                            if epoch_bg.load(Ordering::SeqCst) != turn_epoch {
                                continue;
                            }
                            *history_lock = result.messages;
                        }

                        if let Some(store) = &store_bg {
                            // Re-check epoch before store writes so a reset
                            // that happened between the history update and
                            // here does not re-persist stale data.
                            if epoch_bg.load(Ordering::SeqCst) != turn_epoch {
                                continue;
                            }
                            for message in &new_messages {
                                let _ = store.append(&session_id_bg, message).await;
                            }
                            let _ = store.flush(&session_id_bg).await;
                        }
                    }
                    Err(_) => {
                        // On error, flush whatever was buffered (the user
                        // message) but only if the epoch hasn't changed.
                        if epoch_bg.load(Ordering::SeqCst) != turn_epoch {
                            continue;
                        }
                        if let Some(store) = &store_bg {
                            let _ = store.flush(&session_id_bg).await;
                        }
                    }
                }
            }
        });

        (
            Self {
                session_id,
                input_tx,
                current_ops_tx,
                ops_backpressure,
                ops_capacity,
                on_event: on_event_session,
                ops_saturated,
                history,
                store,
                epoch,
                closed,
            },
            handle,
        )
    }
}

impl Drop for AgentSession {
    fn drop(&mut self) {
        // Mark closed to prevent background loop from persisting after this
        // point (best-effort, since the flag is checked asynchronously).
        self.closed.store(true, Ordering::SeqCst);

        if let Some(store) = self.store.clone() {
            let session_id = self.session_id.clone();
            if let Ok(handle) = tokio::runtime::Handle::try_current() {
                handle.spawn(async move {
                    let _ = store.flush(&session_id).await;
                });
            } else {
                std::thread::spawn(move || {
                    if let Ok(rt) = tokio::runtime::Builder::new_current_thread()
                        .enable_all()
                        .build()
                    {
                        let _ = rt.block_on(async move { store.flush(&session_id).await });
                    }
                });
            }
        }
    }
}

fn uuid_like() -> String {
    format!(
        "{}",
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or_default()
    )
}
