//! ACP connector state management.
//!
//! This module defines the state structure for tracking ACP connection lifecycle,
//! including protocol negotiation, authentication status, and capabilities.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// Connection state for an ACP connector.
///
/// Tracks the lifecycle of an ACP connection from initial connection through
/// initialization, optional authentication, and ready state.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConnectionState {
    /// Not yet initialized (initial state).
    Uninitialized,
    /// Currently performing initialization handshake.
    Initializing,
    /// Initialization complete, ready for optional authentication.
    Initialized,
    /// Currently authenticating.
    Authenticating,
    /// Fully connected and ready for session operations.
    Ready,
    /// Connection lost or transport error.
    Disconnected,
    /// Terminal error state.
    Error(String),
}

/// Complete ACP connector state.
///
/// This structure tracks all state related to an active ACP connection,
/// including protocol version, capabilities, authentication status, and
/// agent information.
#[derive(Debug, Clone)]
pub struct AcpConnectorState {
    /// Current connection state.
    pub connection_state: ConnectionState,

    /// Negotiated protocol version (set after successful initialize).
    pub protocol_version: Option<u32>,

    /// Client capabilities that were advertised to the agent.
    pub client_capabilities: Option<ClientCapabilities>,

    /// Agent capabilities received from the agent.
    pub agent_capabilities: Option<AgentCapabilities>,

    /// Agent implementation info (for display/logging).
    pub agent_info: Option<ImplementationInfo>,

    /// Authentication methods supported by the agent.
    pub auth_methods: Vec<String>,

    /// Whether authentication has been completed successfully.
    pub authenticated: bool,

    /// Active sessions (keyed by session ID).
    ///
    /// Tracks all sessions created through this connector. Each session
    /// maintains its own state including working directory, MCP servers,
    /// current mode, and cancellation status.
    pub sessions: HashMap<String, SessionState>,
}

impl AcpConnectorState {
    /// Create a new uninitialized ACP connector state.
    pub fn new() -> Self {
        Self {
            connection_state: ConnectionState::Uninitialized,
            protocol_version: None,
            client_capabilities: None,
            agent_capabilities: None,
            agent_info: None,
            auth_methods: Vec::new(),
            authenticated: false,
            sessions: HashMap::new(),
        }
    }

    /// Add a new session to the state.
    pub fn add_session(&mut self, session_state: SessionState) {
        self.sessions
            .insert(session_state.session_id.clone(), session_state);
    }

    /// Get a session by ID.
    pub fn get_session(&self, session_id: &str) -> Option<&SessionState> {
        self.sessions.get(session_id)
    }

    /// Get a mutable reference to a session by ID.
    pub fn get_session_mut(&mut self, session_id: &str) -> Option<&mut SessionState> {
        self.sessions.get_mut(session_id)
    }

    /// Remove a session by ID.
    pub fn remove_session(&mut self, session_id: &str) -> Option<SessionState> {
        self.sessions.remove(session_id)
    }

    /// Get all active session IDs.
    pub fn session_ids(&self) -> Vec<String> {
        self.sessions.keys().cloned().collect()
    }

    /// Check if the connector is ready for session operations.
    pub fn is_ready(&self) -> bool {
        matches!(self.connection_state, ConnectionState::Ready)
    }

    /// Check if authentication is required.
    pub fn requires_auth(&self) -> bool {
        !self.auth_methods.is_empty() && !self.authenticated
    }

    /// Transition to error state with message.
    pub fn set_error(&mut self, message: impl Into<String>) {
        self.connection_state = ConnectionState::Error(message.into());
    }
}

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

/// Shared thread-safe wrapper for ACP connector state.
pub type SharedAcpState = Arc<RwLock<AcpConnectorState>>;

/// Client capabilities advertised to the agent.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ClientCapabilities {
    /// Filesystem capabilities.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub fs: Option<FsCapabilities>,

    /// Terminal/command execution support.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub terminal: Option<bool>,

    /// Extension metadata.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub _meta: Option<serde_json::Value>,
}

impl ClientCapabilities {
    /// Create default client capabilities (safe defaults: read-only filesystem, no terminal).
    pub fn default_safe() -> Self {
        Self {
            fs: Some(FsCapabilities {
                read_text_file: Some(true),
                write_text_file: Some(false),
            }),
            terminal: Some(false),
            _meta: None,
        }
    }

    /// Create client capabilities with all features enabled.
    pub fn all_enabled() -> Self {
        Self {
            fs: Some(FsCapabilities {
                read_text_file: Some(true),
                write_text_file: Some(true),
            }),
            terminal: Some(true),
            _meta: None,
        }
    }
}

impl Default for ClientCapabilities {
    /// Default capabilities: CURRENTLY DISABLED (handlers not implemented).
    ///
    /// **IMPORTANT**: These capabilities are currently disabled because the request handlers
    /// in the ACP connector are not yet implemented. Advertising capabilities without
    /// functioning handlers violates ACP compliance and will cause agent requests to fail.
    ///
    /// - **Filesystem Operations** (`fs`):
    ///   - `read_text_file`: Disabled (handler not implemented)
    ///   - `write_text_file`: Disabled (handler not implemented)
    ///
    /// - **Terminal Operations** (`terminal`):
    ///   - Disabled (handler not implemented)
    ///
    /// **Future Implementation**: While the underlying tools exist in `packages/dirigent_tools/`,
    /// the ACP connector's request handler loop does not yet route agent requests to these
    /// tools. Once the handlers are implemented and tested, these capabilities can be enabled.
    ///
    /// **Note**: Use `ClientCapabilities::all_enabled()` to override if you have implemented
    /// custom handlers, but this is NOT recommended until proper request handling is in place.
    fn default() -> Self {
        Self {
            fs: Some(FsCapabilities {
                read_text_file: Some(false),  // Not implemented yet
                write_text_file: Some(false), // Not implemented yet
            }),
            terminal: Some(false),  // Not implemented yet
            _meta: None,
        }
    }
}

/// Filesystem capabilities.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct FsCapabilities {
    /// Support for reading text files.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub read_text_file: Option<bool>,

    /// Support for writing text files.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub write_text_file: Option<bool>,
}

/// Agent capabilities received from the agent.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
pub struct AgentCapabilities {
    /// Support for loading existing sessions.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub load_session: Option<bool>,

    /// Prompt-related capabilities.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prompt_capabilities: Option<PromptCapabilities>,

    /// MCP (Model Context Protocol) server support.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mcp: Option<McpCapabilities>,

    /// Extension metadata.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub _meta: Option<serde_json::Value>,
}

impl AgentCapabilities {
    /// Check if the agent supports loading sessions.
    pub fn supports_load_session(&self) -> bool {
        self.load_session.unwrap_or(false)
    }

    /// Check if the agent supports image content in prompts.
    pub fn supports_image(&self) -> bool {
        self.prompt_capabilities
            .as_ref()
            .and_then(|pc| pc.image)
            .unwrap_or(false)
    }

    /// Check if the agent supports audio content in prompts.
    pub fn supports_audio(&self) -> bool {
        self.prompt_capabilities
            .as_ref()
            .and_then(|pc| pc.audio)
            .unwrap_or(false)
    }

    /// Check if the agent supports embedded context (resource blocks).
    pub fn supports_embedded_context(&self) -> bool {
        self.prompt_capabilities
            .as_ref()
            .and_then(|pc| pc.embedded_context)
            .unwrap_or(false)
    }
}

/// Prompt-related capabilities.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
pub struct PromptCapabilities {
    /// Support for image content blocks.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub image: Option<bool>,

    /// Support for audio content blocks.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audio: Option<bool>,

    /// Support for embedded context (resource blocks).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub embedded_context: Option<bool>,
}

/// MCP server capabilities.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct McpCapabilities {
    /// Support for HTTP transport MCP servers.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub http: Option<bool>,

    /// Support for SSE transport MCP servers.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub sse: Option<bool>,
}

/// Implementation information (for client or agent).
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ImplementationInfo {
    /// Implementation name (e.g., "dirigent", "claude").
    pub name: String,

    /// Human-readable title.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub title: Option<String>,

    /// Version string (e.g., "0.1.0").
    #[serde(skip_serializing_if = "Option::is_none")]
    pub version: Option<String>,
}

/// Session state tracking.
///
/// Tracks the state of a single session including working directory,
/// MCP servers, current mode, and cancellation status.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionState {
    /// Session ID.
    pub session_id: String,
    /// Working directory.
    pub cwd: String,
    /// MCP servers configured for this session.
    pub mcp_servers: Vec<McpServer>,
    /// Current mode (e.g., "code", "chat"), if applicable.
    pub current_mode: Option<String>,
    /// Current model, if applicable.
    pub current_model: Option<String>,
    /// Whether a prompt is currently in progress.
    pub prompt_in_progress: bool,
    /// Whether cancellation is in progress.
    pub cancelling: bool,
    /// Whether this session is currently loading (replaying history).
    pub loading: bool,
}

impl SessionState {
    /// Create a new session state.
    pub fn new(session_id: String, cwd: String, mcp_servers: Vec<McpServer>) -> Self {
        Self {
            session_id,
            cwd,
            mcp_servers,
            current_mode: None,
            current_model: None,
            prompt_in_progress: false,
            cancelling: false,
            loading: false,
        }
    }

    /// Check if the session is ready for new prompts.
    pub fn is_ready(&self) -> bool {
        !self.prompt_in_progress && !self.cancelling && !self.loading
    }

    /// Mark a prompt as started.
    pub fn start_prompt(&mut self) {
        self.prompt_in_progress = true;
    }

    /// Mark a prompt as completed.
    pub fn complete_prompt(&mut self) {
        self.prompt_in_progress = false;
        self.cancelling = false;
    }

    /// Initiate cancellation.
    pub fn start_cancellation(&mut self) {
        self.cancelling = true;
    }

    /// Mark the session as loading (history replay).
    pub fn start_loading(&mut self) {
        self.loading = true;
    }

    /// Mark loading as complete.
    pub fn complete_loading(&mut self) {
        self.loading = false;
    }

    /// Update the current mode.
    pub fn set_mode(&mut self, mode: String) {
        self.current_mode = Some(mode);
    }

    /// Update the current model.
    pub fn set_model(&mut self, model: String) {
        self.current_model = Some(model);
    }
}

/// MCP (Model Context Protocol) server configuration.
///
/// Specifies how to connect to an MCP server that provides additional
/// tools and context to the agent.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum McpServer {
    /// Stdio transport - spawn a process and communicate via stdin/stdout.
    Stdio {
        /// Name of the MCP server (for display).
        name: String,
        /// Command to execute.
        command: String,
        /// Command-line arguments.
        args: Vec<String>,
        /// Environment variables.
        env: Vec<EnvVariable>,
    },
    /// HTTP transport - connect to an HTTP endpoint.
    Http {
        /// Name of the MCP server (for display).
        name: String,
        /// Base URL of the HTTP endpoint.
        url: String,
        /// HTTP headers to include in requests.
        headers: Vec<HttpHeader>,
    },
    /// SSE (Server-Sent Events) transport - connect to an SSE endpoint.
    Sse {
        /// Name of the MCP server (for display).
        name: String,
        /// URL of the SSE endpoint.
        url: String,
        /// HTTP headers to include in connection.
        headers: Vec<HttpHeader>,
    },
}

/// Environment variable for MCP server.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct EnvVariable {
    pub name: String,
    pub value: String,
}

/// HTTP header for MCP server.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct HttpHeader {
    pub name: String,
    pub value: String,
}

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

    #[test]
    fn test_connection_state_transitions() {
        let mut state = AcpConnectorState::new();
        assert_eq!(state.connection_state, ConnectionState::Uninitialized);
        assert!(!state.is_ready());

        state.connection_state = ConnectionState::Initializing;
        assert!(!state.is_ready());

        state.connection_state = ConnectionState::Initialized;
        assert!(!state.is_ready());

        state.connection_state = ConnectionState::Ready;
        assert!(state.is_ready());
    }

    #[test]
    fn test_requires_auth() {
        let mut state = AcpConnectorState::new();
        assert!(!state.requires_auth());

        state.auth_methods = vec!["api_key".to_string()];
        assert!(state.requires_auth());

        state.authenticated = true;
        assert!(!state.requires_auth());
    }

    #[test]
    fn test_client_capabilities_default_safe() {
        let caps = ClientCapabilities::default_safe();
        assert!(caps.fs.is_some());
        assert_eq!(caps.fs.as_ref().unwrap().read_text_file, Some(true));
        assert_eq!(caps.fs.as_ref().unwrap().write_text_file, Some(false));
        assert_eq!(caps.terminal, Some(false));
    }

    #[test]
    fn test_client_capabilities_all_enabled() {
        let caps = ClientCapabilities::all_enabled();
        assert!(caps.fs.is_some());
        assert_eq!(caps.fs.as_ref().unwrap().read_text_file, Some(true));
        assert_eq!(caps.fs.as_ref().unwrap().write_text_file, Some(true));
        assert_eq!(caps.terminal, Some(true));
    }

    #[test]
    fn test_agent_capabilities_checks() {
        let caps = AgentCapabilities {
            load_session: Some(true),
            prompt_capabilities: Some(PromptCapabilities {
                image: Some(true),
                audio: Some(false),
                embedded_context: Some(true),
            }),
            mcp: None,
            _meta: None,
        };

        assert!(caps.supports_load_session());
        assert!(caps.supports_image());
        assert!(!caps.supports_audio());
        assert!(caps.supports_embedded_context());
    }

    #[test]
    fn test_serialization() {
        let caps = ClientCapabilities {
            fs: Some(FsCapabilities {
                read_text_file: Some(true),
                write_text_file: Some(true),
            }),
            terminal: Some(true),
            _meta: Some(serde_json::json!({"custom": "data"})),
        };

        let json = serde_json::to_string(&caps).unwrap();
        let deserialized: ClientCapabilities = serde_json::from_str(&json).unwrap();

        assert_eq!(caps, deserialized);
    }

    #[test]
    fn test_implementation_info() {
        let info = ImplementationInfo {
            name: "test-impl".to_string(),
            title: Some("Test Implementation".to_string()),
            version: Some("1.0.0".to_string()),
        };

        let json = serde_json::to_string(&info).unwrap();
        let deserialized: ImplementationInfo = serde_json::from_str(&json).unwrap();

        assert_eq!(info, deserialized);
    }
}
