//! Initialize method implementation for ACP.
//!
//! This module implements the `initialize` JSON-RPC method, which is the first
//! required interaction in the ACP protocol. Both client and agent must agree
//! on protocol version and exchange capability information.

use crate::acp::connector_state::{
    AgentCapabilities, ClientCapabilities, ImplementationInfo,
};
use crate::acp::transport::{JsonRpcRequest, JsonRpcResult, Transport, TransportError};
use serde::{Deserialize, Serialize};

/// Current ACP protocol version supported by this implementation.
pub const PROTOCOL_VERSION: u32 = 1;

/// Initialize request sent to the agent.
///
/// This is the first message sent after establishing a transport connection.
/// The client advertises its capabilities and requests the agent's capabilities.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct InitializeRequest {
    /// Protocol version supported by the client.
    pub protocol_version: u32,

    /// Capabilities that the client supports.
    pub client_capabilities: ClientCapabilities,

    /// Information about the client implementation.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub client_info: Option<ImplementationInfo>,
}

impl InitializeRequest {
    /// Create a new initialize request with the current protocol version and capabilities.
    pub fn new(
        client_capabilities: ClientCapabilities,
        client_info: Option<ImplementationInfo>,
    ) -> Self {
        Self {
            protocol_version: PROTOCOL_VERSION,
            client_capabilities,
            client_info,
        }
    }

    /// Convert to JSON-RPC request.
    pub fn to_jsonrpc(&self, id: u64) -> JsonRpcRequest {
        JsonRpcRequest::new(
            id,
            "initialize",
            Some(serde_json::to_value(self).expect("Failed to serialize InitializeRequest")),
        )
    }
}

/// Initialize response from the agent.
///
/// The agent responds with its supported protocol version, capabilities,
/// and optional authentication requirements.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct InitializeResponse {
    /// Protocol version supported by the agent.
    pub protocol_version: u32,

    /// Capabilities that the agent supports.
    pub agent_capabilities: AgentCapabilities,

    /// Information about the agent implementation.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub agent_info: Option<ImplementationInfo>,

    /// Authentication methods supported by the agent (empty if no auth required).
    #[serde(default)]
    pub auth_methods: Vec<String>,
}

impl InitializeResponse {
    /// Parse from JSON-RPC result.
    pub fn from_jsonrpc(result: &serde_json::Value) -> Result<Self, serde_json::Error> {
        serde_json::from_value(result.clone())
    }
}

/// Result of initialization attempt.
#[derive(Debug)]
pub enum InitializeResult {
    /// Initialization succeeded with negotiated version and agent capabilities.
    Success {
        protocol_version: u32,
        agent_capabilities: AgentCapabilities,
        agent_info: Option<ImplementationInfo>,
        auth_methods: Vec<String>,
    },
    /// Version mismatch - client and agent don't support compatible versions.
    VersionMismatch {
        client_version: u32,
        agent_version: u32,
    },
    /// Other error during initialization.
    Error(String),
}

/// Perform initialization handshake with the agent.
///
/// This sends an initialize request and processes the response, performing
/// version negotiation and capability exchange.
///
/// # Arguments
///
/// * `transport` - The transport to use for communication
/// * `client_capabilities` - Capabilities that the client supports
/// * `client_info` - Optional client implementation information
///
/// # Returns
///
/// Result of the initialization attempt, including negotiated version and
/// agent capabilities on success.
pub async fn initialize(
    transport: &mut dyn Transport,
    client_capabilities: ClientCapabilities,
    client_info: Option<ImplementationInfo>,
) -> Result<InitializeResult, TransportError> {
    // Create initialize request
    let request = InitializeRequest::new(client_capabilities, client_info);
    let jsonrpc_request = request.to_jsonrpc(1);

    // Send request and await response
    let result = transport.send_request(jsonrpc_request).await?;

    // Handle response
    match result {
        JsonRpcResult::Success(response) => {
            // Parse initialize response
            let init_response = InitializeResponse::from_jsonrpc(&response.result)
                .map_err(|e| TransportError::SerializationError(e))?;

            // Check protocol version compatibility
            if init_response.protocol_version != PROTOCOL_VERSION {
                return Ok(InitializeResult::VersionMismatch {
                    client_version: PROTOCOL_VERSION,
                    agent_version: init_response.protocol_version,
                });
            }

            // Version matches, initialization successful
            Ok(InitializeResult::Success {
                protocol_version: init_response.protocol_version,
                agent_capabilities: init_response.agent_capabilities,
                agent_info: init_response.agent_info,
                auth_methods: init_response.auth_methods,
            })
        }
        JsonRpcResult::Error(error_response) => {
            // Initialization failed with JSON-RPC error
            Ok(InitializeResult::Error(error_response.error.message))
        }
    }
}

/// Validate that client and agent protocol versions are compatible.
///
/// Currently, we require exact version match. In the future, this could
/// be enhanced to support backward compatibility.
pub fn is_version_compatible(client_version: u32, agent_version: u32) -> bool {
    client_version == agent_version
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::acp::connector_state::FsCapabilities;

    #[test]
    fn test_initialize_request_creation() {
        let caps = ClientCapabilities {
            fs: Some(FsCapabilities {
                read_text_file: Some(true),
                write_text_file: Some(false),
            }),
            terminal: Some(false),
            _meta: None,
        };

        let request = InitializeRequest::new(caps.clone(), None);

        assert_eq!(request.protocol_version, PROTOCOL_VERSION);
        assert_eq!(request.client_capabilities, caps);
        assert!(request.client_info.is_none());
    }

    #[test]
    fn test_initialize_request_with_client_info() {
        let caps = ClientCapabilities::default_safe();
        let info = ImplementationInfo {
            name: "dirigent".to_string(),
            title: Some("Dirigent ACP Client".to_string()),
            version: Some("0.1.0".to_string()),
        };

        let request = InitializeRequest::new(caps, Some(info.clone()));

        assert!(request.client_info.is_some());
        assert_eq!(request.client_info.unwrap(), info);
    }

    #[test]
    fn test_initialize_request_serialization() {
        let caps = ClientCapabilities::default_safe();
        let request = InitializeRequest::new(caps, None);

        let json = serde_json::to_string(&request).unwrap();
        let deserialized: InitializeRequest = serde_json::from_str(&json).unwrap();

        assert_eq!(request, deserialized);
    }

    #[test]
    fn test_initialize_request_to_jsonrpc() {
        let caps = ClientCapabilities::default_safe();
        let request = InitializeRequest::new(caps, None);

        let jsonrpc = request.to_jsonrpc(42);

        assert_eq!(jsonrpc.method, "initialize");
        assert_eq!(jsonrpc.id, serde_json::Value::Number(42.into()));
        assert!(jsonrpc.params.is_some());
    }

    #[test]
    fn test_initialize_response_parsing() {
        let json = serde_json::json!({
            "protocol_version": 1,
            "agent_capabilities": {
                "load_session": true,
                "prompt_capabilities": {
                    "image": true,
                    "audio": false,
                    "embedded_context": true
                }
            },
            "agent_info": {
                "name": "test-agent",
                "version": "1.0.0"
            },
            "auth_methods": []
        });

        let response = InitializeResponse::from_jsonrpc(&json).unwrap();

        assert_eq!(response.protocol_version, 1);
        assert!(response.agent_capabilities.load_session.unwrap());
        assert_eq!(response.auth_methods.len(), 0);
    }

    #[test]
    fn test_initialize_response_with_auth() {
        let json = serde_json::json!({
            "protocol_version": 1,
            "agent_capabilities": {},
            "auth_methods": ["api_key", "oauth"]
        });

        let response = InitializeResponse::from_jsonrpc(&json).unwrap();

        assert_eq!(response.auth_methods.len(), 2);
        assert!(response.auth_methods.contains(&"api_key".to_string()));
        assert!(response.auth_methods.contains(&"oauth".to_string()));
    }

    #[test]
    fn test_version_compatibility() {
        assert!(is_version_compatible(1, 1));
        assert!(!is_version_compatible(1, 2));
        assert!(!is_version_compatible(2, 1));
    }

    #[test]
    fn test_initialize_response_missing_optional_fields() {
        let json = serde_json::json!({
            "protocol_version": 1,
            "agent_capabilities": {}
        });

        let response = InitializeResponse::from_jsonrpc(&json).unwrap();

        assert_eq!(response.protocol_version, 1);
        assert!(response.agent_info.is_none());
        assert_eq!(response.auth_methods.len(), 0);
    }
}
