from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.responses import FileResponse
from pydantic import BaseModel

from translator import load_translator

# Point this at YOUR weights repo on the Hugging Face Hub.
WEIGHTS_REPO = "ME22B191/neural_translation_weights"
CHECKPOINT_FILE = "checkpoint_baseline.pt"  # <-- the exact filename you uploaded

ml = {}  # holds the one loaded model for the whole process lifetime


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Runs once at startup: download weights + vocab, assemble the model.
    ml["model"] = load_translator(WEIGHTS_REPO, checkpoint_file=CHECKPOINT_FILE)
    yield
    ml.clear()  # runs once at shutdown


app = FastAPI(title="Neural Translator", lifespan=lifespan)


class TranslationRequest(BaseModel):
    text: str


class TranslationResponse(BaseModel):
    translation: str


@app.get("/")
def home():
    return FileResponse("index.html")


@app.get("/health")
def health_check():
    return {"status": "alive", "model_loaded": "model" in ml}


@app.post("/translate", response_model=TranslationResponse)
def translate(request: TranslationRequest):      # plain def -> runs in threadpool
    # infer() already does .eval() + torch.no_grad() internally.
    english = ml["model"].infer(request.text)
    return TranslationResponse(translation=english)