environment.py

import uuid, json
from typing import Any, Dict, List, Optional
from fastapi import FastAPI, Query
from pydantic import BaseModel, Field
from tasks import (
    TASKS, TASK_SEQUENCE,
    grade_easy, grade_medium, grade_hard,
    grade_medium_lambda, grade_hard_rds     # NEW
)

app = FastAPI(title="AWS Security Auditor", version="1.0.0")

class AuditAction(BaseModel):
    findings:        List[str] = Field(...)
    severity:        List[str] = Field(default=[])
    recommendations: List[str] = Field(default=[])
    config_patch:    dict      = Field(default={})

class AuditObservation(BaseModel):
    config:           str
    task_description: str
    step:             int
    max_steps:        int
    last_reward:      float
    feedback:         Optional[str]
    task_name:        str
    difficulty:       str

class StepResult(BaseModel):
    observation: AuditObservation
    reward:      float
    done:        bool
    info:        Dict[str, Any] = {}

class EpisodeState(BaseModel):
    episode_id:   str
    step:         int
    task_name:    str
    difficulty:   str
    total_reward: float
    best_reward:  float
    done:         bool

_episode: Dict[str, Any] = {
    "id": None, "task": None, "step": 0,
    "done": False, "rewards": [], "last_reward": 0.0
}

def _build_observation(task, step, reward, feedback):
    return AuditObservation(
        config=task["config"], task_description=task["description"],
        step=step, max_steps=task["max_steps"],
        last_reward=reward, feedback=feedback,
        task_name=task["name"], difficulty=task["difficulty"]
    )

def _feedback_message(reward, task_name):
    if reward == 0.0:   return "No issues identified yet."
    elif reward < 0.35: return f"Score {reward:.2f} — several critical misconfigurations missing."
    elif reward < 0.60: return f"Score {reward:.2f} — good progress. Review encryption and logging."
    elif reward < 0.85: return f"Score {reward:.2f} — almost complete. Check severity labels."
    else:               return f"Score {reward:.2f} — excellent audit!"

@app.post("/reset", response_model=StepResult)
async def reset(task: str = Query(default="easy_security_group")):
    global _episode
    task_name = task if task in TASKS else TASK_SEQUENCE[0]
    task_data = TASKS[task_name]
    _episode  = {
        "id": str(uuid.uuid4()), "task": task_data,
        "step": 0, "done": False, "rewards": [], "last_reward": 0.0
    }
    obs = _build_observation(task_data, 0, 0.0, None)
    return StepResult(observation=obs, reward=0.0, done=False, info={"task": task_name})

@app.post("/step", response_model=StepResult)
async def step(action: AuditAction):
    global _episode
    if not _episode["task"] or _episode["done"]:
        task_data = TASKS[TASK_SEQUENCE[0]]
        obs = _build_observation(task_data, 0, 0.0, "Call /reset first.")
        return StepResult(observation=obs, reward=0.0, done=True, info={"error": "not started"})

    _episode["step"] += 1
    task_data = _episode["task"]
    cur_step  = _episode["step"]

    # ── Grading router ──────────────────────────────────────────────────────
    if task_data["name"] == "easy_security_group":
        reward, breakdown = grade_easy(action.findings, action.severity, action.recommendations, action.config_patch)
    elif task_data["name"] == "medium_s3_policy":
        reward, breakdown = grade_medium(action.findings, action.severity, action.recommendations, action.config_patch)
    elif task_data["name"] == "medium_lambda_iam":
        reward, breakdown = grade_medium_lambda(action.findings, action.severity, action.recommendations, action.config_patch)
    elif task_data["name"] == "hard_rds_cloudtrail":
        reward, breakdown = grade_hard_rds(action.findings, action.severity, action.recommendations, action.config_patch)
    else:
        reward, breakdown = grade_hard(action.findings, action.severity, action.recommendations, action.config_patch)

    _episode["rewards"].append(reward)
    _episode["last_reward"] = reward
    done = (reward >= 0.85) or (cur_step >= task_data["max_steps"])
    _episode["done"] = done
    obs = _build_observation(task_data, cur_step, reward, _feedback_message(reward, task_data["name"]))
    return StepResult(observation=obs, reward=reward, done=done, info={"breakdown": breakdown})

@app.get("/state", response_model=EpisodeState)
async def state():
    rewards = _episode.get("rewards", [])
    return EpisodeState(
        episode_id=_episode.get("id") or "not-started",
        step=_episode.get("step", 0),
        task_name=_episode["task"]["name"] if _episode["task"] else "none",
        difficulty=_episode["task"]["difficulty"] if _episode["task"] else "none",
        total_reward=sum(rewards), best_reward=max(rewards) if rewards else 0.0,
        done=_episode.get("done", False)
    )

@app.get("/health")
async def health():
    return {"status": "healthy", "environment": "aws-security-auditor", "version": "1.0.0"}

@app.get("/tasks")
async def list_tasks():
    return {
        "tasks": [
            {"name": t["name"], "difficulty": t["difficulty"], "max_steps": t["max_steps"]}
            for t in TASKS.values()
        ]
    }

@app.get("/metadata")
async def metadata():
    return {
        "name": "aws-security-auditor",
        "description": "An OpenEnv-compatible RL environment for training AI agents to audit AWS cloud infrastructure configurations.",
        "version": "1.0.0",
        "tasks": list(TASKS.keys())
    }

@app.get("/schema")
async def schema():
    return {
        "action": {
            "type": "object",
            "properties": {
                "findings":        {"type": "array", "items": {"type": "string"}},
                "severity":        {"type": "array", "items": {"type": "string"}},
                "recommendations": {"type": "array", "items": {"type": "string"}},
                "config_patch":    {"type": "object"}
            }
        },
        "observation": {
            "type": "object",
            "properties": {
                "config":           {"type": "string"},
                "task_description": {"type": "string"},
                "step":             {"type": "integer"},
                "max_steps":        {"type": "integer"},
                "last_reward":      {"type": "number"},
                "feedback":         {"type": "string"},
                "task_name":        {"type": "string"},
                "difficulty":       {"type": "string"}
            }
        },
        "state": {
            "type": "object",
            "properties": {
                "episode_id":   {"type": "string"},
                "step":         {"type": "integer"},
                "task_name":    {"type": "string"},
                "difficulty":   {"type": "string"},
                "total_reward": {"type": "number"},
                "best_reward":  {"type": "number"},
                "done":         {"type": "boolean"}
            }
        }
    }

@app.post("/mcp")
async def mcp(request: dict):
    method = request.get("method", "")
    req_id = request.get("id", 1)
    if method == "tools/list":
        result = {"tools": [
            {"name": "reset", "description": "Reset environment and start a new episode"},
            {"name": "step",  "description": "Submit audit findings and get reward"},
            {"name": "state", "description": "Get current episode state"}
        ]}
    elif method == "tools/call":
        result = {"content": [{"type": "text", "text": "Use /reset, /step, /state endpoints directly."}]}
    else:
        result = {"message": "OpenEnv MCP interface ready"}
    return {"jsonrpc": "2.0", "id": req_id, "result": result}