# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
FastAPI application for the BrowserForge environment.

This module creates an HTTP server that exposes the BrowserEnvironment
over HTTP and WebSocket endpoints, compatible with EnvClient.

Endpoints:
    - POST /reset: Reset the environment
    - POST /step: Execute an action
    - GET /state: Get current environment state
    - GET /schema: Get action/observation schemas
    - WS /ws: WebSocket endpoint for persistent sessions

Usage:
    # Development (with auto-reload):
    uvicorn server.app:app --reload --host 0.0.0.0 --port 8000

    # Production:
    uvicorn server.app:app --host 0.0.0.0 --port 8000 --workers 4

    # Or run directly:
    python -m server.app
"""

import logging
import os
import time
from uuid import uuid4

from fastapi import Request
from fastapi.responses import JSONResponse, RedirectResponse

try:
    from openenv.core.env_server.http_server import create_app
    import openenv.core.env_server.http_server as openenv_http_server
    import openenv.core.env_server.serialization as openenv_serialization
except Exception as e:  # pragma: no cover
    raise ImportError(
        "openenv is required for the web interface. Install dependencies with '\n    uv sync\n'"
    ) from e

try:
    from ..demo import router as dashboard_router
    from ..models import BrowserAction, BrowserObservation
    from .browser_env_environment import BrowserEnvironment
except ImportError:
    from demo import router as dashboard_router
    from models import BrowserAction, BrowserObservation
    from server.browser_env_environment import BrowserEnvironment


def _serialize_observation_with_metadata(observation):
    """Preserve observation.metadata in OpenEnv responses.

    openenv-core currently drops metadata in serialize_observation(), which hides
    executed_action/browsergym_action from remote clients. We patch it here for
    this environment so diagnostics and training notebooks can inspect the real
    action path.
    """

    obs_dict = observation.model_dump(exclude={"reward", "done"})
    return {
        "observation": obs_dict,
        "reward": observation.reward,
        "done": observation.done,
    }


# Patch both module references used by create_app/http_server route handlers.
openenv_serialization.serialize_observation = _serialize_observation_with_metadata
openenv_http_server.serialize_observation = _serialize_observation_with_metadata


# Create the app with web interface and README integration
app = create_app(
    BrowserEnvironment,
    BrowserAction,
    BrowserObservation,
    env_name="browserforge",
    max_concurrent_envs=1,
)
app.include_router(dashboard_router)


@app.get("/", include_in_schema=False)
def root() -> RedirectResponse:
    return RedirectResponse(url="/command-center")


def _configure_logging() -> logging.Logger:
    level_name = os.getenv("BROWSER_ENV_LOG_LEVEL", "INFO").upper()
    level = getattr(logging, level_name, logging.INFO)
    root_logger = logging.getLogger()
    if not root_logger.handlers:
        logging.basicConfig(
            level=level,
            format="%(asctime)s %(levelname)s %(name)s %(message)s",
        )
    else:
        root_logger.setLevel(level)
    logger = logging.getLogger("browser_env.server")
    logger.setLevel(level)
    return logger


logger = _configure_logging()


@app.on_event("startup")
async def _log_startup() -> None:
    logger.info(
        "startup env=browserforge log_level=%s llm_provider=%s llm_enabled=%s miniwob_url_set=%s hf_space=%s",
        os.getenv("BROWSER_ENV_LOG_LEVEL", "INFO").upper(),
        os.getenv("BROWSER_ENV_LLM_PROVIDER", "openai_compatible"),
        os.getenv("BROWSER_ENV_DISABLE_LLM") != "1",
        bool(os.getenv("MINIWOB_URL")),
        bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID")),
    )


@app.middleware("http")
async def _request_logging_middleware(request: Request, call_next):
    request_id = request.headers.get("x-request-id") or str(uuid4())[:8]
    started = time.perf_counter()
    try:
        response = await call_next(request)
    except Exception:
        elapsed_ms = (time.perf_counter() - started) * 1000.0
        logger.exception(
            "request_failed request_id=%s method=%s path=%s elapsed_ms=%.1f",
            request_id,
            request.method,
            request.url.path,
            elapsed_ms,
        )
        return JSONResponse(
            status_code=500,
            content={"detail": "Internal Server Error", "request_id": request_id},
        )

    elapsed_ms = (time.perf_counter() - started) * 1000.0
    logger.info(
        "request request_id=%s method=%s path=%s status=%s elapsed_ms=%.1f",
        request_id,
        request.method,
        request.url.path,
        response.status_code,
        elapsed_ms,
    )
    response.headers["x-request-id"] = request_id
    return response


def main():
    """
    Entry point for direct execution via uv run or python -m.

    This function enables running the server without Docker:
        uv run --project . server
        uv run --project . server --port 8001
        python -m browser_env.server.app

    For production deployments, consider using uvicorn directly with
    multiple workers:
        uvicorn browser_env.server.app:app --workers 4
    """
    import argparse
    import uvicorn

    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8000)
    args = parser.parse_args()

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


if __name__ == "__main__":
    main()