"""
HF Space (Arko006 account) — the ONLY component that loads the ESM-2 model.
Downloads the fine-tuned checkpoint from the Arko007 HF *model* repo at
container startup (never bundled into this Space's own git history) and
exposes a small, shared-secret-gated scoring API for the Render orchestrator.
"""

import os
import threading
from contextlib import asynccontextmanager
from typing import Optional

import torch
from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel
from huggingface_hub import hf_hub_download

from model import TransformerDMSRegressor, DEFAULT_MODEL_NAME

CHECKPOINT_REPO = os.getenv("CHECKPOINT_REPO", "Arko007/esm2-cancer-nlr-35M")
CHECKPOINT_FILE = os.getenv("CHECKPOINT_FILE", "best.pt")
SHARED_SECRET = os.getenv("MODEL_SERVING_SHARED_SECRET", "")

_state = {"model": None, "loading": True, "revision": None, "error": None}
_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def _load_model():
    try:
        model = TransformerDMSRegressor(DEFAULT_MODEL_NAME)
        checkpoint_path = hf_hub_download(repo_id=CHECKPOINT_REPO, filename=CHECKPOINT_FILE)
        state_dict = torch.load(checkpoint_path, map_location=_device)
        model.load_state_dict(state_dict["model_state_dict"])
        model.to(_device)
        model.eval()
        _state["model"] = model
        _state["revision"] = f"{CHECKPOINT_REPO}/{CHECKPOINT_FILE} (epoch {state_dict.get('epoch', '?')})"
    except Exception as e:
        # No fine-tuned checkpoint yet is a legitimate startup state (e.g.
        # before the first Kaggle training run has completed) — fall back to
        # the base pretrained ESM-2 model so /score still returns real
        # zero-shot LLR scores (nlr_fitness will be None until fine-tuned).
        print(f"Could not load fine-tuned checkpoint ({e}); falling back to base ESM-2 zero-shot scoring.")
        try:
            model = TransformerDMSRegressor(DEFAULT_MODEL_NAME)
            model.to(_device)
            model.eval()
            _state["model"] = model
            _state["revision"] = f"{DEFAULT_MODEL_NAME} (base, not fine-tuned)"
        except Exception as inner_e:
            _state["error"] = str(inner_e)
    finally:
        _state["loading"] = False


@asynccontextmanager
async def lifespan(app: FastAPI):
    threading.Thread(target=_load_model, daemon=True).start()
    yield


app = FastAPI(title="Cancer Mutation Predictor — ESM-2 model serving", lifespan=lifespan)


class ScoreRequest(BaseModel):
    sequence: str
    position: int
    ref_aa: str
    alt_aa: str


def _check_secret(x_internal_key: Optional[str]):
    if SHARED_SECRET and x_internal_key != SHARED_SECRET:
        raise HTTPException(status_code=401, detail="Invalid or missing X-Internal-Key")


@app.get("/health")
def health():
    return {
        "status": "ok" if not _state["loading"] else "warming",
        "model_loaded": _state["model"] is not None,
        "checkpoint": _state["revision"],
        "error": _state["error"],
    }


@app.post("/score")
def score(req: ScoreRequest, x_internal_key: Optional[str] = Header(None)):
    _check_secret(x_internal_key)
    if _state["model"] is None:
        raise HTTPException(status_code=503, detail="Model still warming up, try again shortly")

    if not (1 <= req.position <= len(req.sequence)):
        raise HTTPException(status_code=400, detail="position out of range for the given sequence")

    import time
    t0 = time.time()
    llr, fitness = _state["model"].score(req.sequence, req.position, req.ref_aa, req.alt_aa, _device)
    return {
        "raw_llr": llr,
        "nlr_fitness": fitness if "not fine-tuned" not in (_state["revision"] or "") else None,
        "model_id": DEFAULT_MODEL_NAME,
        "checkpoint_revision": _state["revision"],
        "inference_ms": round((time.time() - t0) * 1000, 1),
    }


@app.post("/score/batch")
def score_batch(reqs: list[ScoreRequest], x_internal_key: Optional[str] = Header(None)):
    _check_secret(x_internal_key)
    if len(reqs) > 50:
        raise HTTPException(status_code=400, detail="Batch limited to 50 variants per call")
    return [score(r, x_internal_key) for r in reqs]