"""HTTP and WebSocket gateway for a2a-mesh.

Provides the HTTP ingress layer for the mesh, handling JSON-RPC requests,
agent registration endpoints, health checks, rate limiting, and auth
middleware. Built on Starlette for async performance.
"""

from __future__ import annotations

import json
import time
from collections import defaultdict
from typing import TYPE_CHECKING, Any

from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Route

from a2a_mesh._logging import get_logger
from a2a_mesh.exceptions import MeshError
from a2a_mesh.models import AgentCard
from a2a_mesh.protocol.a2a import (
    INTERNAL_ERROR,
    INVALID_REQUEST,
    METHOD_NOT_FOUND,
    build_jsonrpc_error,
    build_jsonrpc_response,
)

if TYPE_CHECKING:
    from a2a_mesh.mesh import Mesh

logger = get_logger(__name__)


class RateLimiter:
    """Simple in-memory token bucket rate limiter.

    Attributes:
        max_requests: Maximum requests per window.
        window_seconds: Time window in seconds.
    """

    def __init__(
        self, max_requests: int = 100, window_seconds: float = 60.0
    ) -> None:
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        self._buckets: dict[str, list[float]] = defaultdict(list)

    def allow(self, key: str) -> bool:
        """Check if a request from the given key is allowed.

        Args:
            key: Rate limit key (e.g., client IP or agent name).

        Returns:
            True if the request is within limits.
        """
        now = time.monotonic()
        bucket = self._buckets[key]
        # Remove expired entries
        cutoff = now - self.window_seconds
        self._buckets[key] = [t for t in bucket if t > cutoff]

        if len(self._buckets[key]) >= self.max_requests:
            return False

        self._buckets[key].append(now)
        return True


def create_gateway(mesh: Mesh) -> Starlette:
    """Create the Starlette ASGI application for the mesh gateway.

    Args:
        mesh: The mesh instance to expose via HTTP.

    Returns:
        A configured Starlette application.
    """
    rate_limiter = RateLimiter()

    async def health(request: Request) -> JSONResponse:
        """Health check endpoint."""
        agent_count = len(mesh.registry.list_agents())
        return JSONResponse({
            "status": "healthy",
            "agents": agent_count,
        })

    async def list_agents(request: Request) -> JSONResponse:
        """List all registered agents."""
        agents = mesh.registry.list_agents()
        return JSONResponse({
            "agents": [
                {
                    "name": a.card.name,
                    "description": a.card.description,
                    "capabilities": a.card.capabilities,
                    "status": a.status.value,
                    "current_load": a.current_load,
                    "version": a.card.version,
                }
                for a in agents
            ]
        })

    async def register_agent(request: Request) -> JSONResponse:
        """Register a new agent via HTTP."""
        try:
            body = await request.json()
            card = AgentCard(**body)
            agent = mesh.register(card)
            return JSONResponse({
                "agent_id": agent.agent_id,
                "name": agent.card.name,
                "status": "registered",
            }, status_code=201)
        except MeshError as exc:
            return JSONResponse(
                {"error": str(exc)}, status_code=409
            )
        except Exception as exc:
            return JSONResponse(
                {"error": f"Invalid request: {exc}"}, status_code=400
            )

    async def jsonrpc_handler(request: Request) -> JSONResponse:
        """Handle JSON-RPC 2.0 requests."""
        client_ip = request.client.host if request.client else "unknown"

        if not rate_limiter.allow(client_ip):
            return JSONResponse(
                build_jsonrpc_error(None, -32000, "Rate limit exceeded"),
                status_code=429,
            )

        try:
            body = await request.json()
        except json.JSONDecodeError:
            return JSONResponse(
                build_jsonrpc_error(None, INVALID_REQUEST, "Invalid JSON"),
                status_code=400,
            )

        request_id = body.get("id")
        method = body.get("method", "")
        params = body.get("params", {})

        try:
            result = await _dispatch_method(mesh, method, params)
            return JSONResponse(
                build_jsonrpc_response(request_id, result)
            )
        except MeshError as exc:
            return JSONResponse(
                build_jsonrpc_error(request_id, INTERNAL_ERROR, str(exc))
            )
        except Exception as exc:
            logger.exception("jsonrpc.unhandled_error", method=method)
            return JSONResponse(
                build_jsonrpc_error(
                    request_id, INTERNAL_ERROR, f"Internal error: {exc}"
                )
            )

    async def traces_endpoint(request: Request) -> JSONResponse:
        """Return recent trace spans."""
        limit = int(request.query_params.get("limit", "50"))
        spans = mesh.tracer.get_traces(limit=limit)
        return JSONResponse({
            "spans": [s.model_dump(mode="json") for s in spans],
            "total_cost": mesh.tracer.total_cost(),
        })

    routes = [
        Route("/health", health, methods=["GET"]),
        Route("/agents", list_agents, methods=["GET"]),
        Route("/agents/register", register_agent, methods=["POST"]),
        Route("/rpc", jsonrpc_handler, methods=["POST"]),
        Route("/traces", traces_endpoint, methods=["GET"]),
    ]

    app = Starlette(routes=routes)
    return app


async def _dispatch_method(
    mesh: Any,
    method: str,
    params: dict[str, Any],
) -> Any:
    """Dispatch a JSON-RPC method to the appropriate mesh operation.

    Args:
        mesh: The mesh instance.
        method: JSON-RPC method name.
        params: Method parameters.

    Returns:
        The method result.

    Raises:
        MeshError: On dispatch failure.
    """
    if method == "tasks/send":
        task_input = params.get("input", "")
        capabilities = params.get("capabilities", [])
        result = await mesh.dispatch(
            task=task_input,
            required_capabilities=capabilities,
        )
        return {"status": "completed", "result": result}

    if method == "agents/list":
        agents = mesh.registry.list_agents()
        return {
            "agents": [
                {
                    "name": a.card.name,
                    "capabilities": a.card.capabilities,
                    "status": a.status.value,
                }
                for a in agents
            ]
        }

    if method == "agents/register":
        card = AgentCard(**params)
        agent = mesh.register(card)
        return {"agent_id": agent.agent_id, "name": agent.card.name}

    from a2a_mesh.exceptions import ProtocolError
    raise ProtocolError(f"Unknown method: {method}")
