app/main.py

"""
Meta-SRE FastAPI Server – OpenEnv Standard API.

Implements the OpenEnv contract exactly:
  POST /reset  → Observation
  POST /step   → (observation, reward, done, info)
  GET  /state  → Observation
  GET  /grade  → EpisodeResult
  GET  /tools  → tool specs (JSON Schema per tool)
  GET  /tasks  → task definitions
  GET  /health → liveness probe

The /env/* routes are strict OpenEnv aliases used by openenv_client.connect().
"""
from __future__ import annotations

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Any, Dict, Optional

from app.engine.manager import EpisodeManager, TASK_DEFINITIONS, DifficultyController
from app.models import Observation, ActionResult, EpisodeResult
from app.tools.definitions import TOOL_SPECS


# ---------------------------------------------------------------------------
# App setup
# ---------------------------------------------------------------------------

app = FastAPI(
    title="Meta-SRE",
    description=(
        "OpenEnv environment: train LLM agents to act as Senior Site Reliability Engineers "
        "debugging realistic Meta production incidents."
    ),
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Single global episode manager (stateful server)
_dc = DifficultyController()
_episode = EpisodeManager(difficulty_controller=_dc)
_started = False


# ---------------------------------------------------------------------------
# Request / response models
# ---------------------------------------------------------------------------

class ResetRequest(BaseModel):
    task_id: Optional[int] = None   # 1-5; None = difficulty-controller picks


class StepRequest(BaseModel):
    tool:   str
    params: Dict[str, Any] = {}


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.get("/health")
def health():
    return {"status": "ok", "version": "1.0.0"}


@app.get("/tools")
def list_tools():
    return {"tools": TOOL_SPECS, "count": len(TOOL_SPECS)}


@app.get("/tasks")
def list_tasks():
    return {"tasks": TASK_DEFINITIONS}


@app.post("/reset", response_model=Observation)
def reset(req: ResetRequest = ResetRequest()):
    global _started
    _started = True
    obs = _episode.reset(task_id=req.task_id)
    return obs


@app.post("/step", response_model=ActionResult)
def step(req: StepRequest):
    global _started
    if not _started:
        raise HTTPException(
            status_code=400,
            detail="Episode not started. Call POST /reset first."
        )
    try:
        result = _episode.step(tool=req.tool, params=req.params)
        return result
    except RuntimeError as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/state", response_model=Observation)
def get_state():
    if not _started:
        raise HTTPException(
            status_code=400,
            detail="Episode not started. Call POST /reset first."
        )
    return _episode._build_observation()


@app.get("/grade", response_model=EpisodeResult)
def grade():
    if not _started:
        raise HTTPException(
            status_code=400,
            detail="Episode not started. Call POST /reset first."
        )
    return _episode.get_episode_result()


# ---------------------------------------------------------------------------
# OpenEnv compatibility shim (env.reset / env.step / env.grade)
# ---------------------------------------------------------------------------

@app.post("/env/reset")
def env_reset(req: ResetRequest = ResetRequest()):
    """OpenEnv spec alias for /reset."""
    return reset(req)


@app.post("/env/step")
def env_step(req: StepRequest):
    """
    OpenEnv standard step — returns the canonical 4-tuple:
      (observation, reward, done, info)
    This is what openenv_client.connect().step() unpacks.
    """
    result: ActionResult = step(req)
    return {
        "observation": result.observation,
        "reward":      result.reward_delta,
        "done":        result.done,
        "info": {
            "tool":            result.tool,
            "output":          result.output,
            "episode_id":      _episode._incident_id,
            "step":            _episode._step,
            "budget_remaining": max(0, _episode._build_observation().budget_remaining),
        },
    }


# ---------------------------------------------------------------------------
# Dev entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=True)