//! Internal state management for ACP connector
//!
//! This module provides thread-safe state tracking for the ACP connector,
//! including protocol version, agent capabilities, and active sessions.

use chrono::Utc;
use dirigent_protocol::{SessionModeState, SessionModelState, SessionOwnership};
use serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use tokio::sync::RwLock;

/// Internal state for ACP connector
///
/// Tracks connection state, protocol information, and active sessions.
/// All fields are protected by RwLock for thread-safe concurrent access.
///
/// # Thread Safety
///
/// This struct is designed to be shared across async tasks via `Arc`.
/// All mutations go through interior mutability patterns (RwLock).
#[derive(Debug)]
pub struct InternalState {
    /// Negotiated protocol version
    ///
    /// Set during initialization handshake. None if not yet initialized.
    protocol_version: Arc<RwLock<Option<u32>>>,

    /// Agent capabilities
    ///
    /// Populated from the initialize response. Contains information about
    /// what the agent supports (streaming, tools, MCP servers, etc.).
    agent_capabilities: Arc<RwLock<Option<Value>>>,

    /// Active sessions
    ///
    /// Map of session ID to session metadata. Updated as sessions are
    /// created, updated, or terminated.
    sessions: Arc<RwLock<HashMap<String, SessionInfo>>>,

    /// Active message IDs per session
    ///
    /// Tracks the current Dirigent message_id being accumulated for each session.
    /// This allows us to translate many ACP chunks (each with different messageIds)
    /// into a single Dirigent message with one consistent message_id.
    active_messages: Arc<RwLock<HashMap<String, String>>>,

    /// Message start times
    ///
    /// Tracks when each message started (when the first chunk arrived).
    /// Maps message_id to timestamp. This allows us to provide accurate
    /// created_at timestamps for messages instead of using finalization time.
    message_start_times: Arc<RwLock<HashMap<String, chrono::DateTime<chrono::Utc>>>>,

    /// Loaded sessions
    ///
    /// Tracks which sessions have been successfully loaded (via session/load
    /// or session/resume) during this connection. Used to decide whether to
    /// load or resume a session: if already loaded, resume; otherwise, load first.
    /// Cleared on reconnect.
    loaded_sessions: Arc<RwLock<HashSet<String>>>,
}

/// Information about an active session
///
/// Lightweight metadata tracked for each session. Does not include
/// full message history (that's stored by the agent).
#[derive(Debug, Clone)]
pub struct SessionInfo {
    /// Unique session identifier
    pub id: String,

    /// Session title (if available)
    pub title: Option<String>,

    /// Current working directory for this session
    pub cwd: String,

    /// Number of messages in this session (approximate)
    pub message_count: u32,

    /// Session model (if available) - legacy field for simple model name
    pub model: Option<String>,

    /// Creation timestamp
    pub created_at: chrono::DateTime<chrono::Utc>,

    /// Last activity timestamp
    pub last_activity: chrono::DateTime<chrono::Utc>,

    /// Session status
    pub status: SessionStatus,

    // ========================================================================
    // ACP Session Metadata (from session/new and session/load responses)
    // ========================================================================

    /// Available models and current model selection
    ///
    /// Captured from ACP `session/new` or `session/load` response.
    /// Note: This field is marked UNSTABLE in the ACP spec but Claude-ACP uses it.
    pub models: Option<SessionModelState>,

    /// Available modes and current mode selection
    ///
    /// Captured from ACP `session/new` or `session/load` response.
    /// This is part of the stable ACP specification.
    pub modes: Option<SessionModeState>,

    /// ACP config options (replaces modes/models in future ACP versions)
    ///
    /// Captured from ACP `session/new` or `session/load` response `configOptions` field.
    pub config_options: Option<Vec<dirigent_protocol::ConfigOption>>,

    /// Ownership model for this session
    ///
    /// Tracks whether this session is internal to Dirigent or originates from
    /// an external client, and how tool calls should be handled.
    pub ownership: SessionOwnership,
}

/// Status of a session
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionStatus {
    /// Session is active and can accept prompts
    Active,

    /// Session is processing a prompt (agent is responding)
    Processing,

    /// Session is idle (no active processing)
    Idle,

    /// Session has ended
    Ended,
}

impl InternalState {
    /// Create a new empty state
    ///
    /// All fields are initialized to None/empty. State is populated
    /// as the connector progresses through initialization and operation.
    pub fn new() -> Self {
        Self {
            protocol_version: Arc::new(RwLock::new(None)),
            agent_capabilities: Arc::new(RwLock::new(None)),
            sessions: Arc::new(RwLock::new(HashMap::new())),
            active_messages: Arc::new(RwLock::new(HashMap::new())),
            message_start_times: Arc::new(RwLock::new(HashMap::new())),
            loaded_sessions: Arc::new(RwLock::new(HashSet::new())),
        }
    }

    /// Get the negotiated protocol version
    ///
    /// Returns None if initialization has not completed.
    pub async fn protocol_version(&self) -> Option<u32> {
        *self.protocol_version.read().await
    }

    /// Set the negotiated protocol version
    ///
    /// Called after successful initialization handshake.
    pub async fn set_protocol_version(&self, version: u32) {
        let mut guard = self.protocol_version.write().await;
        *guard = Some(version);
    }

    /// Get the agent capabilities
    ///
    /// Returns None if initialization has not completed.
    pub async fn agent_capabilities(&self) -> Option<Value> {
        self.agent_capabilities.read().await.clone()
    }

    /// Set the agent capabilities
    ///
    /// Called after successful initialization handshake.
    pub async fn set_agent_capabilities(&self, capabilities: Value) {
        let mut guard = self.agent_capabilities.write().await;
        *guard = Some(capabilities);
    }

    /// Add or update a session
    ///
    /// If the session ID already exists, the info is updated.
    /// Otherwise, a new entry is created.
    pub async fn upsert_session(&self, info: SessionInfo) {
        let mut sessions = self.sessions.write().await;
        sessions.insert(info.id.clone(), info);
    }

    /// Remove a session
    ///
    /// Called when a session is terminated or deleted.
    pub async fn remove_session(&self, session_id: &str) {
        let mut sessions = self.sessions.write().await;
        sessions.remove(session_id);
    }

    /// Get a specific session
    ///
    /// Returns None if the session is not tracked.
    pub async fn get_session(&self, session_id: &str) -> Option<SessionInfo> {
        let sessions = self.sessions.read().await;
        sessions.get(session_id).cloned()
    }

    /// List all tracked sessions
    ///
    /// Returns a snapshot of all session info. The returned Vec is
    /// independent of the internal state and can be safely modified.
    pub async fn list_sessions(&self) -> Vec<SessionInfo> {
        let sessions = self.sessions.read().await;
        sessions.values().cloned().collect()
    }

    /// Get the number of tracked sessions
    pub async fn session_count(&self) -> usize {
        let sessions = self.sessions.read().await;
        sessions.len()
    }

    /// Update message count for a session
    ///
    /// Convenience method to increment or set message count without
    /// replacing the entire SessionInfo.
    pub async fn update_message_count(&self, session_id: &str, count: u32) {
        let mut sessions = self.sessions.write().await;
        if let Some(info) = sessions.get_mut(session_id) {
            info.message_count = count;
            info.last_activity = chrono::Utc::now();
        }
    }

    /// Check if a session is active
    ///
    /// Returns true if the session exists and is not in Ended status.
    pub async fn is_session_active(&self, session_id: &str) -> bool {
        let sessions = self.sessions.read().await;
        sessions
            .get(session_id)
            .map(|info| info.status != SessionStatus::Ended)
            .unwrap_or(false)
    }

    /// Update session status
    ///
    /// Convenience method to update session status and last activity time.
    pub async fn update_session_status(&self, session_id: &str, status: SessionStatus) {
        let mut sessions = self.sessions.write().await;
        if let Some(info) = sessions.get_mut(session_id) {
            info.status = status;
            info.last_activity = Utc::now();
        }
    }

    /// Touch session (update last activity timestamp)
    ///
    /// Called when any activity occurs on a session.
    pub async fn touch_session(&self, session_id: &str) {
        let mut sessions = self.sessions.write().await;
        if let Some(info) = sessions.get_mut(session_id) {
            info.last_activity = Utc::now();
        }
    }

    /// Get or create active message ID for a session
    ///
    /// If no message is currently being accumulated for this session, generates
    /// a new UUID-based Dirigent message_id and stores it. Returns the active
    /// message_id (either existing or newly created).
    ///
    /// This allows us to translate multiple ACP chunks (each with different
    /// ACP messageIds) into a single Dirigent message with one consistent ID.
    pub async fn get_or_create_active_message(&self, session_id: &str) -> String {
        let mut active_messages = self.active_messages.write().await;

        // Check if this is a new message BEFORE inserting
        let is_new_message = !active_messages.contains_key(session_id);

        let message_id = active_messages
            .entry(session_id.to_string())
            .or_insert_with(|| format!("msg-{}", uuid::Uuid::now_v7()))
            .clone();

        // Track message start time when creating new message
        if is_new_message {
            let mut start_times = self.message_start_times.write().await;
            start_times.insert(message_id.clone(), Utc::now());
        }

        message_id
    }

    /// Clear active message ID for a session
    ///
    /// Called when a message is completed (after receiving prompt response).
    /// This allows the next message to get a fresh Dirigent message_id.
    pub async fn clear_active_message(&self, session_id: &str) -> Option<String> {
        let mut active_messages = self.active_messages.write().await;
        active_messages.remove(session_id)
    }

    /// Get active message ID for a session (if any)
    ///
    /// Returns None if no message is currently being accumulated.
    pub async fn get_active_message(&self, session_id: &str) -> Option<String> {
        let active_messages = self.active_messages.read().await;
        active_messages.get(session_id).cloned()
    }

    /// Get message start time
    ///
    /// Returns the timestamp when the message started (first chunk arrived).
    /// Returns None if the message is not tracked.
    pub async fn get_message_start_time(&self, message_id: &str) -> Option<chrono::DateTime<chrono::Utc>> {
        let start_times = self.message_start_times.read().await;
        start_times.get(message_id).copied()
    }

    /// Clear message start time
    ///
    /// Called after a message is finalized and archived.
    /// This prevents memory leaks from accumulating start times.
    pub async fn clear_message_start_time(&self, message_id: &str) {
        let mut start_times = self.message_start_times.write().await;
        start_times.remove(message_id);
    }

    /// Clear all state
    ///
    /// Resets the state to its initial empty condition. Called on
    /// reconnection or when recovering from errors.
    pub async fn clear(&self) {
        let mut protocol_version = self.protocol_version.write().await;
        *protocol_version = None;

        let mut capabilities = self.agent_capabilities.write().await;
        *capabilities = None;

        let mut sessions = self.sessions.write().await;
        sessions.clear();

        let mut active_messages = self.active_messages.write().await;
        active_messages.clear();

        let mut start_times = self.message_start_times.write().await;
        start_times.clear();

        let mut loaded = self.loaded_sessions.write().await;
        loaded.clear();
    }

    /// Check if the connector is initialized
    ///
    /// Returns true if we have completed the initialization handshake
    /// (protocol version and capabilities are set).
    pub async fn is_initialized(&self) -> bool {
        self.protocol_version().await.is_some()
    }

    /// Check if the upstream agent supports session/load
    ///
    /// Returns true if `agentCapabilities.loadSession` is `true`.
    pub async fn agent_supports_load_session(&self) -> bool {
        self.agent_capabilities
            .read()
            .await
            .as_ref()
            .and_then(|caps| caps.get("loadSession"))
            .and_then(|v| v.as_bool())
            .unwrap_or(false)
    }

    /// Check if the upstream agent supports session/list
    ///
    /// Checks `sessionCapabilities.list` (spec-correct, object presence) with
    /// fallback to legacy `listSessions` (boolean).
    pub async fn agent_supports_list_sessions(&self) -> bool {
        self.agent_capabilities.read().await.as_ref()
            .map(|caps| {
                // Spec-correct: sessionCapabilities.list (object presence)
                let nested = caps.get("sessionCapabilities")
                    .and_then(|sc| sc.get("list"))
                    .is_some();
                // Legacy compat: listSessions (boolean)
                let flat = caps.get("listSessions")
                    .and_then(|v| v.as_bool())
                    .unwrap_or(false);
                nested || flat
            })
            .unwrap_or(false)
    }

    /// Check if the upstream agent supports session/resume
    ///
    /// Returns true if `agentCapabilities.sessionCapabilities.resume` exists.
    pub async fn agent_supports_session_resume(&self) -> bool {
        self.agent_capabilities
            .read()
            .await
            .as_ref()
            .and_then(|caps| caps.get("sessionCapabilities"))
            .and_then(|sc| sc.get("resume"))
            .is_some()
    }

    /// Check if the upstream agent supports session/close
    ///
    /// Returns true if `agentCapabilities.sessionCapabilities.close` exists.
    pub async fn agent_supports_session_close(&self) -> bool {
        self.agent_capabilities
            .read()
            .await
            .as_ref()
            .and_then(|caps| caps.get("sessionCapabilities"))
            .and_then(|sc| sc.get("close"))
            .is_some()
    }

    /// Mark a session as loaded
    ///
    /// Called after a successful session/load or session/resume. Subsequent
    /// operations on this session can use session/resume instead of session/load.
    pub async fn mark_session_loaded(&self, session_id: &str) {
        let mut loaded = self.loaded_sessions.write().await;
        loaded.insert(session_id.to_string());
    }

    /// Check if a session has been loaded in this connection
    ///
    /// Returns true if the session was previously loaded via session/load or
    /// session/resume. Used to decide whether to load or resume.
    pub async fn is_session_loaded(&self, session_id: &str) -> bool {
        let loaded = self.loaded_sessions.read().await;
        loaded.contains(session_id)
    }

    /// Unmark a session as loaded
    ///
    /// Called when a session is closed. The next access will require a
    /// full session/load again.
    pub async fn unmark_session_loaded(&self, session_id: &str) {
        let mut loaded = self.loaded_sessions.write().await;
        loaded.remove(session_id);
    }
}

impl Default for InternalState {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_state_creation() {
        let state = InternalState::new();

        assert_eq!(state.protocol_version().await, None);
        assert_eq!(state.agent_capabilities().await, None);
        assert_eq!(state.session_count().await, 0);
        assert!(!state.is_initialized().await);
    }

    #[tokio::test]
    async fn test_protocol_version() {
        let state = InternalState::new();

        state.set_protocol_version(1).await;

        assert_eq!(state.protocol_version().await, Some(1));
        assert!(state.is_initialized().await);
    }

    #[tokio::test]
    async fn test_agent_capabilities() {
        let state = InternalState::new();

        let caps = serde_json::json!({
            "streaming": true,
            "tools": ["bash", "edit"]
        });

        state.set_agent_capabilities(caps.clone()).await;

        let retrieved = state.agent_capabilities().await;
        assert_eq!(retrieved, Some(caps));
    }

    #[tokio::test]
    async fn test_session_management() {
        let state = InternalState::new();

        let now = Utc::now();
        let session1 = SessionInfo {
            id: "session-1".to_string(),
            title: Some("Test Session".to_string()),
            cwd: ".".to_string(),
            message_count: 0,
            model: Some("claude-3".to_string()),
            created_at: now,
            last_activity: now,
            status: SessionStatus::Active,
            models: None,
            modes: None,
            config_options: None,
            ownership: SessionOwnership::default(),
        };

        let session2 = SessionInfo {
            id: "session-2".to_string(),
            title: None,
            cwd: "/tmp".to_string(),
            message_count: 5,
            model: None,
            created_at: now,
            last_activity: now,
            status: SessionStatus::Idle,
            models: None,
            modes: None,
            config_options: None,
            ownership: SessionOwnership::default(),
        };

        // Add sessions
        state.upsert_session(session1.clone()).await;
        state.upsert_session(session2.clone()).await;

        assert_eq!(state.session_count().await, 2);

        // Get specific session
        let retrieved = state.get_session("session-1").await;
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().id, "session-1");

        // List all sessions
        let all_sessions = state.list_sessions().await;
        assert_eq!(all_sessions.len(), 2);

        // Remove a session
        state.remove_session("session-1").await;
        assert_eq!(state.session_count().await, 1);
        assert!(state.get_session("session-1").await.is_none());
    }

    #[tokio::test]
    async fn test_session_upsert() {
        let state = InternalState::new();

        let now = Utc::now();
        let session = SessionInfo {
            id: "session-1".to_string(),
            title: Some("Original".to_string()),
            cwd: ".".to_string(),
            message_count: 1,
            model: None,
            created_at: now,
            last_activity: now,
            status: SessionStatus::Active,
            models: None,
            modes: None,
            config_options: None,
            ownership: SessionOwnership::default(),
        };

        // Add session
        state.upsert_session(session.clone()).await;
        assert_eq!(state.session_count().await, 1);

        // Update session (same ID)
        let updated = SessionInfo {
            id: "session-1".to_string(),
            title: Some("Updated".to_string()),
            cwd: ".".to_string(),
            message_count: 2,
            model: Some("claude-3".to_string()),
            created_at: now,
            last_activity: now,
            status: SessionStatus::Processing,
            models: None,
            modes: None,
            config_options: None,
            ownership: SessionOwnership::default(),
        };

        state.upsert_session(updated.clone()).await;
        assert_eq!(state.session_count().await, 1);

        let retrieved = state.get_session("session-1").await.unwrap();
        assert_eq!(retrieved.title, Some("Updated".to_string()));
        assert_eq!(retrieved.message_count, 2);
    }

    #[tokio::test]
    async fn test_update_message_count() {
        let state = InternalState::new();

        let now = Utc::now();
        let session = SessionInfo {
            id: "session-1".to_string(),
            title: None,
            cwd: ".".to_string(),
            message_count: 0,
            model: None,
            created_at: now,
            last_activity: now,
            status: SessionStatus::Active,
            models: None,
            modes: None,
            config_options: None,
            ownership: SessionOwnership::default(),
        };

        state.upsert_session(session).await;

        state.update_message_count("session-1", 10).await;

        let retrieved = state.get_session("session-1").await.unwrap();
        assert_eq!(retrieved.message_count, 10);
    }

    #[tokio::test]
    async fn test_clear_state() {
        let state = InternalState::new();

        // Populate state
        state.set_protocol_version(1).await;
        state.set_agent_capabilities(serde_json::json!({})).await;
        let now = Utc::now();
        state.upsert_session(SessionInfo {
            id: "session-1".to_string(),
            title: None,
            cwd: ".".to_string(),
            message_count: 0,
            model: None,
            created_at: now,
            last_activity: now,
            status: SessionStatus::Active,
            models: None,
            modes: None,
            config_options: None,
            ownership: dirigent_protocol::SessionOwnership::internal(),
        }).await;

        assert!(state.is_initialized().await);
        assert_eq!(state.session_count().await, 1);

        // Clear everything
        state.clear().await;

        assert!(!state.is_initialized().await);
        assert_eq!(state.session_count().await, 0);
        assert_eq!(state.protocol_version().await, None);
        assert_eq!(state.agent_capabilities().await, None);
    }

    #[tokio::test]
    async fn test_concurrent_access() {
        let state = Arc::new(InternalState::new());

        let mut handles = vec![];

        // Spawn 10 tasks that all add sessions
        for i in 0..10 {
            let state_clone = Arc::clone(&state);
            let handle = tokio::spawn(async move {
                let now = Utc::now();
                let session = SessionInfo {
                    id: format!("session-{}", i),
                    title: None,
                    cwd: ".".to_string(),
                    message_count: 0,
                    model: None,
                    created_at: now,
                    last_activity: now,
                    status: SessionStatus::Active,
                    models: None,
                    modes: None,
                    config_options: None,
                    ownership: dirigent_protocol::SessionOwnership::internal(),
                };
                state_clone.upsert_session(session).await;
            });
            handles.push(handle);
        }

        // Wait for all tasks
        for handle in handles {
            handle.await.unwrap();
        }

        // Verify all sessions were added
        assert_eq!(state.session_count().await, 10);
    }

    #[tokio::test]
    async fn test_is_session_active() {
        let state = InternalState::new();
        let now = Utc::now();

        // Active session
        state.upsert_session(SessionInfo {
            id: "active".to_string(),
            title: None,
            cwd: ".".to_string(),
            message_count: 0,
            model: None,
            created_at: now,
            last_activity: now,
            status: SessionStatus::Active,
            models: None,
            modes: None,
            config_options: None,
            ownership: dirigent_protocol::SessionOwnership::internal(),
        }).await;

        // Ended session
        state.upsert_session(SessionInfo {
            id: "ended".to_string(),
            title: None,
            cwd: ".".to_string(),
            message_count: 0,
            model: None,
            created_at: now,
            last_activity: now,
            status: SessionStatus::Ended,
            models: None,
            modes: None,
            config_options: None,
            ownership: SessionOwnership::default(),
        }).await;

        assert!(state.is_session_active("active").await);
        assert!(!state.is_session_active("ended").await);
        assert!(!state.is_session_active("nonexistent").await);
    }

    #[tokio::test]
    async fn test_update_session_status() {
        let state = InternalState::new();
        let now = Utc::now();

        state.upsert_session(SessionInfo {
            id: "session-1".to_string(),
            title: None,
            cwd: ".".to_string(),
            message_count: 0,
            model: None,
            created_at: now,
            last_activity: now,
            status: SessionStatus::Active,
            models: None,
            modes: None,
            config_options: None,
            ownership: dirigent_protocol::SessionOwnership::internal(),
        }).await;

        // Change status to Processing
        state.update_session_status("session-1", SessionStatus::Processing).await;

        let session = state.get_session("session-1").await.unwrap();
        assert_eq!(session.status, SessionStatus::Processing);
        assert!(session.last_activity > now);
    }

    #[tokio::test]
    async fn test_touch_session() {
        let state = InternalState::new();
        let now = Utc::now();

        state.upsert_session(SessionInfo {
            id: "session-1".to_string(),
            title: None,
            cwd: ".".to_string(),
            message_count: 0,
            model: None,
            created_at: now,
            last_activity: now,
            status: SessionStatus::Active,
            models: None,
            modes: None,
            config_options: None,
            ownership: dirigent_protocol::SessionOwnership::internal(),
        }).await;

        // Small delay to ensure timestamp difference
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

        // Touch the session
        state.touch_session("session-1").await;

        let session = state.get_session("session-1").await.unwrap();
        assert!(session.last_activity > now);
    }

    #[tokio::test]
    async fn test_agent_supports_load_session_true() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({
            "loadSession": true
        })).await;
        assert!(state.agent_supports_load_session().await);
    }

    #[tokio::test]
    async fn test_agent_supports_load_session_false() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({
            "loadSession": false
        })).await;
        assert!(!state.agent_supports_load_session().await);
    }

    #[tokio::test]
    async fn test_agent_supports_load_session_missing() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({})).await;
        assert!(!state.agent_supports_load_session().await);
    }

    #[tokio::test]
    async fn test_agent_supports_list_sessions_legacy_true() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({
            "listSessions": true
        })).await;
        assert!(state.agent_supports_list_sessions().await);
    }

    #[tokio::test]
    async fn test_agent_supports_list_sessions_legacy_false() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({
            "listSessions": false
        })).await;
        assert!(!state.agent_supports_list_sessions().await);
    }

    #[tokio::test]
    async fn test_agent_supports_list_sessions_nested() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({
            "sessionCapabilities": {
                "list": {}
            }
        })).await;
        assert!(state.agent_supports_list_sessions().await);
    }

    #[tokio::test]
    async fn test_agent_supports_list_sessions_nested_with_options() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({
            "sessionCapabilities": {
                "list": { "supportsCwd": true }
            }
        })).await;
        assert!(state.agent_supports_list_sessions().await);
    }

    #[tokio::test]
    async fn test_agent_supports_list_sessions_both_formats() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({
            "listSessions": true,
            "sessionCapabilities": {
                "list": {}
            }
        })).await;
        assert!(state.agent_supports_list_sessions().await);
    }

    #[tokio::test]
    async fn test_agent_supports_list_sessions_missing() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({})).await;
        assert!(!state.agent_supports_list_sessions().await);
    }

    #[tokio::test]
    async fn test_agent_supports_list_sessions_no_caps() {
        let state = InternalState::new();
        assert!(!state.agent_supports_list_sessions().await);
    }

    #[tokio::test]
    async fn test_agent_supports_session_resume_true() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({
            "sessionCapabilities": {
                "resume": {}
            }
        })).await;
        assert!(state.agent_supports_session_resume().await);
    }

    #[tokio::test]
    async fn test_agent_supports_session_resume_missing() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({
            "sessionCapabilities": {}
        })).await;
        assert!(!state.agent_supports_session_resume().await);
    }

    #[tokio::test]
    async fn test_agent_supports_session_resume_no_session_caps() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({})).await;
        assert!(!state.agent_supports_session_resume().await);
    }

    #[tokio::test]
    async fn test_agent_supports_session_close_true() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({
            "sessionCapabilities": {
                "close": {}
            }
        })).await;
        assert!(state.agent_supports_session_close().await);
    }

    #[tokio::test]
    async fn test_agent_supports_session_close_missing() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({
            "sessionCapabilities": {}
        })).await;
        assert!(!state.agent_supports_session_close().await);
    }

    #[tokio::test]
    async fn test_agent_supports_session_close_no_session_caps() {
        let state = InternalState::new();
        state.set_agent_capabilities(serde_json::json!({})).await;
        assert!(!state.agent_supports_session_close().await);
    }

    #[tokio::test]
    async fn test_loaded_session_tracking() {
        let state = InternalState::new();

        // Initially not loaded
        assert!(!state.is_session_loaded("sess-1").await);

        // Mark as loaded
        state.mark_session_loaded("sess-1").await;
        assert!(state.is_session_loaded("sess-1").await);
        assert!(!state.is_session_loaded("sess-2").await);

        // Unmark
        state.unmark_session_loaded("sess-1").await;
        assert!(!state.is_session_loaded("sess-1").await);
    }

    #[tokio::test]
    async fn test_loaded_sessions_cleared_on_reconnect() {
        let state = InternalState::new();

        state.mark_session_loaded("sess-1").await;
        state.mark_session_loaded("sess-2").await;
        assert!(state.is_session_loaded("sess-1").await);
        assert!(state.is_session_loaded("sess-2").await);

        // Clear simulates reconnect
        state.clear().await;

        assert!(!state.is_session_loaded("sess-1").await);
        assert!(!state.is_session_loaded("sess-2").await);
    }
}
