server/app.py

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Annotated, Any

from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import HTMLResponse
from pydantic import BaseModel, Field

from server.environment import SecretsAuditEnvironment


app = FastAPI(title="SecretsAuditEnv", version="4.0")
environment = SecretsAuditEnvironment(repo_root=Path(__file__).resolve().parents[1])


class ResetRequest(BaseModel):
    task_id: int | str | None = None
    id: int | str | None = None
    difficulty: str | None = None


class StepRequest(BaseModel):
    action: str = Field(..., min_length=1)


def _parse_task_id(raw: Any) -> int:
    if raw is None:
        return 1
    if isinstance(raw, int):
        task_id = raw
    else:
        text = str(raw).strip()
        if text.startswith("task_"):
            text = text.split("_", 1)[1]
        try:
            task_id = int(text)
        except ValueError as exc:
            raise HTTPException(status_code=422, detail="Invalid task_id value.") from exc
    if not 1 <= task_id <= 13:
        raise HTTPException(status_code=422, detail="task_id must be between 1 and 13.")
    return task_id


_DIFFICULTY_TO_TASK_ID = {"easy": 1, "medium": 6, "hard": 11}


@app.get("/")
def root() -> dict:
    return {
        "name": "SecretsAuditEnv",
        "version": "4.0",
        "status": "ok",
        "docs": "/docs",
        "web_ui": "/web",
        "tasks_endpoint": "/tasks",
    }


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


@app.get("/state")
def get_state() -> dict:
    return environment.state()


@app.get("/tasks")
def get_tasks() -> dict:
    return {"tasks": environment.list_tasks()}


@app.post("/reset")
def reset(
    request: ResetRequest | None = None,
    task_id: Annotated[str | int | None, Query()] = None,
) -> dict:
    try:
        raw_task_id: Any = task_id
        if raw_task_id is None and request is not None:
            if request.difficulty and request.task_id is None and request.id is None:
                raw_task_id = _DIFFICULTY_TO_TASK_ID.get(request.difficulty, 1)
            else:
                raw_task_id = request.task_id if request.task_id is not None else request.id
        resolved_task_id = _parse_task_id(raw_task_id)
        return environment.reset(resolved_task_id)
    except FileNotFoundError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc


@app.post("/step")
def step(request: StepRequest) -> dict:
    try:
        state = environment.step(request.action)
        session = (state.get("session") or {})
        return {
            **state,
            "reward": session.get("reward", 0.0),
            "cumulative_reward": session.get("reward", 0.0),
            "done": float(session.get("reward", 0.0)) >= 1.0,
            "step": session.get("step_count", 0),
            "observation": session.get("observation", ""),
        }
    except RuntimeError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@app.get("/web", response_class=HTMLResponse)
async def web_ui() -> str:
    html_path = Path(__file__).parent / "web_ui.html"
    return html_path.read_text(encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description="Run the SecretsAuditEnv server.")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=7860)
    args = parser.parse_args()

    import uvicorn

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


if __name__ == "__main__":
    main()