import uvicorn
from fastapi import Request
from fastapi.responses import JSONResponse
from pydantic import ValidationError

from app.main import app as _app  # the existing app
from env.environment import env_instance
from env.models import Action

app = _app

@app.post("/reset")
async def env_reset(request: Request):
    data = {}
    try:
        data = await request.json()
    except Exception:
        pass
    obs = env_instance.reset(data)
    return JSONResponse(content={"observation": obs.model_dump()})

@app.post("/step")
async def env_step(request: Request):
    try:
        data = await request.json()
        action_data = data.get("action", {})
        action_obj = Action(**action_data)
        
        obs, reward, done, info = env_instance.step(action_obj.model_dump())
        
        return JSONResponse(content={
            "observation": obs.model_dump(),
            "reward": reward.model_dump(),
            "done": done,
            "info": info
        })
    except ValidationError as e:
        return JSONResponse(status_code=400, content={"error": "Invalid action format", "details": e.errors()})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

@app.get("/state")
def env_state():
    return JSONResponse(content=env_instance.state())

def main():
    uvicorn.run("server:app", host="0.0.0.0", port=8000)

if __name__ == "__main__":
    main()