app.py

from fastapi import FastAPI

from agent_contract import (
    ALLOWED_CATEGORIES,
    ALLOWED_SEVERITIES,
    BASE_AGENT_INSTRUCTIONS,
    RESPONSE_EXAMPLE,
)
from environment import SecurityEnv
from grader import grade_response

app = FastAPI(
    docs_url="/",
    redoc_url=None,
    openapi_url="/openapi.json"
)

from fastapi.middleware.cors import CORSMiddleware

env = None

def get_env():
    global env
    if env is None:
        env = SecurityEnv()
    return env
    
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.api_route("/reset", methods=["GET", "POST"])
def reset():
    return get_env().reset()

@app.post("/step")
def step(action: dict):
    # ✅ Safe extraction (prevents crashes)
    safe_action = {
        "category": str(action.get("category", "")).strip().lower(),
        "severity": str(action.get("severity", "")).strip().lower(),
        "action": str(action.get("action", "")).strip()
    }

    return get_env().step(safe_action)


@app.get("/state")
def state():
    return get_env().state()


@app.get("/tasks")
def tasks():
    return {
        "tasks": [
            {
                "name": "easy",
                "description": "Detect normal vs attack",
            },
            {
                "name": "medium",
                "description": "Classify category",
            },
            {
                "name": "hard",
                "description": "Category + severity + action",
            },
        ],
        "output_contract": {
            "instructions": BASE_AGENT_INSTRUCTIONS,
            "allowed_categories": ALLOWED_CATEGORIES,
            "allowed_severities": ALLOWED_SEVERITIES,
            "response_example": RESPONSE_EXAMPLE,
        },
    }


@app.post("/grader")
def grader(data: dict):
    predicted = data["predicted"]
    expected = data["expected"]
    score = grade_response(predicted, expected)

    return {"score": score}


@app.get("/baseline")
def baseline():
    sample = get_env().reset()
    action = {
        "category": "normal",
        "severity": "low",
        "action": "monitor",
    }
    result = get_env().step(action)

    return {
        "observation": sample,
        "result": result,
    }