main.py

# main.py — OpenEnv-Compliant FastAPI Server
# Autonomy Calibration Environment v1.0
# OpenEnv India Hackathon 2026 — by Rhythm

from __future__ import annotations
import os
from collections import Counter
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import subprocess
import sys
from pydantic import BaseModel

from models import Action, Observation, Reward, StepResult, ResetRequest
from environment.scenarios import SCENARIOS

def _build_registry():
    from tasks.email_triage import EmailTriageTask
    from tasks.devops_incident import DevOpsIncidentTask
    from tasks.financial_request import FinancialRequestTask
    return {
        "email_triage": EmailTriageTask,
        "devops_incident": DevOpsIncidentTask,
        "financial_request": FinancialRequestTask,
    }

TASK_REGISTRY = _build_registry()

# ─── App ─────────────────────────────────────────────────────────────────────

app = FastAPI(
    title="Autonomy Calibration Environment",
    description=(
        "OpenEnv-compliant RL environment training agents to calibrate autonomy "
        "across Email Triage, DevOps Incident Response, and Financial Request Handling."
    ),
    version="2.0.0",
)

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

# ─── Session State ────────────────────────────────────────────────────────────
# Single-session in-memory state. Sufficient for hackathon + HF Spaces.

_session: dict = {
    "task_name": None,
    "task": None,
    "step": 0,
    "history": [],
    "done": True,
    "seed": None,
    "episode_id": None,
}

# Global episode log for /api/history
_episode_log: list[dict] = []


# ─── Helpers ──────────────────────────────────────────────────────────────────

def _get_task(name: str):
    if name not in TASK_REGISTRY:
        raise HTTPException(
            status_code=400,
            detail=f"Unknown task '{name}'. Valid: {list(TASK_REGISTRY.keys())}"
        )
    return TASK_REGISTRY[name]()


# ─── API: Reset ───────────────────────────────────────────────────────────────

@app.post("/reset")
@app.post("/api/reset", response_model=Observation)
def reset(body: ResetRequest = ResetRequest()):
    try:
        task = _get_task(body.task)
        obs = task.reset(seed=body.seed)
        # Store seed in session and observation
        obs.seed = body.seed
        _session["task_name"] = body.task
        _session["task"] = task
        _session["step"] = 0
        _session["history"] = []
        _session["done"] = False
        _session["seed"] = body.seed
        # Create DB episode and store ID
        import database as db
        _session["episode_id"] = db.create_episode(body.task, body.seed)
        return obs
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── API: Step ───────────────────────────────────────────────────────────────

@app.post("/step")
@app.post("/api/step", response_model=StepResult)
async def step_env(action: Action):
    task = _session.get("task")
    if task is None or _session.get("done"):
        raise HTTPException(status_code=400, detail="No active episode. Call /api/reset first.")
    try:
        obs, reward, done, info = task.step(action)
        step_idx = _session["step"]
        _session["step"] += 1
        _session["done"] = done
        step_entry = {
            "step": step_idx,
            "action": action.type,
            "reward": reward.value,
            "done": done,
        }
        _session["history"].append(step_entry)
        # Persist step to SQLite
        import database as db
        db.log_step(
            episode_id=_session["episode_id"],
            step_index=step_idx,
            decision=action.type,
            reward=reward.value,
            done=done,
        )
        if done:
            episode_score = info.get("episode_score")
            db.close_episode(_session["episode_id"], episode_score or 0.0)
            _episode_log.append({
                "episode_id": _session["episode_id"],
                "task": _session["task_name"],
                "seed": _session["seed"],
                "episode_score": episode_score,
                "steps": _session["step"],
                "history": list(_session["history"]),
            })
        return StepResult(observation=obs, reward=reward, done=done, info=info)
    except RuntimeError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── API: State ───────────────────────────────────────────────────────────────

@app.get("/api/state")
def state():
    task = _session.get("task")
    if task is None:
        return {"status": "not_started"}
    return {
        "status": "done" if _session["done"] else "active",
        "task": _session["task_name"],
        "step": _session["step"],
        **task.state(),
    }

# ─── API: Training ────────────────────────────────────────────────────────────

def run_training():
    """Runs the training script in a separate process and pipes output to logs."""
    try:
        print("🚀 GRPO TRAINING CORE: Initializing...")
        # Redirect stdout and stderr to the main process ones so they appear in HF Logs
        process = subprocess.Popen(
            [sys.executable, "train_rl.py"],
            stdout=sys.stdout, 
            stderr=sys.stderr,
            bufsize=1, # Line buffered
            universal_newlines=True
        )
        print(f"✅ Background process PID {process.pid} spawned.")
    except Exception as e:
        print(f"❌ Error during background training: {e}")

@app.post("/api/train")
def start_training(background_tasks: BackgroundTasks):
    # Basic check for GPU presence (useful for logs)
    try:
        import torch
        has_gpu = torch.cuda.is_available()
        device_name = torch.cuda.get_device_name(0) if has_gpu else "CPU"
    except ImportError:
        has_gpu = False
        device_name = "CPU"
    
    background_tasks.add_task(run_training)
    
    return {
        "status": "started",
        "message": "GRPO Training started in background.",
        "using_gpu": has_gpu,
        "device": device_name
    }

@app.post("/api/upload")
def upload_to_hub(repo_id: str = "JOY0021/autonomy-agent-v2"):
    """Pushes the trained folder to the HF Hub model repo, creating it if needed."""
    try:
        import os
        from huggingface_hub import HfApi, create_repo
        
        token = os.getenv("HF_TOKEN")
        api = HfApi(token=token)
        
        # 1. Create repo if it doesn't exist
        print(f"📦 Ensuring repo {repo_id} exists...")
        create_repo(repo_id=repo_id, repo_type="model", exist_ok=True, token=token)
        
        # 2. Upload the folder
        print(f"📡 Uploading autonomy-agent-v2 to {repo_id}...")
        api.upload_folder(
            folder_path="autonomy-agent-v2",
            repo_id=repo_id,
            repo_type="model",
        )
        return {"status": "success", "message": f"Model live at https://huggingface.co/{repo_id}"}
    except Exception as e:
        print(f"❌ Upload Error: {e}")
        return {"status": "error", "message": str(e)}


# ─── API: Health ─────────────────────────────────────────────────────────────

@app.get("/api/health")
def health():
    difficulty_dist = Counter(s["difficulty"] for s in SCENARIOS)
    decision_dist = Counter(s["best_decision"] for s in SCENARIOS)
    return {
        "status": "ok",
        "environment": "autonomy-calibration-env",
        "version": "2.0.0",
        "tasks": list(TASK_REGISTRY.keys()),
        "autonomy_action_space": ["ACT", "ASK", "STOP", "RECOVER"],
        "autonomy_scenarios": len(SCENARIOS),
        "autonomy_difficulty_distribution": dict(difficulty_dist),
        "autonomy_decision_distribution": dict(decision_dist),
        "reward_range": [0.01, 0.99],
    }


# ─── API: History ─────────────────────────────────────────────────────────────

@app.get("/api/history")
def history():
    total = len(_episode_log)
    scores = [e["episode_score"] for e in _episode_log if e.get("episode_score") is not None]
    return {
        "total_episodes": total,
        "avg_score": round(sum(scores) / len(scores), 4) if scores else 0.0,
        "episodes": _episode_log,
    }


@app.delete("/api/history")
def clear_history():
    _episode_log.clear()
    return {"status": "cleared"}


# ─── API: Observability (Step 4) ─────────────────────────────────────────────

@app.get("/api/episodes")
def episodes(limit: int = 20):
    """List recent episodes from SQLite with metadata."""
    import database as db
    try:
        rows = db.list_episodes(limit=limit)
        return {"episodes": rows, "count": len(rows)}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/replay/{episode_id}")
def replay(episode_id: int):
    """
    Return the full step history for a past episode.
    Can be used to reproduce the episode by feeding steps back into reset + step.
    """
    import database as db
    try:
        data = db.get_episode(episode_id)
        return {
            "episode_id": episode_id,
            "episode": data["episode"],
            "steps": data["steps"],
            "total_steps": len(data["steps"]),
        }
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/grade")
def grade_current():
    """Run deterministic grader on the current session's episode history."""
    task = _session.get("task")
    if task is None:
        raise HTTPException(status_code=400, detail="No active episode.")
    score = task.grade_episode(_session["history"])
    return {
        "task": _session["task_name"],
        "seed": _session["seed"],
        "episode_id": _session["episode_id"],
        "score": score,
        "steps_completed": _session["step"],
        "done": _session["done"],
    }


@app.get("/api/grade/{episode_id}")
def grade_episode(episode_id: int):
    """Run deterministic grader on a completed historical episode."""
    import database as db
    try:
        data = db.get_episode(episode_id)
        ep = data["episode"]
        steps = data["steps"]
        # Reconstruct history format expected by grade_episode()
        history = [
            {"step": s["step_index"], "action": s["decision"],
             "reward": {"value": s["reward"]}}
            for s in steps
        ]
        total_reward = sum(s["reward"] for s in steps)
        from utils import clamp
        score = clamp(total_reward)
        return {
            "episode_id": episode_id,
            "task": ep["task"],
            "seed": ep["seed"],
            "score": score,
            "total_steps": len(steps),
            "started_at": ep["started_at"],
            "ended_at": ep["ended_at"],
            "steps": steps,
        }
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))



# ─── Static UI ────────────────────────────────────────────────────────────────

if os.path.exists("static"):
    app.mount("/static", StaticFiles(directory="static"), name="static")

    @app.get("/")
    def serve_ui():
        return FileResponse("static/index.html")
else:
    @app.get("/")
    def serve_fallback():
        return {
            "message": "Autonomy Calibration Environment API v1.0",
            "docs": "/docs",
            "tasks": list(TASK_REGISTRY.keys()),
        }
def main():
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)

if __name__ == "__main__":
    main()