from environment.server import CannonWallEnvironment
import uvicorn
from fastapi import FastAPI
from fastapi.responses import RedirectResponse

# Import Gradio and mount it
import gradio as gr
from ui.demo import demo as gradio_demo

api = FastAPI(title="Cannon & Wall — RedBlue Arena")
env = CannonWallEnvironment()

# ── Environment API routes ──────────────────────────────────────────────────

@api.post("/reset")
def reset(stage: int = 1):
    return env.reset(stage=stage)

@api.post("/step")
def step(action: dict):
    return env.step(action)

@api.get("/state")
def state():
    return env.state

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

# Redirect bare root to /ui/ so judges land on the demo
@api.get("/")
def root():
    return RedirectResponse(url="/ui/")

# ── Mount Gradio at /ui ─────────────────────────────────────────────────────
api = gr.mount_gradio_app(api, gradio_demo, path="/ui")

if __name__ == "__main__":
    uvicorn.run(api, host="0.0.0.0", port=7860)