"""The server: OpenEnv's app plus read-only JSON about the tasks on the same port.

OpenEnv provides `/reset`, `/step`, `/state`, `/schema`, `/ws`, `/mcp`, `/health`, the Task API and its web UI at
`/web`. We add:

    /api/tasks                       task summaries with the references' costs
    /api/tasks/{id}                  a public task (what the agent sees, structured)
    /api/tasks/{id}/reference        the optimum and the baselines - for people; no tool shows them
    /api/episodes                    live episodes on this server (newest first)
    /api/tasks/{id}/plans            the plans the replay can show: optimum, baselines, the Corps, live policies, and
                                     every model's plan in the committed runs (results/rollouts, or FLOOD_RUNS_DIR)
    /api/tasks/{id}/replay?plan=...  a plan on the true river: releases, storage, gauge flows, flood categories, cost
    /twin/                           the 3D replay of the basin (USGS terrain and hydrography) for any plan

    uvicorn flood_openenv.server:app --port 8000
"""

from __future__ import annotations

import os
from pathlib import Path

from fastapi import HTTPException
from fastapi.responses import RedirectResponse
from fastapi.staticfiles import StaticFiles
from openenv.core.env_server.http_server import create_app
from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation

from flood_core import TaskPack, load_pack

try:
    from .environment import STORE, FloodOpsEnvironment, FloodState
except ImportError:
    from environment import STORE, FloodOpsEnvironment, FloodState

ENV_NAME = "flood_ops"
WEB = Path(__file__).resolve().parent / "web"
_UP = Path(__file__).resolve().parents
RUNS_DIR = Path(os.environ.get("FLOOD_RUNS_DIR") or (_UP[4] if len(_UP) > 4 else _UP[-1]) / "results" / "rollouts")


def mount_api(app, pack: TaskPack):
    @app.get("/api/tasks")
    def tasks():
        return [{"task_id": t["task_id"], "split": t["split"], "tier": t["tier"], "start_local": t["start_local"],
                 "steps": t["steps"], **{k: v for k, v in t["reference"].items() if k.endswith("_cost")}}
                for t in pack.tasks]

    def _get(tid: str) -> dict:
        try:
            return pack.get(tid)
        except KeyError:
            raise HTTPException(404, f"no task {tid}") from None

    @app.get("/api/tasks/{tid}")
    def task(tid: str):
        return pack.public(_get(tid))

    @app.get("/api/tasks/{tid}/reference")
    def reference(tid: str):
        return _get(tid)["reference"]

    @app.get("/api/episodes")
    def episodes(n: int = 50):
        return [e.summary() for e in STORE.recent(n)]

    from flood_core.replay import plans as _plans
    from flood_core.replay import replay as _replay

    @app.get("/api/tasks/{tid}/plans")
    def task_plans(tid: str):
        return _plans(_get(tid), RUNS_DIR)

    @app.get("/api/tasks/{tid}/replay")
    def task_replay(tid: str, plan: str = "optimum"):
        try:
            return _replay(_get(tid), plan, RUNS_DIR)
        except KeyError as e:
            raise HTTPException(404, str(e)) from None

    @app.get("/twin", include_in_schema=False)
    def twin_slash():
        return RedirectResponse("/twin/")

    app.mount("/twin/", StaticFiles(directory=WEB / "twin", html=True), name="twin")

    @app.middleware("http")
    async def no_stale_pages(request, call_next):
        # browsers cache static pages without Cache-Control and keep showing an old replay after a deploy
        resp = await call_next(request)
        if request.url.path.startswith("/twin"):
            resp.headers["Cache-Control"] = "no-cache"
        return resp


def build_app(pack: TaskPack | None = None, mode: str | None = None):
    """FLOOD_MODE=live serves live mode (forecast scenarios, block commits); the default is plan mode."""
    pack = pack or load_pack()
    mode = mode or os.environ.get("FLOOD_MODE", "plan")
    if mode == "live":
        try:
            from .live_env import FloodLiveEnvironment as Env
        except ImportError:
            from live_env import FloodLiveEnvironment as Env
    else:
        Env = FloodOpsEnvironment

    def factory():
        return Env(pack)

    kw = {}
    if os.environ.get("ENABLE_WEB_INTERFACE", "").lower() == "true":
        try:
            from .gradio_ui import build_ui
        except ImportError:
            from gradio_ui import build_ui
        kw = dict(gradio_builder=build_ui, custom_tab_name="Explore", custom_tab_primary=True, show_default_tab=True,
                  title_override=f"FloodGate v1 ({mode} mode)")
    app = create_app(factory, CallToolAction, CallToolObservation, env_name=ENV_NAME,
                     max_concurrent_envs=int(os.environ.get("MAX_CONCURRENT_ENVS", "64")), state_cls=FloodState, **kw)
    mount_api(app, pack)

    @app.get("/healthz", include_in_schema=False)
    def healthz():
        return {"status": "ok", "mode": mode, "tasks": {s: pack.count(s) for s in pack.splits()}}

    if not kw:
        @app.get("/", include_in_schema=False)
        def root():
            return RedirectResponse("/twin/")

    return app


_app = None


def __getattr__(name: str):
    global _app
    if name == "app":
        if _app is None:
            _app = build_app()
        return _app
    raise AttributeError(name)