src/driving_env/server.py

"""FastAPI HTTP server exposing the OpenEnv API."""
from __future__ import annotations

import os
from typing import Optional
from fastapi import FastAPI, HTTPException, Body

from .env import DrivingEnv
from .models import (
    ResetRequest,
    StepRequest,
    StepResponse,
    StateResponse,
    SpecResponse,
)

app = FastAPI(title="Driving OpenEnv", version="0.1.0")
_env = DrivingEnv()


@app.get("/")
def root() -> dict:
    return {
        "name": "driving_sim",
        "version": "0.1.0",
        "endpoints": ["/reset", "/step", "/state", "/spec", "/health"],
        "api_base_url": os.environ.get("API_BASE_URL", ""),
        "model_name": os.environ.get("MODEL_NAME", ""),
    }


@app.get("/health")
def health() -> dict:
    return {"status": "ok"}


@app.get("/spec", response_model=SpecResponse)
def spec() -> SpecResponse:
    return SpecResponse()


@app.post("/reset", response_model=StateResponse)
def reset(req: Optional[ResetRequest] = Body(default=None)) -> StateResponse:
    req = req or ResetRequest()
    try:
        _env.reset(task=req.task, seed=req.seed)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    return _env.state()


@app.get("/reset", response_model=StateResponse)
def reset_get(task: str = "lane_follow", seed: Optional[int] = None) -> StateResponse:
    try:
        _env.reset(task=task, seed=seed)  # type: ignore[arg-type]
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    return _env.state()


@app.post("/step", response_model=StepResponse)
def step(req: Optional[StepRequest] = Body(default=None)) -> StepResponse:
    if req is None:
        raise HTTPException(status_code=400, detail="Request body with 'action' is required.")
    return _env.step(req.action)


@app.get("/state", response_model=StateResponse)
def state() -> StateResponse:
    return _env.state()