"""Protocol translator implementation for MCP Proxy.

This module provides translation between stdio JSON-RPC messages and HTTP MCP requests,
handling method mapping, request correlation, and protocol-specific differences with
comprehensive error logging and diagnostics.
"""

import json
import logging
from typing import Dict, Optional, Union
from datetime import datetime

from ..models import (
    JsonRpcMessage,
    JsonRpcError,
    HttpMcpRequest,
    HttpMcpResponse,
    MessageCorrelation,
    MessageStatus,
)
from ..stdio.jsonrpc import MessageIdGenerator, MessageCorrelationTracker
from ..exceptions import (
    ProtocolError,
    UnsupportedMethodError,
    InvalidJsonRpcError,
    MessageCorrelationError,
)
from ..logging_config import get_logger

logger = get_logger(__name__)


class ProtocolTranslatorImpl:
    """Translates between stdio JSON-RPC and HTTP MCP protocols.
    
    This translator handles:
    - Converting JSON-RPC messages to HTTP POST requests
    - Mapping MCP methods between protocols
    - Maintaining request correlation across protocols
    - Managing session state and headers
    
    Attributes:
        base_url: Base URL for the HTTP MCP server
        session_id: Current MCP session ID
        correlation_tracker: Tracks message correlations
        id_generator: Generates unique HTTP request IDs
    """
    
    # MCP HTTP Streamable transport sends ALL requests to the same endpoint (base URL)
    # The method name is included in the JSON-RPC body, not the URL path
    # This is different from REST-style APIs where each method has its own endpoint
    # Use empty string "" instead of "/" to avoid trailing slash issues
    METHOD_ENDPOINT_MAP = {
        "initialize": "",
        "ping": "",
        "tools/list": "",
        "tools/call": "",
        "resources/list": "",
        "resources/read": "",
        "prompts/list": "",
        "prompts/get": "",
        "completion/complete": "",
        "logging/setLevel": "",
        "notifications/initialized": "",
        "notifications/cancelled": "",
        "notifications/progress": "",
    }
    
    # HTTP method for each MCP method (all are POST for JSON-RPC over HTTP)
    METHOD_HTTP_METHOD_MAP = {
        "initialize": "POST",
        "ping": "POST",
        "tools/list": "POST",
        "tools/call": "POST",
        "resources/list": "POST",
        "resources/read": "POST",
        "prompts/list": "POST",
        "prompts/get": "POST",
        "completion/complete": "POST",
        "logging/setLevel": "POST",
        "notifications/initialized": "POST",
        "notifications/cancelled": "POST",
        "notifications/progress": "POST",
    }
    
    def __init__(self, base_url: str, session_id: Optional[str] = None):
        """Initialize the protocol translator.
        
        Args:
            base_url: Base URL for the HTTP MCP server
            session_id: Optional initial session ID
        """
        self._base_url = base_url.rstrip("/")
        self._session_id = session_id
        self._correlation_tracker = MessageCorrelationTracker()
        self._id_generator = MessageIdGenerator(prefix="http")
        
        logger.debug(
            "ProtocolTranslatorImpl initialized with base_url=%s, session_id=%s",
            self._base_url, self._session_id
        )
    
    @property
    def session_id(self) -> Optional[str]:
        """Get the current session ID."""
        return self._session_id
    
    @session_id.setter
    def session_id(self, value: Optional[str]) -> None:
        """Set the session ID."""
        self._session_id = value
        logger.debug("Session ID updated to: %s", value)
    
    @property
    def correlation_tracker(self) -> MessageCorrelationTracker:
        """Get the correlation tracker."""
        return self._correlation_tracker
    
    def _extract_json_from_sse(self, sse_body: str) -> str:
        """Extract JSON data from SSE format response.
        
        SSE format looks like:
            event: message
            data: {"jsonrpc":"2.0",...}
        
        This method extracts the JSON from the "data:" line(s).
        
        Args:
            sse_body: Raw SSE response body
            
        Returns:
            Extracted JSON string, or empty string if no data found
        """
        if not sse_body:
            return ""
        
        data_lines = []
        for line in sse_body.split("\n"):
            line = line.strip()
            if line.startswith("data:"):
                # Extract content after "data:"
                data_content = line[5:].strip()
                if data_content:
                    data_lines.append(data_content)
        
        # Join multiple data lines (SSE spec allows multi-line data)
        return "\n".join(data_lines) if data_lines else ""
    
    def _get_endpoint_for_method(self, method: str) -> str:
        """Get the HTTP endpoint for an MCP method.
        
        For MCP HTTP Streamable transport, all methods go to the same endpoint (base URL).
        The method name is included in the JSON-RPC body.
        
        Args:
            method: MCP method name
            
        Returns:
            HTTP endpoint path (always "" for MCP HTTP Streamable to avoid trailing slash)
        """
        # MCP HTTP Streamable sends all requests to the base URL
        # Return "" for all methods - the method is in the JSON-RPC body
        return self.METHOD_ENDPOINT_MAP.get(method, "")
    
    def _get_http_method_for_mcp_method(self, method: str) -> str:
        """Get the HTTP method for an MCP method.
        
        For MCP HTTP Streamable transport, all requests use POST.
        
        Args:
            method: MCP method name
            
        Returns:
            HTTP method (always POST for MCP)
        """
        # MCP HTTP Streamable uses POST for all JSON-RPC requests
        return self.METHOD_HTTP_METHOD_MAP.get(method, "POST")
    
    def _build_http_body(self, message: JsonRpcMessage) -> Optional[str]:
        """Build HTTP request body from JSON-RPC message.
        
        For POST requests, the body contains the JSON-RPC message parameters.
        
        Args:
            message: JSON-RPC message
            
        Returns:
            JSON string for request body, or None for GET requests
        """
        if not message.is_request():
            raise InvalidJsonRpcError(
                "Cannot build HTTP body from non-request message",
                details={"message": message.model_dump()}
            )
        
        # Build request body with params
        body_data = {
            "jsonrpc": "2.0",
            "id": message.id,
            "method": message.method,
        }
        
        # Include params if present
        if message.params is not None:
            body_data["params"] = message.params
        
        return json.dumps(body_data)
    
    async def translate_stdio_to_http(
        self,
        message: JsonRpcMessage
    ) -> HttpMcpRequest:
        """Translate stdio JSON-RPC message to HTTP MCP request.
        
        This method:
        1. Validates the JSON-RPC message is a request
        2. Maps the MCP method to an HTTP endpoint
        3. Generates a unique HTTP request ID
        4. Creates correlation between stdio and HTTP IDs
        5. Builds the HTTP request with appropriate headers and body
        
        Args:
            message: JSON-RPC message from stdio
            
        Returns:
            HTTP MCP request ready to send
            
        Raises:
            InvalidJsonRpcError: If message is not a valid request
            UnsupportedMethodError: If method is not supported
            MessageCorrelationError: If correlation already exists
        """
        # Validate this is a request message
        if not message.is_request():
            raise InvalidJsonRpcError(
                "Can only translate request messages to HTTP",
                details={"message": message.model_dump()}
            )
        
        method = message.method
        if not method:
            raise InvalidJsonRpcError(
                "Request message must have a method",
                details={"message": message.model_dump()}
            )
        
        logger.debug(
            "Translating stdio request to HTTP: method=%s, id=%s",
            method, message.id
        )
        
        # Get endpoint and HTTP method for this MCP method
        endpoint = self._get_endpoint_for_method(method)
        http_method = self._get_http_method_for_mcp_method(method)
        
        # Generate unique HTTP request ID
        http_request_id = self._id_generator.generate()
        
        # Create correlation if this is not a notification
        if message.id is not None:
            self._correlation_tracker.add_correlation(
                stdio_request_id=message.id,
                http_request_id=http_request_id,
                method=method
            )
            logger.debug(
                "Created correlation: stdio_id=%s -> http_id=%s",
                message.id, http_request_id
            )
        
        # Build request body
        body = self._build_http_body(message)
        
        # Build headers - MCP HTTP Streamable requires Accept header for both JSON and SSE
        headers = {
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        }
        
        # Add session ID header if available
        if self._session_id:
            headers["Mcp-Session-Id"] = self._session_id
        
        # Create HTTP request
        http_request = HttpMcpRequest(
            method=http_method,
            url=f"{self._base_url}{endpoint}",
            headers=headers,
            body=body,
            session_id=self._session_id,
        )
        
        logger.info(
            "Translated stdio request to HTTP: method=%s, endpoint=%s, "
            "stdio_id=%s, http_id=%s",
            method, endpoint, message.id, http_request_id
        )
        
        return http_request
    
    async def translate_http_to_stdio(
        self,
        response: HttpMcpResponse,
        stdio_request_id: Optional[Union[str, int]] = None
    ) -> JsonRpcMessage:
        """Translate HTTP MCP response to stdio JSON-RPC message.
        
        This method:
        1. Parses the HTTP response body as JSON (or extracts from SSE format)
        2. Handles successful responses by extracting the result
        3. Handles error responses by translating to JSON-RPC errors
        4. Maintains correlation with the original stdio request
        
        Args:
            response: HTTP MCP response
            stdio_request_id: Optional stdio request ID for correlation lookup
            
        Returns:
            JSON-RPC message for stdio
            
        Raises:
            ProtocolError: If response cannot be translated
            MessageCorrelationError: If correlation cannot be found
        """
        logger.debug(
            "Translating HTTP response to stdio: status=%s, content_type=%s, stdio_id=%s",
            response.status, response.content_type, stdio_request_id
        )
        
        # If stdio_request_id not provided, we can't correlate
        # This is acceptable for notifications or streaming events
        request_id = stdio_request_id
        
        # Handle HTTP errors by translating to JSON-RPC errors
        if not response.is_success():
            return self._translate_http_error_to_stdio(response, request_id)
        
        # Extract body data - handle SSE format if content type is text/event-stream
        body_to_parse = response.body
        if response.content_type and "text/event-stream" in response.content_type:
            # SSE format: extract JSON from "data:" lines
            body_to_parse = self._extract_json_from_sse(response.body)
        
        # Parse response body
        try:
            body_data = json.loads(body_to_parse) if body_to_parse else {}
        except json.JSONDecodeError as e:
            logger.error("Failed to parse HTTP response body: %s", e)
            # Return JSON-RPC parse error
            return JsonRpcMessage(
                jsonrpc="2.0",
                id=request_id,
                error=JsonRpcError(
                    code=-32700,
                    message="Parse error",
                    data={"details": f"Invalid JSON in HTTP response: {str(e)}"}
                )
            )
        
        # Check if the response body is already a JSON-RPC message
        if "jsonrpc" in body_data:
            # Response is already in JSON-RPC format, use it directly
            # but ensure the ID matches our correlation
            if request_id is not None:
                body_data["id"] = request_id
            
            try:
                stdio_message = JsonRpcMessage.model_validate(body_data)
                logger.info(
                    "Translated HTTP response to stdio: status=%s, stdio_id=%s",
                    response.status, request_id
                )
                return stdio_message
            except Exception as e:
                logger.error("Failed to validate JSON-RPC message: %s", e)
                return JsonRpcMessage(
                    jsonrpc="2.0",
                    id=request_id,
                    error=JsonRpcError(
                        code=-32603,
                        message="Internal error",
                        data={"details": f"Invalid JSON-RPC format: {str(e)}"}
                    )
                )
        
        # Response is not in JSON-RPC format, wrap it as a result
        stdio_message = JsonRpcMessage(
            jsonrpc="2.0",
            id=request_id,
            result=body_data
        )
        
        logger.info(
            "Translated HTTP response to stdio: status=%s, stdio_id=%s",
            response.status, request_id
        )
        
        return stdio_message
    
    def _translate_http_error_to_stdio(
        self,
        response: HttpMcpResponse,
        request_id: Optional[Union[str, int]] = None
    ) -> JsonRpcMessage:
        """Translate HTTP error response to JSON-RPC error message.
        
        Maps HTTP status codes to appropriate JSON-RPC error codes:
        - 400 Bad Request -> -32602 Invalid params
        - 401 Unauthorized -> -32000 Server error (authentication)
        - 404 Not Found -> -32601 Method not found
        - 500+ Server errors -> -32603 Internal error
        - Other errors -> -32603 Internal error
        
        Args:
            response: HTTP error response
            request_id: Optional stdio request ID
            
        Returns:
            JSON-RPC error message
        """
        logger.debug(
            "Translating HTTP error to stdio: status=%s, stdio_id=%s",
            response.status, request_id
        )
        
        # Try to parse error details from response body
        error_data = None
        error_message = f"HTTP {response.status}"
        
        try:
            if response.body:
                body_data = json.loads(response.body)
                
                # Check if body contains a JSON-RPC error
                if "error" in body_data:
                    error_obj = body_data["error"]
                    if isinstance(error_obj, dict):
                        # Use the JSON-RPC error directly
                        return JsonRpcMessage(
                            jsonrpc="2.0",
                            id=request_id,
                            error=JsonRpcError(
                                code=error_obj.get("code", -32603),
                                message=error_obj.get("message", error_message),
                                data=error_obj.get("data")
                            )
                        )
                
                # Extract error message from common formats
                if "message" in body_data:
                    error_message = body_data["message"]
                elif "error" in body_data and isinstance(body_data["error"], str):
                    error_message = body_data["error"]
                
                # Store full body as error data
                error_data = body_data
        except json.JSONDecodeError:
            # If body is not JSON, use it as error message
            if response.body:
                error_message = response.body[:200]  # Limit length
                error_data = {"raw_response": response.body[:500]}
        
        # Map HTTP status to JSON-RPC error code
        if response.status == 400:
            error_code = -32602  # Invalid params
        elif response.status == 401:
            error_code = -32000  # Server error (authentication)
            error_message = f"Authentication failed: {error_message}"
        elif response.status == 404:
            error_code = -32601  # Method not found
            error_message = f"Method not found: {error_message}"
        elif response.status >= 500:
            error_code = -32603  # Internal error
            error_message = f"Server error: {error_message}"
        else:
            error_code = -32603  # Internal error
            error_message = f"HTTP error {response.status}: {error_message}"
        
        # Create JSON-RPC error message
        stdio_message = JsonRpcMessage(
            jsonrpc="2.0",
            id=request_id,
            error=JsonRpcError(
                code=error_code,
                message=error_message,
                data=error_data
            )
        )
        
        logger.info(
            "Translated HTTP error to stdio: status=%s, code=%s, stdio_id=%s",
            response.status, error_code, request_id
        )
        
        return stdio_message
    
    async def translate_sse_event_to_stdio(
        self,
        event_data: str,
        event_type: Optional[str] = None
    ) -> Optional[JsonRpcMessage]:
        """Translate Server-Sent Event to stdio JSON-RPC message.
        
        SSE events from the HTTP MCP server are translated to JSON-RPC
        notifications or responses depending on the event type.
        
        Args:
            event_data: SSE event data (JSON string)
            event_type: Optional SSE event type
            
        Returns:
            JSON-RPC message for stdio, or None if event should be ignored
            
        Raises:
            ProtocolError: If event cannot be translated
        """
        logger.debug(
            "Translating SSE event to stdio: type=%s, data_length=%s",
            event_type, len(event_data) if event_data else 0
        )
        
        # Parse event data
        try:
            event_obj = json.loads(event_data) if event_data else {}
        except json.JSONDecodeError as e:
            logger.error("Failed to parse SSE event data: %s", e)
            raise ProtocolError(
                "Invalid JSON in SSE event",
                details={"event_type": event_type, "error": str(e)}
            )
        
        # Check if event is already a JSON-RPC message
        if "jsonrpc" in event_obj:
            try:
                stdio_message = JsonRpcMessage.model_validate(event_obj)
                logger.info(
                    "Translated SSE event to stdio: type=%s, method=%s",
                    event_type, stdio_message.method
                )
                return stdio_message
            except Exception as e:
                logger.error("Failed to validate JSON-RPC message from SSE: %s", e)
                raise ProtocolError(
                    "Invalid JSON-RPC format in SSE event",
                    details={"event_type": event_type, "error": str(e)}
                )
        
        # Handle different event types
        if event_type == "message":
            # Regular message event - wrap as notification
            stdio_message = JsonRpcMessage(
                jsonrpc="2.0",
                method="notifications/message",
                params=event_obj
            )
            logger.info("Translated SSE message event to stdio notification")
            return stdio_message
        
        elif event_type == "error":
            # Error event - create error notification
            error_message = event_obj.get("message", "Unknown error")
            stdio_message = JsonRpcMessage(
                jsonrpc="2.0",
                method="notifications/error",
                params={
                    "error": error_message,
                    "details": event_obj
                }
            )
            logger.info("Translated SSE error event to stdio notification")
            return stdio_message
        
        elif event_type in ["ping", "heartbeat"]:
            # Heartbeat events can be ignored or logged
            logger.debug("Ignoring SSE heartbeat event")
            return None
        
        else:
            # Unknown event type - wrap as generic notification
            stdio_message = JsonRpcMessage(
                jsonrpc="2.0",
                method="notifications/event",
                params={
                    "type": event_type,
                    "data": event_obj
                }
            )
            logger.info(
                "Translated SSE event to stdio notification: type=%s",
                event_type
            )
            return stdio_message
    
    def correlate_messages(
        self,
        request_id: Union[str, int],
        response_id: str
    ) -> None:
        """Correlate request and response messages.
        
        This method is used to manually correlate messages when needed.
        In most cases, correlation is handled automatically during translation.
        
        Args:
            request_id: Stdio request ID
            response_id: HTTP response ID
            
        Raises:
            MessageCorrelationError: If correlation fails
        """
        # Check if correlation already exists
        correlation = self._correlation_tracker.get_correlation(request_id)
        if correlation:
            logger.warning(
                "Correlation already exists for request_id=%s, updating",
                request_id
            )
            # Update the HTTP request ID
            correlation.http_request_id = response_id
        else:
            # Create new correlation
            self._correlation_tracker.add_correlation(
                stdio_request_id=request_id,
                http_request_id=response_id,
                method="unknown"
            )
        
        logger.debug(
            "Manually correlated messages: stdio_id=%s -> http_id=%s",
            request_id, response_id
        )
    
    def get_correlation(
        self,
        stdio_request_id: Union[str, int]
    ) -> Optional[MessageCorrelation]:
        """Get correlation for a stdio request ID.
        
        Args:
            stdio_request_id: Stdio request ID
            
        Returns:
            MessageCorrelation if found, None otherwise
        """
        return self._correlation_tracker.get_correlation(stdio_request_id)
    
    def mark_correlation_completed(
        self,
        stdio_request_id: Union[str, int]
    ) -> None:
        """Mark a correlation as completed.
        
        Args:
            stdio_request_id: Stdio request ID
            
        Raises:
            MessageCorrelationError: If correlation not found
        """
        self._correlation_tracker.mark_completed(stdio_request_id)
        logger.debug("Marked correlation completed: stdio_id=%s", stdio_request_id)
    
    def mark_correlation_failed(
        self,
        stdio_request_id: Union[str, int]
    ) -> None:
        """Mark a correlation as failed.
        
        Args:
            stdio_request_id: Stdio request ID
            
        Raises:
            MessageCorrelationError: If correlation not found
        """
        self._correlation_tracker.mark_failed(stdio_request_id)
        logger.debug("Marked correlation failed: stdio_id=%s", stdio_request_id)
    
    def clear_correlations(self) -> None:
        """Clear all correlations.
        
        Useful for cleanup and testing.
        """
        self._correlation_tracker.clear()
        logger.debug("Cleared all correlations")
    
    def get_pending_correlations(self) -> list[MessageCorrelation]:
        """Get all pending correlations.
        
        Returns:
            List of pending MessageCorrelation objects
        """
        return self._correlation_tracker.get_pending_correlations()
