api/main.py

"""
FastAPI Server — REST API for the DevOps RL Agent.

Endpoints for running episodes, viewing replays, checking stats,
and triggering training steps.
"""

from __future__ import annotations

import os
import asyncio
import threading
import uuid
from pathlib import Path
from typing import Dict, List, Optional

from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from agent.baseline_agent import BaselineAgent
from agent.devops_agent import DevOpsAgent
from devops_env.env import DevOpsEnv
from replay.buffer import ReplayBuffer
from scenarios.registry import ScenarioRegistry
from training.curriculum import CurriculumScheduler

# --- App Setup ---
app = FastAPI(
    title="DevOps RL Agent API",
    description="REST API for the reinforcement-learning-powered terminal troubleshooting agent.",
    version="1.0.0",
)

# Serve frontend static files
FRONTEND_DIR = Path(__file__).parent.parent / "frontend"
if FRONTEND_DIR.exists():
    app.mount("/app", StaticFiles(directory=str(FRONTEND_DIR), html=True), name="frontend")

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

# --- Shared State ---
DB_URL = os.environ.get("REPLAY_DB_URL", "sqlite:///replay_buffer.db")
replay_buffer = ReplayBuffer(DB_URL)
registry = ScenarioRegistry()
registry.register_defaults()
curriculum = CurriculumScheduler()
agent = DevOpsAgent(model_name="rule-based")
openenv_sessions: Dict[str, DevOpsEnv] = {}
openenv_lock = threading.Lock()


# --- Request/Response Models ---
class RunEpisodeRequest(BaseModel):
    """Request body for running an episode."""
    scenario_id: Optional[str] = None
    level: Optional[int] = None


class TrainStepRequest(BaseModel):
    """Request body for triggering a training step."""
    num_episodes: int = 10
    level: Optional[int] = None


class EpisodeResponse(BaseModel):
    """Response for an episode run."""
    episode_id: str
    scenario_id: str
    level: int
    solved: bool
    total_reward: float
    total_steps: int
    steps: List[Dict]


class OpenEnvAction(BaseModel):
    """Structured action payload for OpenEnv-style stepping."""
    command: str


class OpenEnvResetRequest(BaseModel):
    """Request for starting a new OpenEnv session."""
    scenario_id: Optional[str] = None
    level: Optional[int] = None
    max_steps: int = 10


class OpenEnvStepRequest(BaseModel):
    """Request for stepping an existing OpenEnv session."""
    session_id: str
    action: OpenEnvAction


class OpenEnvCloseRequest(BaseModel):
    """Request for closing an OpenEnv session."""
    session_id: str


def _openenv_pop_session(session_id: str) -> DevOpsEnv | None:
    """Remove and return an OpenEnv session from the in-memory store."""
    with openenv_lock:
        return openenv_sessions.pop(session_id, None)


def _openenv_get_session(session_id: str) -> DevOpsEnv | None:
    """Get an OpenEnv session without removing it."""
    with openenv_lock:
        return openenv_sessions.get(session_id)


# --- Endpoints ---
@app.get("/")
async def root():
    """Health check endpoint."""
    return {
        "service": "DevOps RL Agent API",
        "status": "running",
        "version": "1.0.0",
    }


@app.post("/episode/run")
async def run_episode(request: RunEpisodeRequest):
    """Run one episode with the current agent.

    Returns the full episode log including step-by-step
    observations, actions, rewards, and error classifications.
    """
    env = None
    try:
        env = DevOpsEnv(
            scenario_registry=registry,
            target_level=request.level,
            target_scenario=request.scenario_id,
        )

        obs, info = env.reset()
        steps = []
        total_reward = 0.0
        done = False
        step_num = 0

        while not done:
            step_num += 1
            action = agent.act(obs)
            obs, reward, terminated, truncated, step_info = env.step(action)
            total_reward += reward

            steps.append({
                "step": step_num,
                "action": action,
                "observation": {
                    "error_log": obs.get("error_log", "")[:500],
                    "command_history": obs.get("command_history", []),
                    "step_count": obs.get("step_count", 0),
                },
                "reward": round(reward, 2),
                "reward_breakdown": {k: round(v, 2) for k, v in step_info.get("reward_breakdown", {}).items()},
                "error_type": obs.get("error_type", "unknown"),
                "execution_result": step_info.get("execution_result", {}),
                "solved": step_info.get("solved", False),
            })

            done = terminated or truncated

        summary = env.get_episode_summary()

        # Store in replay buffer
        episode_id = replay_buffer.store_episode(
            scenario_id=summary["scenario_id"],
            level=summary["level"],
            steps=steps,
            total_reward=total_reward,
            solved=summary["solved"],
        )

        return {
            "episode_id": episode_id,
            "scenario_id": summary["scenario_id"],
            "level": summary["level"],
            "solved": summary["solved"],
            "total_reward": round(total_reward, 2),
            "total_steps": step_num,
            "steps": steps,
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if env is not None:
            env.close()


@app.post("/reset")
async def openenv_reset(request: OpenEnvResetRequest):
    """OpenEnv-compatible reset endpoint.

    Creates a server-managed environment session and returns
    the initial observation.
    """
    env = None
    try:
        env = DevOpsEnv(
            scenario_registry=registry,
            target_level=request.level,
            target_scenario=request.scenario_id,
            max_steps=request.max_steps,
        )
        options = {"scenario_id": request.scenario_id} if request.scenario_id else None
        observation, info = env.reset(options=options)
        session_id = str(uuid.uuid4())

        with openenv_lock:
            openenv_sessions[session_id] = env

        return {
            "session_id": session_id,
            "observation": observation,
            "info": info,
        }
    except Exception as e:
        if env is not None:
            env.close()
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/step")
async def openenv_step(request: OpenEnvStepRequest):
    """OpenEnv-compatible step endpoint for a server-managed session."""
    env = _openenv_get_session(request.session_id)
    if env is None:
        raise HTTPException(status_code=404, detail=f"Session {request.session_id} not found")

    try:
        observation, reward, terminated, truncated, info = env.step(request.action.command)
        done = terminated or truncated

        if done:
            session = _openenv_pop_session(request.session_id)
            if session is not None:
                session.close()

        return {
            "session_id": request.session_id,
            "observation": observation,
            "reward": reward,
            "terminated": terminated,
            "truncated": truncated,
            "done": done,
            "info": info,
        }
    except RuntimeError as e:
        # RuntimeError generally means terminal episode; clean up stale session.
        session = _openenv_pop_session(request.session_id)
        if session is not None:
            session.close()
        raise HTTPException(status_code=409, detail=str(e))
    except Exception as e:
        session = _openenv_pop_session(request.session_id)
        if session is not None:
            session.close()
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/close")
async def openenv_close(request: OpenEnvCloseRequest):
    """Close and remove an OpenEnv session explicitly."""
    env = _openenv_pop_session(request.session_id)
    if env is None:
        raise HTTPException(status_code=404, detail=f"Session {request.session_id} not found")
    env.close()
    return {"session_id": request.session_id, "closed": True}


@app.get("/episode/{episode_id}")
async def get_episode(episode_id: str):
    """Get a stored episode by its UUID."""
    episode = replay_buffer.get_episode(episode_id)
    if not episode:
        raise HTTPException(status_code=404, detail=f"Episode {episode_id} not found")
    return episode


@app.get("/stats")
async def get_stats():
    """Get aggregate statistics: solve rates, mean rewards, training progress."""
    stats = replay_buffer.get_stats()
    stats["curriculum"] = curriculum.get_status()

    # Update curriculum from stats
    for lvl in [1, 2, 3]:
        if lvl in stats.get("levels", {}):
            lvl_stats = stats["levels"][lvl]
            curriculum.update_stats(
                level=lvl,
                solve_rate=lvl_stats["solve_rate"],
                episodes=lvl_stats["count"],
            )

    return stats


@app.get("/replay/{episode_id}")
async def get_replay(episode_id: str):
    """Get step-by-step replay data for an episode.

    Returns formatted data optimized for the Replay Viewer frontend.
    """
    episode = replay_buffer.get_episode(episode_id)
    if not episode:
        raise HTTPException(status_code=404, detail=f"Episode {episode_id} not found")

    return {
        "episode_id": episode["episode_id"],
        "scenario_id": episode["scenario_id"],
        "level": episode["level"],
        "solved": episode["solved"],
        "total_reward": episode["total_reward"],
        "total_steps": episode["total_steps"],
        "timestamp": episode["timestamp"],
        "steps": episode["steps"],
    }


@app.post("/train/step")
async def trigger_training_step(request: TrainStepRequest):
    """Trigger a batch of training rollout episodes.

    Runs the specified number of episodes and returns aggregate results.
    """
    results = []
    for _ in range(request.num_episodes):
        env = None
        try:
            env = DevOpsEnv(
                scenario_registry=registry,
                target_level=request.level if request.level is not None else curriculum.sample_level(),
            )
            obs, info = env.reset()
            total_reward = 0.0
            done = False
            steps = []

            while not done:
                action = agent.act(obs)
                obs, reward, terminated, truncated, step_info = env.step(action)
                total_reward += reward
                steps.append({
                    "step": step_info.get("step_count", len(steps) + 1),
                    "action": action,
                    "reward": reward,
                    "reward_breakdown": step_info.get("reward_breakdown", {}),
                    "error_type": obs.get("error_type", "unknown"),
                    "observation": {"error_log": obs.get("error_log", "")[:300]},
                    "result": step_info.get("execution_result", {}),
                })
                done = terminated or truncated

            summary = env.get_episode_summary()

            ep_id = replay_buffer.store_episode(
                scenario_id=summary["scenario_id"],
                level=summary["level"],
                steps=steps,
                total_reward=total_reward,
                solved=summary["solved"],
            )
            results.append({
                "episode_id": ep_id,
                "scenario_id": summary["scenario_id"],
                "solved": summary["solved"],
                "total_reward": round(total_reward, 2),
            })
        except Exception as e:
            results.append({"error": str(e)})
        finally:
            if env is not None:
                env.close()

    return {
        "episodes_run": len(results),
        "episodes_solved": sum(1 for r in results if r.get("solved", False)),
        "mean_reward": round(
            sum(r.get("total_reward", 0) for r in results) / max(len(results), 1), 2
        ),
        "results": results,
    }


@app.get("/scenarios")
async def list_scenarios():
    """List all available scenarios with their solve rates."""
    scenarios = []
    stats = replay_buffer.get_stats()
    scenario_stats = stats.get("scenarios", {})

    for scenario in registry.get_all():
        sc_stats = scenario_stats.get(scenario.id, {})
        scenarios.append({
            "id": scenario.id,
            "level": scenario.level,
            "description": scenario.description,
            "hint_commands": scenario.hint_commands,
            "solve_rate": sc_stats.get("solve_rate", 0.0),
            "attempts": sc_stats.get("count", 0),
        })

    return {"scenarios": scenarios}


@app.get("/recent")
async def get_recent_episodes(n: int = Query(default=20, le=100)):
    """Get the most recent episodes."""
    return {"episodes": replay_buffer.get_recent(n)}