from __future__ import annotations

from typing import Optional

import uvicorn
from fastapi import Body, FastAPI
from pydantic import BaseModel

from env.environment import CodeReviewOpenEnv
from env.models import Action, Observation, TaskName


app = FastAPI(
    title="Autonomous Code Review & Merge Decision Environment",
    version="1.0.0",
    description="A deterministic OpenEnv-compatible environment for AI-driven pull request review workflows.",
)
env = CodeReviewOpenEnv()


class ResetRequest(BaseModel):
    task: Optional[TaskName] = None


class ResetResponse(BaseModel):
    observation: Observation


class StepResponse(BaseModel):
    observation: Observation
    reward: float
    done: bool
    info: dict


@app.post("/reset", response_model=ResetResponse)
def reset_environment(request: Optional[ResetRequest] = Body(default=None)) -> ResetResponse:
    observation = env.reset(task=request.task if request else None)
    return ResetResponse(observation=observation)


@app.post("/step", response_model=StepResponse)
def step_environment(action: Action) -> StepResponse:
    observation, reward, done, info = env.step(action)
    return StepResponse(observation=observation, reward=reward, done=done, info=info)


@app.get("/state")
def get_state() -> dict:
    return env.state()


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


if __name__ == "__main__":
    main()