//! Transport abstraction for MCP communication
//!
//! Provides traits for implementing different transport mechanisms
//!
//! # Bidirectional Communication
//!
//! MCP supports bidirectional communication where:
//! - Client→Server: Tools/call, resources/read, prompts/get, etc.
//! - Server→Client: roots/list, sampling/createMessage, etc.
//!
//! The transport layer handles both directions using [`IncomingMessage`] to
//! distinguish between client requests and responses to server-initiated requests.

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;

use crate::protocol::types::{JsonRpcRequest, JsonRpcResponse};
use crate::server::multiplexer::JsonRpcClientRequest;

pub(crate) mod sealed {
    pub trait Sealed {}
}

/// Transport errors
#[derive(Debug, Error)]
pub enum TransportError {
    /// IO error
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    /// Serialization error
    #[error("Serialization error: {0}")]
    Serialization(#[from] serde_json::Error),

    /// Connection closed
    #[error("Connection closed")]
    Closed,

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

/// JSON-RPC notification (server->client, no id)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcNotification {
    /// JSON-RPC version
    pub jsonrpc: String,
    /// Method name
    pub method: String,
    /// Optional params
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<Value>,
}

impl JsonRpcNotification {
    /// Create a new notification
    pub fn new(method: impl Into<String>, params: Option<Value>) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            method: method.into(),
            params,
        }
    }
}

/// Incoming message from the transport
///
/// Can be either a client request (has method) or a response to a server-initiated
/// request (has result/error). Notifications are treated as requests without id.
#[derive(Debug, Clone)]
pub enum IncomingMessage {
    /// A request from the client (has method field)
    Request(JsonRpcRequest),

    /// A response to a server-initiated request (has result or error)
    Response(JsonRpcResponse),
}

impl IncomingMessage {
    /// Parse a raw JSON message into the appropriate type
    pub fn parse(json: &str) -> Result<Self, serde_json::Error> {
        // First, parse as generic Value to inspect structure
        let value: Value = serde_json::from_str(json)?;

        // If it has a "method" field, it's a request
        if value.get("method").is_some() {
            let request: JsonRpcRequest = serde_json::from_value(value)?;
            Ok(IncomingMessage::Request(request))
        } else {
            // Otherwise, it's a response (has result or error)
            let response: JsonRpcResponse = serde_json::from_value(value)?;
            Ok(IncomingMessage::Response(response))
        }
    }
}

/// Transport trait for reading and writing MCP messages
///
/// Supports bidirectional communication:
/// - Client→Server requests (standard MCP operations)
/// - Server→Client requests (roots/list, sampling, etc.)
#[async_trait]
pub trait Transport: sealed::Sealed + Send + Sync {
    /// Read a JSON-RPC request from the transport
    ///
    /// Note: For bidirectional transports, prefer `read_incoming()` which
    /// can distinguish between client requests and responses to server requests.
    async fn read_message(&mut self) -> Result<JsonRpcRequest, TransportError>;

    /// Read any incoming message (request or response)
    ///
    /// This enables bidirectional communication by distinguishing between:
    /// - Client requests (have method field)
    /// - Responses to server-initiated requests (have result/error)
    ///
    /// Default implementation wraps `read_message()` for backward compatibility.
    async fn read_incoming(&mut self) -> Result<IncomingMessage, TransportError> {
        // Default: assume all incoming messages are requests
        // Transports supporting bidirectional should override this
        let request = self.read_message().await?;
        Ok(IncomingMessage::Request(request))
    }

    /// Write a JSON-RPC response to the transport
    async fn write_message(&mut self, response: JsonRpcResponse) -> Result<(), TransportError>;

    /// Send a notification (server->client, no response expected)
    async fn send_notification(
        &mut self,
        notification: JsonRpcNotification,
    ) -> Result<(), TransportError>;

    /// Send a request to the client (server->client)
    ///
    /// Used for roots/list, sampling/createMessage, etc.
    /// The response will come back via `read_incoming()`.
    async fn send_request(&mut self, request: JsonRpcClientRequest) -> Result<(), TransportError>;

    /// Check if this transport supports bidirectional communication
    fn supports_bidirectional(&self) -> bool {
        false
    }

    /// Shutdown the transport gracefully
    async fn shutdown(&mut self) -> Result<(), TransportError> {
        Ok(())
    }
}

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

    // Mock transport for testing
    struct MockTransport {
        requests: Vec<JsonRpcRequest>,
        responses: Vec<JsonRpcResponse>,
        notifications: Vec<JsonRpcNotification>,
        closed: bool,
    }

    impl MockTransport {
        fn new() -> Self {
            Self {
                requests: Vec::new(),
                responses: Vec::new(),
                notifications: Vec::new(),
                closed: false,
            }
        }

        fn add_request(&mut self, request: JsonRpcRequest) {
            self.requests.push(request);
        }
    }

    impl sealed::Sealed for MockTransport {}

    #[async_trait]
    impl Transport for MockTransport {
        async fn read_message(&mut self) -> Result<JsonRpcRequest, TransportError> {
            if self.closed {
                return Err(TransportError::Closed);
            }

            self.requests.pop().ok_or(TransportError::Closed)
        }

        async fn write_message(&mut self, response: JsonRpcResponse) -> Result<(), TransportError> {
            if self.closed {
                return Err(TransportError::Closed);
            }

            self.responses.push(response);
            Ok(())
        }

        async fn send_notification(
            &mut self,
            notification: JsonRpcNotification,
        ) -> Result<(), TransportError> {
            if self.closed {
                return Err(TransportError::Closed);
            }

            self.notifications.push(notification);
            Ok(())
        }

        async fn send_request(
            &mut self,
            _request: JsonRpcClientRequest,
        ) -> Result<(), TransportError> {
            if self.closed {
                return Err(TransportError::Closed);
            }
            // Mock: just accept the request
            Ok(())
        }

        async fn shutdown(&mut self) -> Result<(), TransportError> {
            self.closed = true;
            Ok(())
        }
    }

    #[tokio::test]
    async fn test_mock_transport() {
        let mut transport = MockTransport::new();

        // Add a request
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(serde_json::json!(1)),
            method: "ping".to_string(),
            params: None,
        };
        transport.add_request(request.clone());

        // Read it back
        let read_request = transport.read_message().await.unwrap();
        assert_eq!(read_request.method, "ping");

        // Write a response
        let response = JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!(1),
            result: Some(serde_json::json!({})),
            error: None,
        };
        transport.write_message(response).await.unwrap();

        assert_eq!(transport.responses.len(), 1);
    }

    #[tokio::test]
    async fn test_transport_shutdown() {
        let mut transport = MockTransport::new();

        transport.shutdown().await.unwrap();

        let result = transport.read_message().await;
        assert!(matches!(result, Err(TransportError::Closed)));
    }
}
