"""
Main proxy implementation for MCP over HTTP Streamable.

This module implements the core proxy logic that bridges stdio transport
with HTTP Streamable transport to a remote MCP server.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import sys
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass
from enum import IntEnum
from typing import Any

import httpx

# Configure logging
LOG_LEVEL = os.getenv("MCP_LOG_LEVEL", "INFO").upper()
logging.basicConfig(
    level=getattr(logging, LOG_LEVEL),
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    stream=sys.stderr,  # Log to stderr to keep stdout clean for MCP messages
)
logger = logging.getLogger(__name__)


class JSONRPCErrorCode(IntEnum):
    """JSON-RPC 2.0 error codes."""

    PARSE_ERROR = -32700
    INVALID_REQUEST = -32600
    METHOD_NOT_FOUND = -32601
    INVALID_PARAMS = -32602
    INTERNAL_ERROR = -32603
    SERVER_ERROR = -32000
    CONNECTION_ERROR = -32001


@dataclass(frozen=True)
class SSEEvent:
    """Represents a parsed SSE event."""

    event_type: str
    data: str


class MCPProxyError(Exception):
    """Base exception for MCP proxy errors."""


class ConfigurationError(MCPProxyError):
    """Raised when configuration is invalid."""


class ConnectionError(MCPProxyError):
    """Raised when connection to remote MCP server fails."""


class MCPProxy:
    """
    MCP Proxy that bridges stdio transport with HTTP Streamable transport.

    The proxy reads JSON-RPC messages from stdin, forwards them to a remote
    MCP server via HTTP POST, and writes SSE responses to stdout.
    """

    def __init__(
        self,
        base_url: str | None = None,
        auth_token: str | None = None,
        timeout: float = 60.0,
    ) -> None:
        """
        Initialize the MCP proxy.

        Args:
            base_url: Base URL of the MCP server (default: from MCP_BASE_URL env var)
            auth_token: Authentication token (default: from X_AUTH_TOKEN env var)
            timeout: Request timeout in seconds (default: 60.0)

        Raises:
            ConfigurationError: If required configuration is missing
        """
        self.base_url = base_url or os.getenv(
            "MCP_BASE_URL", "https://mcp.ubidots.com/mcp"
        )
        self.auth_token = auth_token or os.getenv("X_AUTH_TOKEN")

        if not self.auth_token:
            raise ConfigurationError(
                "X_AUTH_TOKEN environment variable is required for authentication"
            )

        self.timeout = timeout
        self.session_id: str | None = None
        self._client: httpx.AsyncClient | None = None
        self._running = False

        logger.info(f"Initialized MCP proxy with base URL: {self.base_url}")

    @property
    def headers(self) -> dict[str, str]:
        """Get HTTP headers for requests."""
        headers = {
            "Authorization": f"Bearer {self.auth_token}",
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        }
        if self.session_id:
            headers["mcp-session-id"] = self.session_id
        return headers

    async def __aenter__(self) -> MCPProxy:
        """Async context manager entry."""
        self._client = httpx.AsyncClient(
            timeout=httpx.Timeout(self.timeout),
            follow_redirects=True,
            http2=True,
        )
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Async context manager exit."""
        if self._client:
            await self._client.aclose()

    def _create_error_response(
        self,
        message_id: Any,
        code: JSONRPCErrorCode,
        message: str,
        data: str | None = None,
    ) -> dict[str, Any]:
        """
        Create a JSON-RPC error response.

        Args:
            message_id: The message ID to respond to
            code: JSON-RPC error code
            message: Error message
            data: Optional error data

        Returns:
            JSON-RPC error response dictionary
        """
        error_response: dict[str, Any] = {
            "jsonrpc": "2.0",
            "id": message_id,
            "error": {
                "code": code,
                "message": message,
            },
        }
        if data is not None:
            error_response["error"]["data"] = data
        return error_response

    async def _write_stdout_message(self, message: dict[str, Any]) -> None:
        """
        Write a JSON-RPC message to stdout.

        Args:
            message: JSON-RPC message to write
        """
        try:
            output = json.dumps(message) + "\n"
            sys.stdout.write(output)
            sys.stdout.flush()
            logger.debug(f"Sent message to stdout: {message}")
        except Exception as e:
            logger.error(f"Error writing to stdout: {e}")

    async def _handle_parse_error(self, error: json.JSONDecodeError) -> None:
        """Handle JSON parsing errors from stdin."""
        logger.error(f"Failed to parse JSON from stdin: {error}")
        error_response = self._create_error_response(
            message_id=None,
            code=JSONRPCErrorCode.PARSE_ERROR,
            message="Parse error",
            data=str(error),
        )
        await self._write_stdout_message(error_response)

    async def _read_stdin_messages(self) -> AsyncIterator[dict[str, Any]]:
        """
        Read JSON-RPC messages from stdin.

        Yields:
            Parsed JSON-RPC message dictionaries

        The MCP protocol uses newline-delimited JSON messages.
        """
        logger.debug("Starting to read messages from stdin")
        loop = asyncio.get_event_loop()

        while self._running:
            try:
                line = await loop.run_in_executor(None, sys.stdin.readline)

                if not line:
                    logger.info("EOF reached on stdin, shutting down")
                    break

                line = line.strip()
                if not line:
                    continue

                try:
                    message = json.loads(line)
                    logger.debug(f"Received message from stdin: {message}")
                    yield message
                except json.JSONDecodeError as e:
                    await self._handle_parse_error(e)

            except Exception as e:
                logger.error(f"Error reading from stdin: {e}")
                break

    def _parse_sse_line(self, line: str) -> tuple[str | None, str | None]:
        """
        Parse a single SSE line.

        Args:
            line: A line from SSE stream

        Returns:
            Tuple of (field, value) or (None, None) if not a valid field
        """
        line = line.rstrip("\r\n")

        # Skip empty lines and comments
        if not line or line.startswith(":"):
            return None, None

        # Parse field:value format
        if ":" not in line:
            return line, ""

        field, _, value = line.partition(":")
        # Strip leading space from value (SSE spec)
        return field, value[1:] if value.startswith(" ") else value

    @asynccontextmanager
    async def _stream_remote_response(
        self, message: dict[str, Any]
    ) -> AsyncIterator[httpx.Response]:
        """
        Stream response from remote MCP server.

        Args:
            message: JSON-RPC message to send

        Yields:
            HTTP response object

        Raises:
            ConnectionError: If client not initialized or base URL not configured
        """
        if not self._client:
            raise ConnectionError("HTTP client not initialized")

        if not self.base_url:
            raise ConnectionError("MCP base URL not configured")

        logger.debug(f"Sending message to remote MCP: {message}")

        async with self._client.stream(
            "POST",
            self.base_url,
            headers=self.headers,
            json=message,
        ) as response:
            yield response

    def _update_session_id(self, response: httpx.Response) -> None:
        """Update session ID from response headers if present."""
        new_session_id = response.headers.get("mcp-session-id")
        if not new_session_id:
            return

        if self.session_id != new_session_id:
            logger.info(f"Session ID: {new_session_id}")
        self.session_id = new_session_id

    async def _handle_http_error(
        self,
        response: httpx.Response,
        message: dict[str, Any],
        is_notification: bool,
    ) -> None:
        """Handle HTTP error responses from remote server."""
        body = await response.aread()
        logger.error(f"HTTP error {response.status_code}: {body.decode()}")

        if is_notification:
            return

        error_response = self._create_error_response(
            message_id=message.get("id"),
            code=JSONRPCErrorCode.SERVER_ERROR,
            message=f"Remote server error: {response.status_code}",
            data=body.decode(),
        )
        await self._write_stdout_message(error_response)

    async def _parse_sse_data(self, data_lines: list[str]) -> None:
        """Parse and write SSE data to stdout."""
        if not data_lines:
            return

        data_str = "\n".join(data_lines)
        if not data_str:
            return

        try:
            response_data = json.loads(data_str)
            await self._write_stdout_message(response_data)
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse SSE data: {e}")

    async def _process_sse_stream(self, response: httpx.Response) -> None:
        """Process Server-Sent Events stream from remote server."""
        current_data: list[str] = []

        async for line in response.aiter_lines():
            field, value = self._parse_sse_line(line)

            # Empty line = end of event
            if field is None and not line.strip():
                await self._parse_sse_data(current_data)
                current_data = []
                continue

            if field == "data":
                current_data.append(value or "")

        # Handle any remaining data after stream ends
        await self._parse_sse_data(current_data)

    async def _process_json_response(self, response: httpx.Response) -> None:
        """Process JSON response from remote server."""
        body = await response.aread()
        if not body:
            return

        try:
            response_data = json.loads(body.decode())
            await self._write_stdout_message(response_data)
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse JSON response: {e}")

    async def _process_response_by_content_type(
        self, response: httpx.Response
    ) -> None:
        """Route response processing based on content type."""
        content_type = response.headers.get("content-type", "")

        if "text/event-stream" in content_type:
            await self._process_sse_stream(response)
        elif "application/json" in content_type:
            await self._process_json_response(response)
        else:
            logger.warning(f"Unexpected content-type: {content_type}")

    async def _send_error_if_not_notification(
        self,
        message: dict[str, Any],
        is_notification: bool,
        code: JSONRPCErrorCode,
        error_message: str,
        error_data: str,
    ) -> None:
        """Send error response if message is not a notification."""
        if is_notification:
            return

        error_response = self._create_error_response(
            message_id=message.get("id"),
            code=code,
            message=error_message,
            data=error_data,
        )
        await self._write_stdout_message(error_response)

    async def _send_to_remote(self, message: dict[str, Any]) -> None:
        """
        Send a JSON-RPC message to the remote MCP server.

        Args:
            message: JSON-RPC message to send

        Raises:
            ConnectionError: If connection to remote server fails
        """
        is_notification = "id" not in message

        try:
            async with self._stream_remote_response(message) as response:
                self._update_session_id(response)

                # Check for HTTP errors
                if response.status_code >= 400:
                    await self._handle_http_error(response, message, is_notification)
                    return

                # Process response based on content type
                await self._process_response_by_content_type(response)

        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP error from remote MCP: {e}")
            await self._send_error_if_not_notification(
                message,
                is_notification,
                JSONRPCErrorCode.SERVER_ERROR,
                f"Remote server error: {e.response.status_code}",
                str(e),
            )

        except httpx.RequestError as e:
            logger.error(f"Request error to remote MCP: {e}")
            await self._send_error_if_not_notification(
                message,
                is_notification,
                JSONRPCErrorCode.CONNECTION_ERROR,
                "Failed to connect to remote MCP server",
                str(e),
            )

        except Exception as e:
            logger.error(f"Unexpected error communicating with remote MCP: {e}")
            await self._send_error_if_not_notification(
                message,
                is_notification,
                JSONRPCErrorCode.INTERNAL_ERROR,
                "Internal error",
                str(e),
            )

    def _log_special_methods(self, method: str | None) -> None:
        """Log special MCP protocol methods."""
        if method == "initialize":
            logger.info("Received initialize request")
        elif method == "initialized":
            logger.info("Received initialized notification")

    async def run(self) -> None:
        """
        Run the proxy main loop.

        This method reads messages from stdin, forwards them to the remote
        MCP server, and writes responses to stdout.
        """
        self._running = True
        logger.info("Starting MCP proxy main loop")

        try:
            async for message in self._read_stdin_messages():
                self._log_special_methods(message.get("method"))
                await self._send_to_remote(message)

        except KeyboardInterrupt:
            logger.info("Received keyboard interrupt, shutting down")
        except Exception as e:
            logger.error(f"Fatal error in proxy main loop: {e}", exc_info=True)
        finally:
            self._running = False
            logger.info("MCP proxy stopped")


async def run_proxy() -> None:
    """
    Main entry point to run the MCP proxy.

    This function creates and runs the proxy with configuration from
    environment variables.
    """
    try:
        async with MCPProxy() as proxy:
            await proxy.run()
    except ConfigurationError as e:
        logger.error(f"Configuration error: {e}")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Fatal error: {e}", exc_info=True)
        sys.exit(1)
