app.py

import os
import logging
from contextlib import asynccontextmanager

import torch
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from maia3_pkg.inference import build_model, predict_moves

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

model = None
model_cfg = None


class PredictRequest(BaseModel):
    fen: str = Field(..., description="FEN string of the chess position")
    self_elo: int = Field(1500, ge=0, le=5000, description="Side-to-move Elo rating")
    oppo_elo: int = Field(1500, ge=0, le=5000, description="Opponent Elo rating")
    temperature: float = Field(1.0, ge=0.0, le=5.0, description="Sampling temperature (0 = argmax)")
    top_p: float = Field(1.0, ge=0.0, le=1.0, description="Nucleus sampling threshold")
    multipv: int = Field(5, ge=1, le=50, description="Number of candidate moves to return")


class PredictResponse(BaseModel):
    fen: str
    turn: str
    self_elo: int
    oppo_elo: int
    top_moves: list[dict]
    wdl: dict


class HealthResponse(BaseModel):
    status: str
    device: str
    model_loaded: bool


@asynccontextmanager
async def lifespan(app: FastAPI):
    global model, model_cfg
    logger.info("Starting Maia3 inference server...")
    try:
        ckpt_path = os.environ.get(
            "MAIA3_CHECKPOINT",
            "/app/maia3-3m.pt",
        )
        device = os.environ.get("MAIA3_DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
        use_amp = os.environ.get("MAIA3_USE_AMP", "1" if device.startswith("cuda") else "0") == "1"
        logger.info(f"Loading model from {ckpt_path} on {device} (amp={use_amp})...")
        model, model_cfg = build_model(ckpt_path, device=device, use_amp=use_amp)
        logger.info("Model loaded successfully")
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        model = None
        model_cfg = None
    yield


app = FastAPI(
    title="Maia3 Chess Prediction API",
    description="Human-like chess move prediction using Maia3-3M (Chessformer)",
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/health", response_model=HealthResponse)
async def health():
    return HealthResponse(
        status="ok" if model is not None else "degraded",
        device=str(model_cfg.device) if model_cfg else "unknown",
        model_loaded=model is not None,
    )


@app.post("/predict", response_model=PredictResponse)
async def predict(req: PredictRequest):
    if model is None or model_cfg is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        result = predict_moves(
            model=model,
            cfg=model_cfg,
            fen=req.fen,
            self_elo=req.self_elo,
            oppo_elo=req.oppo_elo,
            temperature=req.temperature,
            top_p=req.top_p,
            multipv=req.multipv,
        )
        return PredictResponse(**result)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.exception("Prediction failed")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/")
async def root():
    return {
        "service": "Maia3 Chess Prediction",
        "model": "Maia3-ablate-3M (Chessformer, 3M params)",
        "endpoints": {
            "health": "GET /health",
            "predict": "POST /predict",
            "docs": "GET /docs",
        },
    }