"""FastAPI server exposing POST /reset, POST /step, GET /state endpoints."""

from __future__ import annotations

from typing import Any, Dict, Optional

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

from env.environment import CustomerSupportEnv
from env.models import Action, ActionType

# ── App ───────────────────────────────────────────────────────────────

app = FastAPI(
    title="OpenEnv Customer Support",
    description="AI Customer Support Ticket Resolution Environment",
    version="1.0.0",
)

# Single environment instance
_env = CustomerSupportEnv()


# ── Request / Response schemas ────────────────────────────────────────

class ResetRequest(BaseModel):
    task_id: str = "easy_classification"
    ticket_index: int = 0


class StepRequest(BaseModel):
    action_type: str
    classification: Optional[str] = None
    response_text: Optional[str] = None
    refund_amount: Optional[float] = None
    escalation_reason: Optional[str] = None


class ResetResponse(BaseModel):
    observation: Dict[str, Any]


class StepResponse(BaseModel):
    observation: Dict[str, Any]
    reward: Dict[str, Any]
    done: bool
    info: Dict[str, Any]


class StateResponse(BaseModel):
    state: Dict[str, Any]


# ── Endpoints ─────────────────────────────────────────────────────────

@app.post("/reset", response_model=ResetResponse)
def reset_endpoint(req: Optional[ResetRequest] = None) -> ResetResponse:
    """Reset the environment and return the initial observation."""
    if req is None:
        req = ResetRequest()
    try:
        obs = _env.reset(task_id=req.task_id, ticket_index=req.ticket_index)
        return ResetResponse(observation=obs.model_dump())
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc))


@app.post("/step", response_model=StepResponse)
def step_endpoint(req: StepRequest) -> StepResponse:
    """Execute one environment step."""
    try:
        action_type = ActionType(req.action_type)
    except ValueError:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Invalid action_type '{req.action_type}'. "
                f"Valid: {[a.value for a in ActionType]}"
            ),
        )

    action = Action(
        action_type=action_type,
        classification=req.classification,
        response_text=req.response_text,
        refund_amount=req.refund_amount,
        escalation_reason=req.escalation_reason,
    )

    try:
        result = _env.step(action)
    except RuntimeError as exc:
        raise HTTPException(status_code=400, detail=str(exc))

    return StepResponse(
        observation=result.observation.model_dump(),
        reward=result.reward.model_dump(),
        done=result.done,
        info=result.info,
    )


@app.get("/state", response_model=StateResponse)
def state_endpoint() -> StateResponse:
    """Return the current environment state."""
    return StateResponse(state=_env.state())


# ── Health check ──────────────────────────────────────────────────────

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


# ── Entry point ───────────────────────────────────────────────────────

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

if __name__ == "__main__":
    main()