"""JSON-RPC message parser and serializer for stdio communication.

This module provides functionality for parsing and serializing JSON-RPC 2.0 messages
with newline-delimited format for stdio communication, along with message correlation
and ID generation.
"""

import json
import uuid
from typing import Optional, Union, Dict, Any
from datetime import datetime

from ..models import JsonRpcMessage, JsonRpcError, MessageCorrelation, MessageStatus
from ..exceptions import InvalidJsonRpcError, MessageCorrelationError


class JsonRpcParser:
    """Parser for JSON-RPC 2.0 messages with newline-delimited format."""

    @staticmethod
    def parse(line: str) -> JsonRpcMessage:
        """Parse a newline-delimited JSON-RPC message.
        
        Args:
            line: A single line containing a JSON-RPC message
            
        Returns:
            Parsed JsonRpcMessage object
            
        Raises:
            InvalidJsonRpcError: If the message is invalid or malformed
        """
        if not line or not line.strip():
            raise InvalidJsonRpcError("Empty message line")
        
        try:
            # Remove any trailing newline characters
            line = line.rstrip('\n\r')
            
            # Parse JSON
            data = json.loads(line)
            
            # Validate and create JsonRpcMessage
            return JsonRpcMessage.model_validate(data)
            
        except json.JSONDecodeError as e:
            raise InvalidJsonRpcError(
                f"Invalid JSON in message: {e}",
                details={"line": line, "error": str(e)}
            )
        except Exception as e:
            raise InvalidJsonRpcError(
                f"Failed to parse JSON-RPC message: {e}",
                details={"line": line, "error": str(e)}
            )

    @staticmethod
    def parse_multiple(lines: str) -> list[JsonRpcMessage]:
        """Parse multiple newline-delimited JSON-RPC messages.
        
        Args:
            lines: Multiple lines containing JSON-RPC messages
            
        Returns:
            List of parsed JsonRpcMessage objects
            
        Raises:
            InvalidJsonRpcError: If any message is invalid
        """
        messages = []
        for line in lines.split('\n'):
            if line.strip():  # Skip empty lines
                messages.append(JsonRpcParser.parse(line))
        return messages


class JsonRpcSerializer:
    """Serializer for JSON-RPC 2.0 messages with newline-delimited format."""

    @staticmethod
    def serialize(message: JsonRpcMessage) -> str:
        """Serialize a JSON-RPC message to newline-delimited format.
        
        Args:
            message: JsonRpcMessage to serialize
            
        Returns:
            Newline-delimited JSON string
        """
        # Use Pydantic's JSON serialization with exclude_none
        json_str = message.model_dump_json(exclude_none=True)
        
        # Add newline delimiter
        return json_str + '\n'

    @staticmethod
    def serialize_multiple(messages: list[JsonRpcMessage]) -> str:
        """Serialize multiple JSON-RPC messages to newline-delimited format.
        
        Args:
            messages: List of JsonRpcMessage objects to serialize
            
        Returns:
            Newline-delimited JSON string with multiple messages
        """
        return ''.join(JsonRpcSerializer.serialize(msg) for msg in messages)


class MessageIdGenerator:
    """Generator for unique message IDs."""

    def __init__(self, prefix: str = "mcp"):
        """Initialize the ID generator.
        
        Args:
            prefix: Prefix for generated IDs (default: "mcp")
        """
        self.prefix = prefix
        self._counter = 0

    def generate(self) -> str:
        """Generate a unique message ID.
        
        Returns:
            Unique message ID string
        """
        self._counter += 1
        # Use UUID for uniqueness combined with counter for ordering
        return f"{self.prefix}-{self._counter}-{uuid.uuid4().hex[:8]}"

    def reset(self) -> None:
        """Reset the counter (useful for testing)."""
        self._counter = 0


class MessageCorrelationTracker:
    """Tracks correlation between stdio and HTTP messages."""

    def __init__(self):
        """Initialize the correlation tracker."""
        self._correlations: Dict[Union[str, int], MessageCorrelation] = {}

    def add_correlation(
        self,
        stdio_request_id: Union[str, int],
        http_request_id: str,
        method: str
    ) -> MessageCorrelation:
        """Add a new message correlation.
        
        Args:
            stdio_request_id: ID from the stdio JSON-RPC request
            http_request_id: ID for the HTTP request
            method: Method name being called
            
        Returns:
            Created MessageCorrelation object
            
        Raises:
            MessageCorrelationError: If correlation already exists
        """
        if stdio_request_id in self._correlations:
            raise MessageCorrelationError(
                f"Correlation already exists for stdio request ID: {stdio_request_id}",
                details={"stdio_request_id": stdio_request_id}
            )

        correlation = MessageCorrelation(
            stdio_request_id=stdio_request_id,
            http_request_id=http_request_id,
            timestamp=datetime.now(),
            method=method,
            status=MessageStatus.PENDING
        )
        
        self._correlations[stdio_request_id] = correlation
        return correlation

    def get_correlation(
        self,
        stdio_request_id: Union[str, int]
    ) -> Optional[MessageCorrelation]:
        """Get correlation by stdio request ID.
        
        Args:
            stdio_request_id: ID from the stdio JSON-RPC request
            
        Returns:
            MessageCorrelation object if found, None otherwise
        """
        return self._correlations.get(stdio_request_id)

    def get_correlation_by_http_id(
        self,
        http_request_id: str
    ) -> Optional[MessageCorrelation]:
        """Get correlation by HTTP request ID.
        
        Args:
            http_request_id: ID from the HTTP request
            
        Returns:
            MessageCorrelation object if found, None otherwise
        """
        for correlation in self._correlations.values():
            if correlation.http_request_id == http_request_id:
                return correlation
        return None

    def mark_completed(self, stdio_request_id: Union[str, int]) -> None:
        """Mark a correlation as completed.
        
        Args:
            stdio_request_id: ID from the stdio JSON-RPC request
            
        Raises:
            MessageCorrelationError: If correlation not found
        """
        correlation = self.get_correlation(stdio_request_id)
        if not correlation:
            raise MessageCorrelationError(
                f"No correlation found for stdio request ID: {stdio_request_id}",
                details={"stdio_request_id": stdio_request_id}
            )
        correlation.mark_completed()

    def mark_failed(self, stdio_request_id: Union[str, int]) -> None:
        """Mark a correlation as failed.
        
        Args:
            stdio_request_id: ID from the stdio JSON-RPC request
            
        Raises:
            MessageCorrelationError: If correlation not found
        """
        correlation = self.get_correlation(stdio_request_id)
        if not correlation:
            raise MessageCorrelationError(
                f"No correlation found for stdio request ID: {stdio_request_id}",
                details={"stdio_request_id": stdio_request_id}
            )
        correlation.mark_failed()

    def remove_correlation(self, stdio_request_id: Union[str, int]) -> None:
        """Remove a correlation from tracking.
        
        Args:
            stdio_request_id: ID from the stdio JSON-RPC request
        """
        self._correlations.pop(stdio_request_id, None)

    def get_pending_correlations(self) -> list[MessageCorrelation]:
        """Get all pending correlations.
        
        Returns:
            List of pending MessageCorrelation objects
        """
        return [
            corr for corr in self._correlations.values()
            if corr.is_pending()
        ]

    def clear(self) -> None:
        """Clear all correlations (useful for testing and cleanup)."""
        self._correlations.clear()


def create_request(
    method: str,
    params: Optional[Any] = None,
    request_id: Optional[Union[str, int]] = None
) -> JsonRpcMessage:
    """Create a JSON-RPC request message.
    
    Args:
        method: Method name to call
        params: Method parameters (optional)
        request_id: Request ID (optional, will be generated if not provided)
        
    Returns:
        JsonRpcMessage request object
    """
    if request_id is None:
        # Generate a unique ID
        request_id = MessageIdGenerator().generate()
    
    return JsonRpcMessage(
        jsonrpc="2.0",
        id=request_id,
        method=method,
        params=params
    )


def create_response(
    request_id: Union[str, int],
    result: Optional[Any] = None,
    error: Optional[JsonRpcError] = None
) -> JsonRpcMessage:
    """Create a JSON-RPC response message.
    
    Args:
        request_id: ID from the original request
        result: Result data (for success response)
        error: Error object (for error response)
        
    Returns:
        JsonRpcMessage response object
        
    Raises:
        ValueError: If both result and error are provided, or neither
    """
    if (result is None and error is None) or (result is not None and error is not None):
        raise ValueError("Must provide either result or error, but not both")
    
    return JsonRpcMessage(
        jsonrpc="2.0",
        id=request_id,
        result=result,
        error=error
    )


def create_notification(
    method: str,
    params: Optional[Any] = None
) -> JsonRpcMessage:
    """Create a JSON-RPC notification message (request without ID).
    
    Args:
        method: Method name to call
        params: Method parameters (optional)
        
    Returns:
        JsonRpcMessage notification object
    """
    return JsonRpcMessage(
        jsonrpc="2.0",
        method=method,
        params=params
    )


def create_error_response(
    request_id: Union[str, int],
    code: int,
    message: str,
    data: Optional[Any] = None
) -> JsonRpcMessage:
    """Create a JSON-RPC error response message.
    
    Args:
        request_id: ID from the original request
        code: Error code
        message: Error message
        data: Additional error data (optional)
        
    Returns:
        JsonRpcMessage error response object
    """
    error = JsonRpcError(
        code=code,
        message=message,
        data=data
    )
    
    return JsonRpcMessage(
        jsonrpc="2.0",
        id=request_id,
        error=error
    )
