"""FastAPI and WebSocket entrypoint for SpreadsheetBench OpenEnv."""

from __future__ import annotations

from collections.abc import Callable
import os

from fastapi import FastAPI
from openenv.core.env_server.http_server import create_app
from openenv.core.env_server.interfaces import Environment
from starlette.types import ASGIApp, Receive, Scope, Send
from starlette.websockets import WebSocketDisconnect

from spreadsheetbench_env.models import (
    SpreadsheetBenchAction,
    SpreadsheetBenchObservation,
    native_tool_specs,
)
from spreadsheetbench_env.server.environment import SpreadsheetBenchEnvironment


_SEND_AFTER_CLOSE = 'Cannot call "send" once a close message has been sent.'


class DisconnectedWebSocketMiddleware:
    """Suppress only OpenEnv errors caused by an already-disconnected client."""

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        try:
            await self.app(scope, receive, send)
        except WebSocketDisconnect:
            if scope["type"] != "websocket":
                raise
        except RuntimeError as error:
            if scope["type"] != "websocket" or str(error) != _SEND_AFTER_CLOSE:
                raise


def _configured_capacity() -> int:
    raw = os.getenv("MAX_CONCURRENT_ENVS", "1")
    try:
        capacity = int(raw)
    except ValueError as error:
        raise ValueError(f"MAX_CONCURRENT_ENVS must be an integer, got {raw!r}") from error
    if capacity < 1:
        raise ValueError("MAX_CONCURRENT_ENVS must be at least 1")
    return capacity


def _configured_default_task_id() -> int | None:
    value = os.getenv("DEFAULT_TASK_ID")
    if value is None or not value.strip():
        return None
    try:
        return int(value)
    except ValueError as error:
        raise ValueError(f"DEFAULT_TASK_ID must be an integer, got {value!r}") from error


def _configured_environment_class(
    default_task_id: int | None,
) -> type[SpreadsheetBenchEnvironment]:
    if default_task_id is None:
        return SpreadsheetBenchEnvironment

    class ConfiguredSpreadsheetBenchEnvironment(SpreadsheetBenchEnvironment):
        def __init__(self) -> None:
            super().__init__(default_task_id=default_task_id)

    ConfiguredSpreadsheetBenchEnvironment.__name__ = "SpreadsheetBenchEnvironment"
    return ConfiguredSpreadsheetBenchEnvironment


def create_spreadsheetbench_app(
    max_concurrent_envs: int | None = None,
    environment_factory: Callable[[], Environment] | None = None,
) -> FastAPI:
    """Create the pinned OpenEnv app without constructing runtime data eagerly."""

    capacity = _configured_capacity() if max_concurrent_envs is None else max_concurrent_envs
    if capacity < 1:
        raise ValueError("MAX_CONCURRENT_ENVS must be at least 1")
    factory = environment_factory or _configured_environment_class(
        _configured_default_task_id()
    )
    app = create_app(
        factory,
        SpreadsheetBenchAction,
        SpreadsheetBenchObservation,
        env_name="spreadsheetbench_env",
        max_concurrent_envs=capacity,
    )
    app.add_api_route("/tools", native_tool_specs, methods=["GET"])
    app.add_middleware(DisconnectedWebSocketMiddleware)
    return app


app = create_spreadsheetbench_app()


def main(host: str = "0.0.0.0", port: int = 8000) -> None:
    """Run the standalone OpenEnv server."""

    import uvicorn

    uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    main()