"""MCP Transport Layer

Implements stdio transport for JSON-RPC 2.0 communication.
"""

import sys
import json
from typing import Dict, Any, Optional, Callable


class StdioTransport:
    """Stdio transport for MCP protocol.

    Reads JSON-RPC messages from stdin, writes responses to stdout.
    Uses Content-Length header framing per MCP spec.
    """

    def __init__(self):
        self._running = False
        self._handler: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None

    def set_handler(self, handler: Callable[[Dict[str, Any]], Dict[str, Any]]):
        """Set the message handler function."""
        self._handler = handler

    def _read_message(self) -> Optional[Dict[str, Any]]:
        """Read a JSON-RPC message from stdin with Content-Length framing."""
        try:
            # Read headers
            headers = {}
            while True:
                line = sys.stdin.readline()
                if not line:
                    return None  # EOF
                line = line.strip()
                if not line:
                    break  # Empty line = end of headers
                if ':' in line:
                    key, value = line.split(':', 1)
                    headers[key.strip().lower()] = value.strip()

            # Get content length
            content_length = int(headers.get('content-length', 0))
            if content_length == 0:
                return None

            # Read body
            body = sys.stdin.read(content_length)
            return json.loads(body)

        except (json.JSONDecodeError, ValueError, KeyError):
            return None

    def _write_message(self, message: Dict[str, Any]):
        """Write a JSON-RPC message to stdout with Content-Length framing."""
        body = json.dumps(message)
        content_length = len(body.encode('utf-8'))

        sys.stdout.write(f"Content-Length: {content_length}\r\n")
        sys.stdout.write("\r\n")
        sys.stdout.write(body)
        sys.stdout.flush()

    def start(self):
        """Start the transport loop."""
        self._running = True

        while self._running:
            message = self._read_message()
            if message is None:
                break

            if self._handler:
                response = self._handler(message)
                if response:
                    self._write_message(response)

    def stop(self):
        """Stop the transport loop."""
        self._running = False


class DirectTransport:
    """Direct in-process transport for testing.

    Allows calling MCP tools directly without stdio.
    """

    def __init__(self, handler: Callable[[Dict[str, Any]], Dict[str, Any]]):
        self._handler = handler

    def call(self, method: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Call an MCP method directly."""
        request = {
            "jsonrpc": "2.0",
            "id": 1,
            "method": method,
            "params": params or {}
        }
        return self._handler(request)

    def call_tool(self, name: str, arguments: Optional[Dict[str, Any]] = None) -> Any:
        """Call a tool directly and return the result."""
        response = self.call("tools/call", {
            "name": name,
            "arguments": arguments or {}
        })

        if "error" in response:
            raise RuntimeError(response["error"].get("message", "Unknown error"))

        result = response.get("result", {})
        content = result.get("content", [])

        if content and len(content) > 0:
            return json.loads(content[0].get("text", "{}"))
        return result
