src/server/app.py

"""OpenEnv server implementing reset/step/state endpoints."""

from typing import Any

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from pydantic import BaseModel

from src.models import Action, Observation, State
from src.openenv_environment import OpenEnvEnvironment

app = FastAPI(title="911 — Dispatch API")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

_env: OpenEnvEnvironment | None = None


# Removed ResetRequest since /reset now dynamically parses the Request to handle null bodies gracefully.

class StepRequest(BaseModel):
    action: dict[str, Any]


class StepResponse(BaseModel):
    observation: dict[str, Any]
    reward: float
    done: bool


@app.exception_handler(RuntimeError)
async def runtime_error_handler(request, exc: RuntimeError):
    from fastapi.responses import JSONResponse

    return JSONResponse(status_code=500, content={"detail": str(exc)})


@app.get("/", include_in_schema=False)
async def root():
    """Serve the live dashboard on the root route for HF Spaces."""
    import os
    current_file = os.path.abspath(__file__)
    base_dir = os.path.dirname(os.path.dirname(os.path.dirname(current_file)))
    dashboard_path = os.path.join(base_dir, "live_dashboard.html")
    
    # Fallback to current working directory if not found
    if not os.path.exists(dashboard_path):
        dashboard_path = os.path.join(os.getcwd(), "live_dashboard.html")

    if os.path.exists(dashboard_path):
        return FileResponse(dashboard_path)
        
    return JSONResponse({
        "status": "healthy", 
        "error": "dashboard not found",
        "debug": dashboard_path
    })


@app.get("/health")
async def health() -> dict[str, str]:
    # OpenEnv runtime validation expects status=healthy
    return {"status": "healthy"}


@app.get("/metadata")
async def metadata() -> dict[str, Any]:
    """OpenEnv metadata endpoint used by runtime validators."""

    return {
        "name": "citywide-dispatch-supervisor",
        "description": (
            "City-wide 911 emergency dispatch supervisor RL environment. "
            "An LLM agent learns to manage simultaneous incidents by dispatching "
            "police, fire, and EMS units across a city grid under realistic constraints."
        ),
        "version": "0.1.0",
        "mode": "simulation",
    }


@app.get("/schema")
async def schema() -> dict[str, Any]:
    """Return JSON schemas for Action/Observation/State."""

    return {
        "action": Action.model_json_schema(),
        "observation": Observation.model_json_schema(),
        "state": State.model_json_schema(),
    }


@app.post("/mcp")
async def mcp_endpoint(request: Request):
    """MCP JSON-RPC passthrough for OpenEnv runtime compatibility."""
    try:
        body = await request.json()
    except Exception:
        return JSONResponse({"error": "invalid JSON"}, status_code=400)

    method = body.get("method", "")
    req_id = body.get("id", 1)

    if method == "reset":
        result = await _env.reset()
        return {"jsonrpc": "2.0", "id": req_id, "result": result.model_dump()}
    elif method == "step":
        action_data = body.get("params", {}).get("action", {})
        action = Action(**action_data)
        obs, reward, done = await _env.step(action)
        return {"jsonrpc": "2.0", "id": req_id, "result": {"observation": obs.model_dump(), "reward": reward, "done": done}}
    elif method == "state":
        result = _env.state()
        return {"jsonrpc": "2.0", "id": req_id, "result": result.model_dump()}
    elif method == "legal_actions":
        actions = _env.legal_actions()
        return {"jsonrpc": "2.0", "id": req_id, "result": [a.model_dump() for a in actions]}
    else:
        return JSONResponse({"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Method not found: {method}"}}, status_code=404)


@app.get("/tasks")
async def list_tasks() -> list[dict[str, str]]:
    """List all available tasks."""
    from src.tasks.registry import TaskRegistry

    return [
        {
            "task_id": t.task_id,
            "name": t.name,
            "description": t.description,
            "difficulty": t.difficulty,
        }
        for t in TaskRegistry.list_tasks()
    ]


@app.post("/reset")
async def reset(request: Request) -> dict[str, Any]:
    try:
        body = await request.json()
    except Exception:
        body = {}
    if body is None:
        body = {}
        
    task_id = body.get("task_id", "single_incident")
    seed = body.get("seed", None)

    global _env
    _env = OpenEnvEnvironment(task_id=task_id, seed=seed)
    obs = await _env.reset()
    return obs.model_dump()


@app.post("/step")
async def step(request: StepRequest) -> StepResponse:
    if _env is None:
        raise RuntimeError("Environment not initialized. Call /reset first.")
    try:
        action = Action.model_validate(request.action)
    except Exception as e:
        raise RuntimeError(f"Invalid action: {e}")
    obs, reward, done = await _env.step(action)
    return StepResponse(
        observation=obs.model_dump(),
        reward=reward,
        done=done,
    )


@app.get("/state")
async def get_state() -> dict[str, Any]:
    if _env is None:
        raise RuntimeError("Environment not initialized. Call /reset first.")
    state = _env.state()
    return state.model_dump()


@app.get("/dashboard/state")
async def get_dashboard_state() -> dict[str, Any]:
    """Extended state for the HTML live dashboard.

    Keeps the existing /state response stable for typed clients.
    """
    if _env is None:
        # Return an empty but valid structure before /reset is called
        return {
            "units": {},
            "incidents": {},
            "episode_id": "not-initialized",
            "step_count": 0,
            "task_id": "none",
            "city_time": 0.0,
            "metadata": {},
            "legal_actions": [],
            "issues": [],
            "observation": None,
        }

    state_dict = _env.state().model_dump()
    legal_actions = [a.model_dump() for a in _env.legal_actions()]
    last_obs = _env.last_observation()
    issues = list(last_obs.issues) if last_obs is not None else []
    obs_dict = last_obs.model_dump() if last_obs is not None else None

    return {
        **state_dict,
        "legal_actions": legal_actions,
        "issues": issues,
        "observation": obs_dict,
    }


def main():
    import uvicorn
    import os

    port = int(os.environ.get("PORT", "7860"))
    uvicorn.run("src.server.app:app", host="0.0.0.0", port=port, reload=False)


if __name__ == "__main__":
    main()