"""
Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com

Route registration helper for AppFactory.
"""

from typing import Any, Dict, List, Union

from fastapi import Body, FastAPI, Request, WebSocket
from fastapi.responses import JSONResponse

from mcp_proxy_adapter.api.handlers import (
    handle_json_rpc,
    handle_batch_json_rpc,
    get_server_health,
    get_commands_list,
    handle_heartbeat,
    execute_command_async,
)
from mcp_proxy_adapter.core.errors import (
    InvalidRequestError,
    MethodNotFoundError,
    ValidationError,
)

try:
    from mcp_proxy_adapter.api.schemas import (
        JsonRpcRequest,
        JsonRpcSuccessResponse,
        JsonRpcErrorResponse,
        HealthResponse,
        CommandListResponse,
        AsyncExecuteRequest,
    )
except Exception:  # pragma: no cover - fallback for simplified environments
    JsonRpcRequest = Dict[str, Any]  # type: ignore
    JsonRpcSuccessResponse = Dict[str, Any]  # type: ignore
    JsonRpcErrorResponse = Dict[str, Any]  # type: ignore
    HealthResponse = Dict[str, Any]  # type: ignore
    CommandListResponse = Dict[str, Any]  # type: ignore
    AsyncExecuteRequest = Dict[str, Any]  # type: ignore

try:
    from mcp_proxy_adapter.api.tools import get_tool_description, execute_tool
except Exception:  # pragma: no cover - optional dependency
    get_tool_description = None
    execute_tool = None

try:
    from mcp_proxy_adapter.api.ws_endpoint import handle_websocket_job_push
except ImportError:  # pragma: no cover
    handle_websocket_job_push = None


def setup_routes(app: FastAPI) -> None:
    """
    Register built-in API routes.

    /ws is always registered when the handler is available. WS is the way
    the server pushes job results to the client (no polling). No config.
    """

    @app.get("/health", response_model=HealthResponse)
    async def health():  # type: ignore
        return await get_server_health()  # type: ignore[misc]

    @app.get("/commands", response_model=CommandListResponse)
    async def commands():  # type: ignore
        return await get_commands_list()  # type: ignore[misc]

    @app.get("/heartbeat")
    async def heartbeat():  # type: ignore
        """Built-in heartbeat endpoint for proxy health checks."""
        return await handle_heartbeat()  # type: ignore[misc]

    @app.post(
        "/api/jsonrpc",
        response_model=Union[JsonRpcSuccessResponse, JsonRpcErrorResponse],
    )
    async def jsonrpc(request: JsonRpcRequest):  # type: ignore
        return await handle_json_rpc(request.dict())  # type: ignore[misc]

    @app.post(
        "/api/jsonrpc/batch",
        response_model=List[Union[JsonRpcSuccessResponse, JsonRpcErrorResponse]],
    )
    async def jsonrpc_batch(requests: List[JsonRpcRequest]):  # type: ignore
        payload = [req.dict() for req in requests]
        return await handle_batch_json_rpc(payload)  # type: ignore[misc]

    @app.post("/api/async")
    async def api_async(request: Request, body: AsyncExecuteRequest) -> Any:
        """
        Accept async command execution; return immediately with deliver_id.
        Result is delivered via WebSocket to subscribers of that deliver_id.
        """
        cmd = body.command or ""
        params = body.params or {}
        deliver_id = body.deliver_id
        if not cmd.strip():
            return JSONResponse(
                status_code=400,
                content={"error": "command is required"},
            )
        request_id = getattr(request.state, "request_id", None)
        try:
            result = await execute_command_async(
                cmd, params, deliver_id, request_id, request
            )
            return result
        except MethodNotFoundError as e:
            return JSONResponse(
                status_code=404,
                content={"error": e.message, "code": e.code, "data": e.data},
            )
        except (ValidationError, InvalidRequestError) as e:
            return JSONResponse(
                status_code=400,
                content={"error": e.message, "code": e.code, "data": e.data},
            )

    @app.post("/tools/{tool_name}")
    async def execute_tool_endpoint(
        tool_name: str,
        request: Request,
        payload: Dict[str, Any] = Body(default_factory=dict),
    ):
        if execute_tool is None:
            raise ValueError("Tool execution is not available in this environment")
        return await execute_tool(tool_name, request, payload)

    @app.get("/tools/{tool_name}")
    async def get_tool_endpoint(tool_name: str):
        if get_tool_description is None:
            raise ValueError("Tool description is not available in this environment")
        return await get_tool_description(tool_name)

    if handle_websocket_job_push is not None:

        @app.websocket("/ws")
        async def websocket_job_push(websocket: WebSocket) -> None:
            await handle_websocket_job_push(websocket, app)
