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

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

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

import os
import logging
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel



# Import environments
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

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")

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

# Global env instance
_env = None


class ResetRequest(BaseModel):
    """Request schema for resetting the environment."""
    task: str = "simple"   # simple | multi_service | critical


class StepRequest(BaseModel):
    """Request schema for stepping through the environment."""
    action: str


def create_env(task: str):
    """Factory method to create environment based on task type."""
    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():
    """Retrieve the current environment instance."""
    global _env
    if _env is None:
        raise HTTPException(status_code=400, detail="Call /reset first")
    return _env


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


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


@app.post("/reset")
def reset(request: Optional[ResetRequest] = None):
    """Reset the environment with a given task."""
    global _env
    try:
        task = request.task if request else "simple"
        logging.info(f"Resetting environment with task: {task}")

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

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

    except Exception as e:
        logging.error(f"Error during reset: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/step")
def step(request: StepRequest):
    """Perform a step in the environment with the given action."""
    env = get_env()

    try:
        obs, reward, done, info = env.step(request.action)
        logging.info(f"Step taken: {request.action}, reward={reward}, done={done}")

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

    except Exception as e:
        logging.error(f"Error during step: {e}")
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/state")
def state():
    """Retrieve the current environment state."""
    env = get_env()
    try:
        return env.state()
    except Exception as e:
        logging.error(f"Error retrieving state: {e}")
        raise HTTPException(status_code=400, detail=str(e))





def main():
    """Entry point for running the FastAPI server."""
    port = int(os.getenv("PORT", 7860))
    logging.info(f"Starting DebugOps server on port {port}")
    uvicorn.run(app, host="0.0.0.0", port=port)


if __name__ == "__main__":
    main()