"""
HF Spaces server — DebugOps Environment (OpenEnv compatible)

Exposes:
  POST /reset
  POST /step
  GET  /state
"""

from __future__ import annotations
from typing import Dict, Any, Optional

import os
import uvicorn
from fastapi import FastAPI, HTTPException, Body

from tasks.task_simple import create_env as create_simple
from tasks.task_multi_service import create_env as create_multi
from tasks.task_critical import create_env as create_critical


app = FastAPI(title="DebugOps AI Environment", version="1.0.0")

# Global environment instance
_env: Optional[Any] = None

def create_env(task: str):
    if task == "simple":
        return create_simple()
    elif task == "multi_service":
        return create_multi()
    elif task == "critical":
        return create_critical()
    else:
        raise ValueError(f"Invalid task: {task}")


def get_env():
    global _env
    if _env is None:
        raise HTTPException(status_code=400, detail="Call /reset first")
    return _env

@app.get("/")
def root():
    return {
        "name": "DebugOps AI Environment",
        "description": "Production debugging RL environment (OpenEnv compatible)",
        "endpoints": ["/reset", "/step", "/state", "/health"],
    }


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


@app.post("/reset")
def reset(payload: Optional[Dict[str, Any]] = Body(default={})):
    """
    Accepts BOTH:
      {}                         (validator case)
      {"task": "critical"}       (manual case)
    """
    global _env

    try:
        task = payload.get("task", "simple") if payload else "simple"

        _env = create_env(task)
        obs = _env.reset()

        return {
            "observation": obs,
            "done": False,
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/step")
@app.post("/step/")
def step(payload: Dict[str, Any] = Body(...)):
    env = get_env()

    try:
        action = payload.get("action")
        if not action:
            raise ValueError("Missing 'action' field")

        obs, reward, done, info = env.step(action)

        return {
            "observation": obs,
            "reward": float(round(reward, 3)),
            "done": done,
            "info": info,
        }

    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/state")
@app.get("/state/")
def state():
    env = get_env()
    try:
        return env.state()
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

def main():
    port = int(os.getenv("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)


if __name__ == "__main__":
    main()