//! JSON-RPC protocol handler for ACP
//!
//! This module provides request/response correlation, timeout handling, and
//! notification routing for JSON-RPC 2.0 messages used in the Agent-Client Protocol.
//!
//! # Architecture
//!
//! The `ProtocolHandler` maintains a registry of pending requests and correlates
//! incoming responses by request ID. It distinguishes between:
//!
//! - **Requests**: Messages with an `id` field that expect a response
//! - **Responses**: Messages with an `id` field and a `result` or `error` field
//! - **Notifications**: Messages with a `method` field but no `id` (server-initiated)
//!
//! # Request/Response Correlation
//!
//! When sending a request:
//! 1. Generate unique request ID using atomic counter
//! 2. Create oneshot channel for response
//! 3. Store channel sender in pending requests map
//! 4. Send request to transport
//! 5. Wait for response with timeout
//!
//! When receiving a message:
//! 1. Check if message has `id` (response) or only `method` (notification)
//! 2. For responses: Look up pending request by ID, send response via oneshot
//! 3. For notifications: Route to notification channel
//!
//! # Example
//!
//! ```no_run
//! use dirigent_core::connectors::acp::protocol::ProtocolHandler;
//! use serde_json::json;
//!
//! # async fn example() -> anyhow::Result<()> {
//! let mut handler = ProtocolHandler::new();
//!
//! // Build and send a request
//! let request = json!({
//!     "jsonrpc": "2.0",
//!     "method": "initialize",
//!     "params": {
//!         "protocolVersion": 1,
//!         "clientCapabilities": {}
//!     }
//! });
//!
//! let response = handler.send_request(request).await?;
//! println!("Response: {:?}", response);
//! # Ok(())
//! # }
//! ```

use serde_json::{json, Value};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use tokio::sync::{mpsc, oneshot, Mutex};
use tokio::time::{timeout, Duration};
use tracing::{debug, info, warn};

/// Default timeout for requests (30 seconds)
const DEFAULT_TIMEOUT_MS: u64 = 30_000;

/// Request ID type (JSON-RPC 2.0 allows string, number, or null)
pub type RequestId = Value;

/// Result type for protocol operations
pub type ProtocolResult<T> = Result<T, ProtocolError>;

/// Protocol handler errors
#[derive(Debug, thiserror::Error)]
pub enum ProtocolError {
    /// Request timed out waiting for response
    #[error("Request timeout after {timeout_ms}ms")]
    Timeout { timeout_ms: u64 },

    /// Received error response from server
    #[error("JSON-RPC error: {message} (code: {code})")]
    JsonRpcError {
        code: i64,
        message: String,
        data: Option<Value>,
    },

    /// Invalid message format
    #[error("Invalid message: {0}")]
    InvalidMessage(String),

    /// Internal error
    #[error("Internal error: {0}")]
    Internal(String),
}

/// Result type for message handling
///
/// Represents the different outcomes when handling an incoming message from the agent.
/// This allows the protocol handler to distinguish between notifications (no response),
/// immediate responses, and agent requests that need external input.
#[derive(Debug)]
pub enum MessageHandlerResult {
    /// No response needed (notification was processed)
    None,

    /// Immediate response to send back to the agent
    Response(Value),

    /// Agent-initiated request requiring external input (e.g., permission prompt)
    ///
    /// The connector should emit an Event::AgentRequest and wait for the client's
    /// response before sending anything back to the agent.
    AgentRequest {
        /// Request ID from the agent (for correlating the response)
        request_id: Value,
        /// Method being requested (e.g., "session/request_permission")
        method: String,
        /// Request parameters
        params: Value,
    },
}

/// JSON-RPC protocol handler
///
/// Manages request/response correlation, timeouts, and notification routing.
///
/// # Thread Safety
///
/// ProtocolHandler is thread-safe and can be shared across tasks using `Arc`.
/// All internal state is protected by async-aware locks.
pub struct ProtocolHandler {
    /// Atomic counter for generating unique request IDs
    next_request_id: Arc<AtomicU64>,

    /// Map of pending requests waiting for responses
    ///
    /// Key: Request ID
    /// Value: Oneshot sender to deliver the response
    pending_requests: Arc<Mutex<HashMap<RequestId, oneshot::Sender<Value>>>>,

    /// Channel for routing notifications (messages without request ID)
    notification_tx: mpsc::UnboundedSender<Value>,

    /// Receiver for notifications (taken by consumer)
    notification_rx: Arc<Mutex<Option<mpsc::UnboundedReceiver<Value>>>>,

    /// Default timeout for requests
    timeout_ms: u64,
}

impl ProtocolHandler {
    /// Create a new protocol handler
    ///
    /// # Example
    ///
    /// ```
    /// use dirigent_core::connectors::acp::protocol::ProtocolHandler;
    ///
    /// let handler = ProtocolHandler::new();
    /// ```
    pub fn new() -> Self {
        Self::with_timeout(DEFAULT_TIMEOUT_MS)
    }

    /// Create a new protocol handler with custom timeout
    ///
    /// # Arguments
    ///
    /// * `timeout_ms` - Request timeout in milliseconds
    ///
    /// # Example
    ///
    /// ```
    /// use dirigent_core::connectors::acp::protocol::ProtocolHandler;
    ///
    /// let handler = ProtocolHandler::with_timeout(60_000); // 60 second timeout
    /// ```
    pub fn with_timeout(timeout_ms: u64) -> Self {
        let (notification_tx, notification_rx) = mpsc::unbounded_channel();

        Self {
            next_request_id: Arc::new(AtomicU64::new(1)),
            pending_requests: Arc::new(Mutex::new(HashMap::new())),
            notification_tx,
            notification_rx: Arc::new(Mutex::new(Some(notification_rx))),
            timeout_ms,
        }
    }

    /// Generate a unique request ID
    ///
    /// Uses an atomic counter to ensure uniqueness across concurrent requests.
    fn generate_request_id(&self) -> RequestId {
        let id = self.next_request_id.fetch_add(1, Ordering::SeqCst);
        json!(id)
    }

    /// Add request ID to a message and prepare for response
    ///
    /// Returns the modified message and a oneshot receiver for the response.
    ///
    /// # Arguments
    ///
    /// * `message` - JSON-RPC request (without `id` field)
    ///
    /// # Returns
    ///
    /// Tuple of (message with ID, response receiver)
    pub async fn prepare_request(&self, mut message: Value) -> (Value, oneshot::Receiver<Value>) {
        let request_id = self.generate_request_id();

        // Add ID to message
        if let Some(obj) = message.as_object_mut() {
            obj.insert("id".to_string(), request_id.clone());
        }

        // Create oneshot channel for response
        let (tx, rx) = oneshot::channel();

        // Store in pending requests
        self.pending_requests.lock().await.insert(request_id, tx);

        (message, rx)
    }

    /// Send a request and wait for response with timeout
    ///
    /// This is a placeholder that will be integrated with the transport layer.
    /// For now, it demonstrates the correlation logic that will be used.
    ///
    /// # Arguments
    ///
    /// * `message` - JSON-RPC request (without `id` field)
    ///
    /// # Returns
    ///
    /// Response value or error
    ///
    /// # Errors
    ///
    /// - `ProtocolError::Timeout`: Request timed out
    /// - `ProtocolError::JsonRpcError`: Server returned error response
    /// - `ProtocolError::Internal`: Internal error occurred
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use dirigent_core::connectors::acp::protocol::ProtocolHandler;
    /// # use serde_json::json;
    /// # async fn example() -> anyhow::Result<()> {
    /// let handler = ProtocolHandler::new();
    ///
    /// let request = json!({
    ///     "jsonrpc": "2.0",
    ///     "method": "initialize",
    ///     "params": { "protocolVersion": 1 }
    /// });
    ///
    /// let response = handler.send_request(request).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn send_request(&self, message: Value) -> ProtocolResult<Value> {
        let (message_with_id, response_rx) = self.prepare_request(message).await;

        // Extract the request ID for cleanup on timeout
        let request_id = message_with_id
            .get("id")
            .cloned()
            .unwrap_or(json!(null));

        // In a real implementation, this would send via transport
        // For now, we just demonstrate the correlation logic
        debug!(request_id = %request_id, "Request prepared (transport send would happen here)");

        // Wait for response with timeout
        match timeout(Duration::from_millis(self.timeout_ms), response_rx).await {
            Ok(Ok(response)) => {
                // Response received successfully
                self.handle_response_result(response)
            }
            Ok(Err(_)) => {
                // Oneshot channel was dropped (should not happen)
                Err(ProtocolError::Internal(
                    "Response channel dropped".to_string(),
                ))
            }
            Err(_) => {
                // Timeout occurred
                // Clean up pending request
                self.pending_requests.lock().await.remove(&request_id);

                Err(ProtocolError::Timeout {
                    timeout_ms: self.timeout_ms,
                })
            }
        }
    }

    /// Handle incoming message from transport
    ///
    /// Routes the message based on its type:
    /// - Request (has `method` and `id`): Return AgentRequest or Response
    /// - Response (has `id` and `result`/`error`): Route to pending request, return None
    /// - Notification (has `method` but no `id`): Route to notification channel, return None
    ///
    /// # Arguments
    ///
    /// * `message` - JSON-RPC message from transport
    ///
    /// # Returns
    ///
    /// - `MessageHandlerResult::None`: Message processed, no response needed
    /// - `MessageHandlerResult::Response(val)`: Immediate response to send to agent
    /// - `MessageHandlerResult::AgentRequest`: Request needs external input (emit event)
    ///
    /// # Example
    ///
    /// ```
    /// # use dirigent_core::connectors::acp::protocol::{ProtocolHandler, MessageHandlerResult};
    /// # use serde_json::json;
    /// # async fn example() {
    /// let handler = ProtocolHandler::new();
    ///
    /// // Handle a response
    /// let response = json!({
    ///     "jsonrpc": "2.0",
    ///     "id": 1,
    ///     "result": { "status": "ok" }
    /// });
    /// let result = handler.handle_message(response).await;
    /// // Returns MessageHandlerResult::None
    ///
    /// // Handle a notification
    /// let notification = json!({
    ///     "jsonrpc": "2.0",
    ///     "method": "session/update",
    ///     "params": { "sessionId": "123" }
    /// });
    /// let result = handler.handle_message(notification).await;
    /// // Returns MessageHandlerResult::None
    /// # }
    /// ```
    pub async fn handle_message(&self, message: Value) -> MessageHandlerResult {
        // Determine message type according to JSON-RPC 2.0 spec:
        // - Request: has "method" AND "id" (expects response)
        // - Notification: has "method" but NO "id" (no response expected)
        // - Response: has "result" or "error" AND "id", NO "method" (answer to our request)

        let has_id = message.get("id").is_some();
        let has_method = message.get("method").is_some();
        let has_result_or_error = message.get("result").is_some() || message.get("error").is_some();

        if has_method && has_id {
            // This is an INCOMING REQUEST from the agent - it expects a response from us
            let method = message.get("method").and_then(|m| m.as_str()).unwrap_or("<unknown>");
            let request_id = message.get("id").cloned().unwrap();

            info!(
                method = method,
                request_id = %request_id,
                "📨 Received request from agent (agent→client)"
            );

            // Handle known request types
            match method {
                "session/request_permission" => {
                    // Return agent request for external handling (permission UI in client)
                    info!(
                        request_id = %request_id,
                        "🔔 Permission request from agent - forwarding to client for user input"
                    );

                    let params = message.get("params").cloned().unwrap_or(json!({}));

                    return MessageHandlerResult::AgentRequest {
                        request_id,
                        method: method.to_string(),
                        params,
                    };
                }
                _ => {
                    warn!(
                        method = method,
                        request_id = %request_id,
                        "⚠️ Received unsupported request from agent, returning error"
                    );
                    // Return error response for unsupported methods
                    return MessageHandlerResult::Response(json!({
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "error": {
                            "code": -32601,
                            "message": format!("Method not supported: {}", method)
                        }
                    }));
                }
            }
        } else if has_id && has_result_or_error {
            // This is a RESPONSE to one of OUR requests
            let request_id = message.get("id").cloned().unwrap();

            let sender = self.pending_requests.lock().await.remove(&request_id);

            if let Some(tx) = sender {
                // Send response to waiting request
                if tx.send(message).is_err() {
                    warn!(
                        request_id = %request_id,
                        "Failed to deliver response (receiver dropped)"
                    );
                }
            } else {
                // Unknown request ID (could be late response after timeout)
                warn!(
                    request_id = %request_id,
                    "Received response for unknown request ID (possible timeout)"
                );
            }
        } else if has_method {
            // This is a NOTIFICATION - route to notification channel
            let method = message.get("method").and_then(|m| m.as_str()).unwrap_or("<unknown>");

            // Extract session_id from params if available (for session/update notifications)
            let session_id = message.get("params")
                .and_then(|p| p.get("sessionId"))
                .and_then(|s| s.as_str());

            // Use masked version for logging
            let masked_message = dirigent_protocol::log_utils::format_for_log(&message);

            if let Some(sid) = session_id {
                info!(method = method, session_id = sid, message = %masked_message, "🔔 Protocol handler received notification, routing to channel");
            } else {
                info!(method = method, message = %masked_message, "🔔 Protocol handler received notification, routing to channel");
            }

            match self.notification_tx.send(message.clone()) {
                Ok(_) => {
                    info!(method = method, "✅ Successfully forwarded notification to channel");
                }
                Err(e) => {
                    warn!(
                        error = %e,
                        method = method,
                        "⚠️ Failed to route notification (no receivers)"
                    );
                }
            }
        } else {
            // Invalid message format
            warn!(
                message = %message,
                "Received message with neither valid request, response, nor notification format"
            );
        }

        MessageHandlerResult::None
    }

    /// Cancel all pending requests with a structured error response.
    ///
    /// Called before the protocol handler is dropped when the transport dies,
    /// so that spawned tasks receive a meaningful error response instead of
    /// a bare `RecvError` from the dropped oneshot sender.
    pub async fn cancel_all_pending(&self, reason: &str) {
        let mut pending = self.pending_requests.lock().await;
        let count = pending.len();
        if count > 0 {
            warn!(count, reason, "Cancelling pending requests");
            for (id, sender) in pending.drain() {
                debug!(request_id = %id, "Cancelling pending request");
                let _ = sender.send(serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": id.to_string(),
                    "error": {
                        "code": -32000,
                        "message": reason
                    }
                }));
            }
        }
    }

    /// Get the number of currently pending requests.
    pub async fn pending_request_count(&self) -> usize {
        self.pending_requests.lock().await.len()
    }

    /// Take the notification receiver
    ///
    /// This can only be called once. Subsequent calls will return None.
    ///
    /// # Returns
    ///
    /// The notification receiver, or None if already taken
    ///
    /// # Example
    ///
    /// ```
    /// # use dirigent_core::connectors::acp::protocol::ProtocolHandler;
    /// # async fn example() {
    /// let handler = ProtocolHandler::new();
    ///
    /// let mut notifications = handler.take_notification_receiver().await.unwrap();
    ///
    /// // Process notifications
    /// while let Some(notif) = notifications.recv().await {
    ///     println!("Notification: {:?}", notif);
    /// }
    /// # }
    /// ```
    pub async fn take_notification_receiver(&self) -> Option<mpsc::UnboundedReceiver<Value>> {
        self.notification_rx.lock().await.take()
    }

    /// Handle a response value, checking for JSON-RPC errors
    ///
    /// # Arguments
    ///
    /// * `response` - Full JSON-RPC response message
    ///
    /// # Returns
    ///
    /// The result value or error
    fn handle_response_result(&self, response: Value) -> ProtocolResult<Value> {
        // Check for error field
        if let Some(error) = response.get("error") {
            let code = error.get("code").and_then(|c| c.as_i64()).unwrap_or(-1);
            let message = error
                .get("message")
                .and_then(|m| m.as_str())
                .unwrap_or("Unknown error")
                .to_string();
            let data = error.get("data").cloned();

            return Err(ProtocolError::JsonRpcError {
                code,
                message,
                data,
            });
        }

        // Return the full response (caller can extract result field)
        Ok(response)
    }
}

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

// Helper functions for building ACP JSON-RPC messages

/// Build an initialize request
///
/// # Arguments
///
/// * `protocol_version` - ACP protocol version to use
/// * `capabilities` - Client capabilities object
///
/// # Returns
///
/// JSON-RPC initialize request (without `id` field)
///
/// # Example
///
/// ```
/// use dirigent_core::connectors::acp::protocol::build_initialize_request;
/// use serde_json::json;
///
/// let request = build_initialize_request(1, json!({}));
/// assert_eq!(request["method"], "initialize");
/// assert_eq!(request["params"]["protocolVersion"], 1);
/// ```
pub fn build_initialize_request(protocol_version: u32, capabilities: Value) -> Value {
    json!({
        "jsonrpc": "2.0",
        "method": "initialize",
        "params": {
            "protocolVersion": protocol_version,
            "clientCapabilities": capabilities
        }
    })
}

/// Build a session/new request
///
/// # Arguments
///
/// * `cwd` - Current working directory for the session
/// * `mcp_servers` - Optional array of MCP server configurations
///
/// # Returns
///
/// JSON-RPC session/new request (without `id` field)
///
/// # Example
///
/// ```
/// use dirigent_core::connectors::acp::protocol::build_session_new_request;
/// use serde_json::json;
///
/// let request = build_session_new_request(".", None);
/// assert_eq!(request["method"], "session/new");
/// assert_eq!(request["params"]["cwd"], ".");
/// ```
pub fn build_session_new_request(cwd: impl Into<String>, mcp_servers: Option<Value>) -> Value {
    json!({
        "jsonrpc": "2.0",
        "method": "session/new",
        "params": {
            "cwd": cwd.into(),
            "mcpServers": mcp_servers.unwrap_or(json!([]))
        }
    })
}

/// Build a session/prompt request
///
/// # Arguments
///
/// * `session_id` - Session ID to send prompt to
/// * `prompt` - Prompt array (content blocks)
///
/// # Returns
///
/// JSON-RPC session/prompt request (without `id` field)
///
/// # Example
///
/// ```
/// use dirigent_core::connectors::acp::protocol::build_session_prompt_request;
/// use serde_json::json;
///
/// let prompt = json!([
///     {
///         "type": "text",
///         "text": "Hello, world!"
///     }
/// ]);
///
/// let request = build_session_prompt_request("session-123", prompt);
/// assert_eq!(request["method"], "session/prompt");
/// assert_eq!(request["params"]["sessionId"], "session-123");
/// ```
pub fn build_session_prompt_request(session_id: impl Into<String>, prompt: Value) -> Value {
    json!({
        "jsonrpc": "2.0",
        "method": "session/prompt",
        "params": {
            "sessionId": session_id.into(),
            "prompt": prompt
        }
    })
}

/// Build a session/cancel request
///
/// # Arguments
///
/// * `session_id` - Session ID to cancel
///
/// # Returns
///
/// JSON-RPC session/cancel request (without `id` field)
///
/// # Example
///
/// ```
/// use dirigent_core::connectors::acp::protocol::build_session_cancel_request;
///
/// let request = build_session_cancel_request("session-123");
/// assert_eq!(request["method"], "session/cancel");
/// assert_eq!(request["params"]["sessionId"], "session-123");
/// ```
/// Build a session/load request
///
/// Loads an existing session and replays its history (if the agent supports it).
/// Per the ACP spec, `sessionId`, `cwd`, and `mcpServers` are all required fields.
///
/// # Arguments
///
/// * `session_id` - The ID of the session to load
/// * `cwd` - The working directory for the session
/// * `mcp_servers` - Optional MCP server configuration (defaults to empty array)
///
/// # Returns
///
/// A JSON-RPC 2.0 request message for session/load
///
/// # Example
///
/// ```
/// use dirigent_core::connectors::acp::protocol::build_session_load_request;
/// use serde_json::json;
///
/// let request = build_session_load_request("session-123", "/home/user", None);
/// assert_eq!(request["method"], "session/load");
/// assert_eq!(request["params"]["sessionId"], "session-123");
/// assert_eq!(request["params"]["cwd"], "/home/user");
/// ```
pub fn build_session_load_request(
    session_id: impl Into<String>,
    cwd: impl Into<String>,
    mcp_servers: Option<Value>,
) -> Value {
    json!({
        "jsonrpc": "2.0",
        "method": "session/load",
        "params": {
            "sessionId": session_id.into(),
            "cwd": cwd.into(),
            "mcpServers": mcp_servers.unwrap_or(json!([]))
        }
    })
}

pub fn build_session_cancel_request(session_id: impl Into<String>) -> Value {
    json!({
        "jsonrpc": "2.0",
        "method": "session/cancel",
        "params": {
            "sessionId": session_id.into()
        }
    })
}

/// Build a session/list request
///
/// Lists sessions available from the agent. Supports optional filtering
/// by working directory and cursor-based pagination.
///
/// # Arguments
///
/// * `cwd` - Optional working directory filter
/// * `cursor` - Optional pagination cursor from previous response
///
/// # Returns
///
/// JSON-RPC session/list request (without `id` field)
///
/// # Example
///
/// ```
/// use dirigent_core::connectors::acp::protocol::build_session_list_request;
///
/// let request = build_session_list_request(None, None);
/// assert_eq!(request["method"], "session/list");
/// ```
pub fn build_session_list_request(
    cwd: Option<&str>,
    cursor: Option<&str>,
) -> Value {
    let mut params = serde_json::Map::new();
    if let Some(cwd) = cwd {
        params.insert("cwd".to_string(), Value::String(cwd.to_string()));
    }
    if let Some(cursor) = cursor {
        params.insert("cursor".to_string(), Value::String(cursor.to_string()));
    }
    json!({
        "jsonrpc": "2.0",
        "method": "session/list",
        "params": params
    })
}

/// Build a session/resume request
///
/// Resumes an existing session without replaying history. The agent restores
/// internal context but does not send `session/update` history notifications.
///
/// # Arguments
///
/// * `session_id` - The session to resume
/// * `cwd` - Optional working directory
/// * `mcp_servers` - Optional MCP server configuration
///
/// # Example
///
/// ```
/// use dirigent_core::connectors::acp::protocol::build_session_resume_request;
///
/// let request = build_session_resume_request("sess-123", None, None);
/// assert_eq!(request["method"], "session/resume");
/// ```
pub fn build_session_resume_request(
    session_id: impl Into<String>,
    cwd: Option<&str>,
    mcp_servers: Option<Value>,
) -> Value {
    let mut params = serde_json::Map::new();
    params.insert("sessionId".to_string(), Value::String(session_id.into()));
    if let Some(cwd) = cwd {
        params.insert("cwd".to_string(), Value::String(cwd.to_string()));
    }
    if let Some(servers) = mcp_servers {
        params.insert("mcpServers".to_string(), servers);
    }
    json!({
        "jsonrpc": "2.0",
        "method": "session/resume",
        "params": params
    })
}

/// Build a session/close JSON-RPC request.
///
/// Tells the agent to release resources for this session. The session
/// remains in session/list and can be session/load-ed again later.
///
/// # Arguments
///
/// * `session_id` - The session to close
///
/// # Example
///
/// ```
/// use dirigent_core::connectors::acp::protocol::build_session_close_request;
///
/// let request = build_session_close_request("session-abc");
/// assert_eq!(request["method"], "session/close");
/// ```
pub fn build_session_close_request(session_id: impl Into<String>) -> Value {
    json!({
        "jsonrpc": "2.0",
        "method": "session/close",
        "params": {
            "sessionId": session_id.into()
        }
    })
}

/// Build a session/set_config_option JSON-RPC request.
///
/// Sets a configuration option for a session. The response contains
/// the complete updated configOptions array.
///
/// # Arguments
///
/// * `session_id` - The session to update
/// * `config_id` - The config option identifier (e.g., "mode", "model")
/// * `value` - The new value to set
///
/// # Example
///
/// ```
/// use dirigent_core::connectors::acp::protocol::build_set_config_option_request;
///
/// let request = build_set_config_option_request("session-abc", "mode", "code");
/// assert_eq!(request["method"], "session/set_config_option");
/// assert_eq!(request["params"]["configId"], "mode");
/// assert_eq!(request["params"]["value"], "code");
/// ```
pub fn build_set_config_option_request(
    session_id: impl Into<String>,
    config_id: impl Into<String>,
    value: impl Into<String>,
) -> Value {
    json!({
        "jsonrpc": "2.0",
        "method": "session/set_config_option",
        "params": {
            "sessionId": session_id.into(),
            "configId": config_id.into(),
            "value": value.into()
        }
    })
}

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

    #[tokio::test]
    async fn test_request_id_generation() {
        let handler = ProtocolHandler::new();

        let id1 = handler.generate_request_id();
        let id2 = handler.generate_request_id();
        let id3 = handler.generate_request_id();

        // IDs should be unique
        assert_ne!(id1, id2);
        assert_ne!(id2, id3);
        assert_ne!(id1, id3);

        // IDs should be numbers
        assert!(id1.is_u64());
        assert!(id2.is_u64());
        assert!(id3.is_u64());
    }

    #[tokio::test]
    async fn test_request_id_uniqueness_concurrent() {
        let handler = Arc::new(ProtocolHandler::new());
        let mut handles = vec![];

        // Generate 1000 IDs concurrently
        for _ in 0..1000 {
            let h = Arc::clone(&handler);
            let handle = tokio::spawn(async move { h.generate_request_id() });
            handles.push(handle);
        }

        let mut ids = vec![];
        for handle in handles {
            ids.push(handle.await.unwrap());
        }

        // Check all IDs are unique
        let unique_count = ids.iter().collect::<std::collections::HashSet<_>>().len();
        assert_eq!(unique_count, 1000, "All IDs should be unique");
    }

    #[tokio::test]
    async fn test_handle_response_with_matching_id() {
        let handler = ProtocolHandler::new();

        // Prepare a request
        let request = json!({
            "jsonrpc": "2.0",
            "method": "test",
            "params": {}
        });

        let (message_with_id, response_rx) = handler.prepare_request(request).await;
        let request_id = message_with_id.get("id").cloned().unwrap();

        // Simulate receiving a response
        let response = json!({
            "jsonrpc": "2.0",
            "id": request_id,
            "result": { "status": "ok" }
        });

        handler.handle_message(response.clone()).await;

        // Verify response was delivered
        let received = response_rx.await.unwrap();
        assert_eq!(received, response);
    }

    #[tokio::test]
    async fn test_handle_notification_routing() {
        let handler = ProtocolHandler::new();

        let mut notifications = handler.take_notification_receiver().await.unwrap();

        // Send a notification
        let notification = json!({
            "jsonrpc": "2.0",
            "method": "session/update",
            "params": { "sessionId": "123" }
        });

        handler.handle_message(notification.clone()).await;

        // Verify notification was routed
        let received = notifications.recv().await.unwrap();
        assert_eq!(received, notification);
    }

    #[tokio::test]
    async fn test_timeout_cleans_up_pending_request() {
        let handler = ProtocolHandler::with_timeout(100); // 100ms timeout

        // Prepare request
        let request = json!({
            "jsonrpc": "2.0",
            "method": "test",
            "params": {}
        });

        let (message_with_id, _response_rx) = handler.prepare_request(request).await;
        let request_id = message_with_id.get("id").cloned().unwrap();

        // Verify request is pending
        assert!(handler
            .pending_requests
            .lock()
            .await
            .contains_key(&request_id));

        // Wait for timeout
        sleep(Duration::from_millis(150)).await;

        // Simulate timeout by trying to send response (should be cleaned up)
        handler
            .handle_message(json!({
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {}
            }))
            .await;

        // Request should still be in pending (timeout cleanup happens in send_request)
        // Let's verify by creating a real timeout scenario
    }

    #[tokio::test]
    async fn test_response_with_unknown_id() {
        let handler = ProtocolHandler::new();

        // Send response with unknown ID (should be logged and discarded)
        let response = json!({
            "jsonrpc": "2.0",
            "id": 99999,
            "result": {}
        });

        // This should not panic
        handler.handle_message(response).await;

        // Verify no pending requests
        assert_eq!(handler.pending_requests.lock().await.len(), 0);
    }

    #[tokio::test]
    async fn test_concurrent_requests_correlation() {
        let handler = Arc::new(ProtocolHandler::new());
        let mut handles = vec![];

        // Send 10 concurrent requests
        for i in 0..10 {
            let h = Arc::clone(&handler);
            let handle = tokio::spawn(async move {
                let request = json!({
                    "jsonrpc": "2.0",
                    "method": "test",
                    "params": { "value": i }
                });

                let (message_with_id, response_rx) = h.prepare_request(request).await;
                let request_id = message_with_id.get("id").cloned().unwrap();

                // Simulate response
                h.handle_message(json!({
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": { "value": i }
                }))
                .await;

                // Receive response
                response_rx.await.unwrap()
            });
            handles.push(handle);
        }

        // Verify all responses correlated correctly
        for (i, handle) in handles.into_iter().enumerate() {
            let response = handle.await.unwrap();
            assert_eq!(response["result"]["value"], i);
        }
    }

    #[tokio::test]
    async fn test_json_rpc_error_response() {
        let handler = ProtocolHandler::new();

        let error_response = json!({
            "jsonrpc": "2.0",
            "id": 1,
            "error": {
                "code": -32601,
                "message": "Method not found",
                "data": { "method": "unknown" }
            }
        });

        let result = handler.handle_response_result(error_response);

        match result {
            Err(ProtocolError::JsonRpcError {
                code,
                message,
                data,
            }) => {
                assert_eq!(code, -32601);
                assert_eq!(message, "Method not found");
                assert!(data.is_some());
            }
            _ => panic!("Expected JsonRpcError"),
        }
    }

    #[tokio::test]
    async fn test_invalid_message_no_id_no_method() {
        let handler = ProtocolHandler::new();

        // Message with neither ID nor method (should be logged and discarded)
        let invalid = json!({
            "jsonrpc": "2.0",
            "params": {}
        });

        // Should not panic
        handler.handle_message(invalid).await;
    }

    // Helper builder tests

    #[test]
    fn test_build_initialize_request() {
        let request = build_initialize_request(1, json!({}));

        assert_eq!(request["jsonrpc"], "2.0");
        assert_eq!(request["method"], "initialize");
        assert_eq!(request["params"]["protocolVersion"], 1);
        assert!(request["params"]["clientCapabilities"].is_object());
        assert!(request.get("id").is_none()); // Should not have ID yet
    }

    #[test]
    fn test_build_session_new_request() {
        let request = build_session_new_request(".", None);

        assert_eq!(request["jsonrpc"], "2.0");
        assert_eq!(request["method"], "session/new");
        assert_eq!(request["params"]["cwd"], ".");
        assert!(request["params"]["mcpServers"].is_array());
    }

    #[test]
    fn test_build_session_new_request_with_servers() {
        let servers = json!([
            { "name": "server1", "url": "http://localhost:8080" }
        ]);

        let request = build_session_new_request("/path", Some(servers.clone()));

        assert_eq!(request["params"]["mcpServers"], servers);
    }

    #[test]
    fn test_build_session_prompt_request() {
        let prompt = json!([
            {
                "type": "text",
                "text": "Hello"
            }
        ]);

        let request = build_session_prompt_request("session-123", prompt.clone());

        assert_eq!(request["jsonrpc"], "2.0");
        assert_eq!(request["method"], "session/prompt");
        assert_eq!(request["params"]["sessionId"], "session-123");
        assert_eq!(request["params"]["prompt"], prompt);
    }

    #[test]
    fn test_build_session_cancel_request() {
        let request = build_session_cancel_request("session-456");

        assert_eq!(request["jsonrpc"], "2.0");
        assert_eq!(request["method"], "session/cancel");
        assert_eq!(request["params"]["sessionId"], "session-456");
    }

    // ========================================================================
    // MessageHandlerResult Tests (T044)
    // ========================================================================

    #[tokio::test]
    async fn test_message_handler_result_none_for_notification() {
        let handler = ProtocolHandler::new();

        // Notification (has method, no id)
        let notification = json!({
            "jsonrpc": "2.0",
            "method": "session/update",
            "params": {
                "sessionId": "session-123",
                "update": {
                    "sessionUpdate": "agent_message_chunk",
                    "content": {"type": "text", "text": "Hello"}
                }
            }
        });

        let result = handler.handle_message(notification).await;

        // Should return None for notifications
        assert!(matches!(result, MessageHandlerResult::None));
    }

    #[tokio::test]
    async fn test_message_handler_result_response_for_unsupported_method() {
        let handler = ProtocolHandler::new();

        // Request with unsupported method (has method and id)
        let request = json!({
            "jsonrpc": "2.0",
            "id": 42,
            "method": "unsupported/method",
            "params": {}
        });

        let result = handler.handle_message(request).await;

        // Should return Response variant with error
        match result {
            MessageHandlerResult::Response(response) => {
                assert_eq!(response.get("jsonrpc"), Some(&json!("2.0")));
                assert_eq!(response.get("id"), Some(&json!(42)));
                assert!(response.get("error").is_some());

                let error = response.get("error").unwrap();
                assert_eq!(error.get("code"), Some(&json!(-32601)));
                assert!(error.get("message").unwrap().as_str().unwrap().contains("not supported"));
            }
            _ => panic!("Expected Response variant, got {:?}", result),
        }
    }

    #[tokio::test]
    async fn test_message_handler_result_agent_request_for_permission() {
        let handler = ProtocolHandler::new();

        // Permission request (has method and id)
        let request = json!({
            "jsonrpc": "2.0",
            "id": 0,
            "method": "session/request_permission",
            "params": {
                "question": "Allow Write tool?",
                "options": [
                    {"id": "allow", "label": "Allow"},
                    {"id": "deny", "label": "Deny"}
                ]
            }
        });

        let result = handler.handle_message(request).await;

        // Should return AgentRequest variant
        match result {
            MessageHandlerResult::AgentRequest { request_id, method, params } => {
                assert_eq!(request_id, json!(0));
                assert_eq!(method, "session/request_permission");
                assert_eq!(params.get("question").unwrap(), "Allow Write tool?");
                assert!(params.get("options").is_some());
            }
            _ => panic!("Expected AgentRequest variant, got {:?}", result),
        }
    }

    #[tokio::test]
    async fn test_message_handler_result_none_for_response_to_our_request() {
        let handler = ProtocolHandler::new();

        // Prepare a request first
        let request = json!({
            "jsonrpc": "2.0",
            "method": "test",
            "params": {}
        });

        let (message_with_id, _response_rx) = handler.prepare_request(request).await;
        let request_id = message_with_id.get("id").cloned().unwrap();

        // Now simulate receiving a response
        let response = json!({
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {"status": "ok"}
        });

        let result = handler.handle_message(response).await;

        // Should return None (response was delivered to pending request)
        assert!(matches!(result, MessageHandlerResult::None));
    }

    #[tokio::test]
    async fn test_message_handler_result_invalid_message() {
        let handler = ProtocolHandler::new();

        // Invalid message (no method, no id, no result/error)
        let invalid = json!({
            "jsonrpc": "2.0",
            "params": {}
        });

        let result = handler.handle_message(invalid).await;

        // Should return None (message logged and discarded)
        assert!(matches!(result, MessageHandlerResult::None));
    }

    #[test]
    fn test_build_session_list_request_no_params() {
        let request = build_session_list_request(None, None);
        assert_eq!(request["jsonrpc"], "2.0");
        assert_eq!(request["method"], "session/list");
        assert!(request["params"].is_object());
    }

    #[test]
    fn test_build_session_list_request_with_cwd() {
        let request = build_session_list_request(Some("/home/user/project"), None);
        assert_eq!(request["params"]["cwd"], "/home/user/project");
    }

    #[test]
    fn test_build_session_list_request_with_cursor() {
        let request = build_session_list_request(None, Some("cursor-token-123"));
        assert_eq!(request["params"]["cursor"], "cursor-token-123");
    }

    #[test]
    fn test_build_session_list_request_with_both() {
        let request = build_session_list_request(Some("/tmp"), Some("page2"));
        assert_eq!(request["params"]["cwd"], "/tmp");
        assert_eq!(request["params"]["cursor"], "page2");
    }

    #[test]
    fn test_build_session_load_request_basic() {
        let request = build_session_load_request("session-123", "/home/user", None);
        assert_eq!(request["jsonrpc"], "2.0");
        assert_eq!(request["method"], "session/load");
        assert_eq!(request["params"]["sessionId"], "session-123");
        assert_eq!(request["params"]["cwd"], "/home/user");
        assert_eq!(request["params"]["mcpServers"], json!([]));
    }

    #[test]
    fn test_build_session_load_request_with_mcp_servers() {
        let servers = json!([{"name": "test-server", "url": "http://localhost:3000"}]);
        let request = build_session_load_request("session-123", "/tmp", Some(servers.clone()));
        assert_eq!(request["params"]["sessionId"], "session-123");
        assert_eq!(request["params"]["cwd"], "/tmp");
        assert_eq!(request["params"]["mcpServers"], servers);
    }

    #[test]
    fn test_build_session_resume_request_basic() {
        let request = build_session_resume_request("sess-123", None, None);
        assert_eq!(request["jsonrpc"], "2.0");
        assert_eq!(request["method"], "session/resume");
        assert_eq!(request["params"]["sessionId"], "sess-123");
        assert!(request["params"].get("cwd").is_none());
        assert!(request["params"].get("mcpServers").is_none());
    }

    #[test]
    fn test_build_session_resume_request_with_cwd() {
        let request = build_session_resume_request("sess-123", Some("/home/user"), None);
        assert_eq!(request["params"]["sessionId"], "sess-123");
        assert_eq!(request["params"]["cwd"], "/home/user");
    }

    #[test]
    fn test_build_session_resume_request_with_mcp_servers() {
        let servers = json!([{"name": "srv"}]);
        let request = build_session_resume_request("sess-123", Some("/tmp"), Some(servers.clone()));
        assert_eq!(request["params"]["sessionId"], "sess-123");
        assert_eq!(request["params"]["cwd"], "/tmp");
        assert_eq!(request["params"]["mcpServers"], servers);
    }

    #[test]
    fn test_build_session_close_request() {
        let request = build_session_close_request("session-abc");
        assert_eq!(request["jsonrpc"], "2.0");
        assert_eq!(request["method"], "session/close");
        assert_eq!(request["params"]["sessionId"], "session-abc");
    }

    #[test]
    fn test_build_set_config_option_request() {
        let request = build_set_config_option_request("session-abc", "mode", "code");
        assert_eq!(request["jsonrpc"], "2.0");
        assert_eq!(request["method"], "session/set_config_option");
        assert_eq!(request["params"]["sessionId"], "session-abc");
        assert_eq!(request["params"]["configId"], "mode");
        assert_eq!(request["params"]["value"], "code");
    }
}
