"""Phase 2 — FastAPI inference service for the HF Docker Space (port 7860).

Loads ALL switchable CLIP models + their indexes once at startup, asserts each
index matches its model, then serves POST /guess (multipart, X-API-Key) and
GET /health (no auth). Pick a model per request with ?model=fast|pro.
"""

import io
import json
import os
import sys
import time
from collections import Counter
from contextlib import asynccontextmanager
from pathlib import Path

import numpy as np
from fastapi import FastAPI, File, Header, HTTPException, Query, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parent / "shared"))
from version import MODELS, DEFAULT_MODEL  # noqa: E402
from embedder import Embedder  # noqa: E402
import guess as guesslib  # noqa: E402
import atlas as atlaslib  # noqa: E402
import region as regionlib  # noqa: E402

DATA = Path(__file__).resolve().parent / "data"
API_KEY = os.environ.get("API_KEY", "")
ALLOWED_ORIGIN = os.environ.get("ALLOWED_ORIGIN", "*")
MAX_BYTES = 4 * 1024 * 1024
RATE_LIMIT_PER_MIN = 30

# Region kNN indexes live in a free HF Dataset repo (the 1GB Space can't hold them all).
# Resolve a region file from local data/ if present, else download from the dataset (cached).
INDEX_DATASET = os.environ.get("INDEX_DATASET", "GeoguessrAngular/geobot-indexes")


def region_file_path(fname):
    local = DATA / fname
    if local.exists():
        return local
    from huggingface_hub import hf_hub_download
    return Path(hf_hub_download(repo_id=INDEX_DATASET, filename=fname, repo_type="dataset"))

Image.MAX_IMAGE_PIXELS = 50_000_000

STATE = {}
_req_times = []  # global in-memory rate limiter timestamps


@asynccontextmanager
async def lifespan(app: FastAPI):
    centroids = json.loads((DATA / "centroids.json").read_text(encoding="utf-8"))
    priors_path = DATA / "priors.json"
    priors = json.loads(priors_path.read_text(encoding="utf-8")) if priors_path.exists() else None

    models = {}
    embedder_cache = {}   # model_id -> Embedder (heavy weights loaded once)
    index_cache = {}      # index_file -> (index, rows, count, model_id, index_version)
    for key, m in MODELS.items():
        if m["model_id"] not in embedder_cache:
            embedder_cache[m["model_id"]] = Embedder(m["model_id"])

        # --- region kNN locator (e.g. Serbia) ---
        if m.get("region_file"):
            try:
                ref_emb, ref_lat, ref_lng = regionlib.load_region(region_file_path(m["region_file"]))
                models[key] = {
                    "type": "region", "embedder": embedder_cache[m["model_id"]],
                    "ref_emb": ref_emb, "ref_lat": ref_lat, "ref_lng": ref_lng,
                    "country_slug": m.get("country_slug", key), "country_name": m.get("country_name", key),
                    "model_id": m["model_id"], "label": m.get("label", key),
                }
                # Difficulty cuts whole LOCATIONS, not frames: ?frac= (0..1) keeps
                # that fraction of the bot's known locations (each pano's frames
                # kept together), so frac=0.05 = the bot knows 5% of locations.
                # All frames of a pano share lat/lng → group rows by coord, permute
                # the locations (seeded), and store each row's location rank. At
                # query time keep rows whose location rank < k. Any % works with no
                # redeploy; the admin panel tunes the fractions server-side.
                n = int(ref_emb.shape[0])
                if n > 60:
                    keys = np.round(np.stack([ref_lat, ref_lng], axis=1), 5)
                    uniq, inv = np.unique(keys, axis=0, return_inverse=True)
                    inv = np.asarray(inv).ravel()
                    nloc = int(uniq.shape[0])
                    loc_order = np.random.default_rng(1234).permutation(nloc)
                    loc_rank = np.empty(nloc, dtype=np.int64)
                    loc_rank[loc_order] = np.arange(nloc)
                    models[key]["row_rank"] = loc_rank[inv].astype(np.int32)
                    models[key]["nloc"] = nloc
                print(f"Loaded model '{key}': region kNN {ref_emb.shape} "
                      f"({models[key].get('nloc', '?')} locs), {m['model_id']}"
                      f"{' +tierable' if 'row_rank' in models[key] else ''}")
            except Exception as e:
                print(f"SKIP model '{key}': failed to load {m['region_file']}: {e}")
            continue

        # --- learned classifier head (Atlas) ---
        if m.get("head_file"):
            W, b, classes = atlaslib.load_head(DATA / m["head_file"])
            models[key] = {
                "type": "head", "embedder": embedder_cache[m["model_id"]],
                "W": W, "b": b, "classes": classes,
                "model_id": m["model_id"], "label": m.get("label", key),
            }
            print(f"Loaded model '{key}': head W{W.shape} {len(classes)} classes, {m['model_id']}")
            continue

        # --- retrieval (index) model ---
        if m["index_file"] not in index_cache:
            meta = json.loads((DATA / m["meta_file"]).read_text(encoding="utf-8"))
            index = np.load(DATA / m["index_file"]).astype(np.float32)
            rows = meta["rows"]
            assert index.shape[0] == len(rows), f"[{key}] index/meta row mismatch"
            index_cache[m["index_file"]] = (
                index, rows, Counter(r["country"] for r in rows),
                meta["model_id"], meta["index_version"])
        index, rows, count, meta_mid, meta_iv = index_cache[m["index_file"]]
        if meta_mid != m["model_id"] or meta_iv != m["index_version"]:
            raise RuntimeError(
                f"[{key}] index/version mismatch: meta has {meta_mid}/{meta_iv}, "
                f"expected {m['model_id']}/{m['index_version']}")

        text_vecs, text_countries = None, None
        if m.get("text_file"):
            text_vecs = np.load(DATA / m["text_file"]).astype(np.float32)
            text_countries = json.loads((DATA / m["text_countries_file"]).read_text(encoding="utf-8"))

        models[key] = {
            "type": "retrieval", "embedder": embedder_cache[m["model_id"]],
            "index": index, "rows": rows, "count": count,
            "model_id": m["model_id"], "index_version": m["index_version"],
            "label": m.get("label", key),
            "text_vecs": text_vecs, "text_countries": text_countries,
        }
        print(f"Loaded model '{key}': index {index.shape}, {m['model_id']} {m['index_version']}"
              f"{', +zeroshot-text' if text_vecs is not None else ''}")

    # Optional script (writing-system) branch for the Atlas head.
    script_vecs = script_names = country_scripts = None
    if (DATA / "script_text.npy").exists() and (DATA / "country_scripts.json").exists():
        script_vecs = np.load(DATA / "script_text.npy").astype(np.float32)
        script_names = json.loads((DATA / "script_names.json").read_text(encoding="utf-8"))
        country_scripts = json.loads((DATA / "country_scripts.json").read_text(encoding="utf-8"))
        print(f"Loaded script branch: {len(script_names)} scripts, {len(country_scripts)} country maps")

    STATE["models"] = models
    STATE["centroids"] = centroids
    STATE["priors"] = priors
    STATE["script_vecs"] = script_vecs
    STATE["script_names"] = script_names
    STATE["country_scripts"] = country_scripts
    STATE["cfg"] = guesslib.Config()
    print(f"Ready. models={list(models)} default={DEFAULT_MODEL} "
          f"centroids={len(centroids)} priors={len(priors) if priors else 0}")
    yield
    STATE.clear()


app = FastAPI(title="GeoBot", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=[ALLOWED_ORIGIN] if ALLOWED_ORIGIN != "*" else ["*"],
                   allow_methods=["*"], allow_headers=["*"])


def _check_key(x_api_key):
    if not API_KEY:
        return  # unset key disables auth (local dev)
    if x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="bad or missing API key")


def _rate_limit():
    now = time.time()
    cutoff = now - 60
    while _req_times and _req_times[0] < cutoff:
        _req_times.pop(0)
    if len(_req_times) >= RATE_LIMIT_PER_MIN:
        raise HTTPException(status_code=429, detail="rate limited")
    _req_times.append(now)


async def _read_image(upload: UploadFile):
    raw = await upload.read()
    if len(raw) > MAX_BYTES:
        raise HTTPException(status_code=413, detail="image too large (>4 MB)")
    try:
        return Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="undecodable image")


@app.get("/health")
def health():
    models = STATE.get("models", {})
    return {
        "status": "ok",
        "default_model": DEFAULT_MODEL,
        "models": {
            k: {"type": v.get("type"), "model_id": v["model_id"], "label": v["label"],
                **({"index_size": int(v["index"].shape[0]), "index_version": v["index_version"]}
                   if v.get("type") == "retrieval" else
                   {"refs": int(v["ref_emb"].shape[0]),
                    "locations": int(v.get("nloc") or 0),
                    "tierable": v.get("row_rank") is not None}
                   if v.get("type") == "region" else
                   {"classes": len(v["classes"])})}
            for k, v in models.items()
        },
    }


@app.post("/guess")
async def do_guess(images: list[UploadFile] = File(None),
                   image0: UploadFile = File(None), image180: UploadFile = File(None),
                   skill: float = Query(3.0, ge=0.0, le=3.0),  # legacy: 1=easy 2=medium 3=full
                   frac: float = Query(None, ge=0.0, le=1.0),   # index density 0..1 (server-computed; wins over skill)
                   model: str = Query(DEFAULT_MODEL),
                   x_api_key: str = Header(None, alias="X-API-Key")):
    _check_key(x_api_key)
    _rate_limit()
    # Accept either N frames under repeated field "images", or legacy image0/image180.
    uploads = [u for u in (images or []) if u is not None]
    if not uploads:
        uploads = [u for u in (image0, image180) if u is not None]
    if not uploads:
        raise HTTPException(status_code=400, detail="no images")
    key = model if model in STATE["models"] else DEFAULT_MODEL
    M = STATE["models"][key]
    t0 = time.time()
    pil = [await _read_image(u) for u in uploads]
    emb = M["embedder"].embed(pil)
    try:
        if M.get("type") == "region":
            # Difficulty = fraction of LOCATIONS the bot knows. Prefer server-sent
            # ?frac=; else map the legacy skill (1->0.05, 2->0.25, 3->full). Keep
            # every frame of the first k locations (k = frac * nloc, >=5 locations).
            re_, rl_, rn_ = M["ref_emb"], M["ref_lat"], M["ref_lng"]
            f = frac if frac is not None else (0.05 if skill <= 1.5 else 0.25 if skill <= 2.5 else 1.0)
            rr, nloc = M.get("row_rank"), M.get("nloc")
            if rr is not None and nloc and f < 0.999:
                k = min(nloc, max(5, int(round(f * nloc))))
                sel = np.where(rr < k)[0]
                re_, rl_, rn_ = re_[sel], rl_[sel], rn_[sel]
            # "human mistake" spread by level: easy often takes another option,
            # medium rarely, hard never (best geo-medoid).
            spread = 0.6 if skill <= 1.5 else 0.3 if skill <= 2.5 else 0.0
            result = regionlib.predict(list(emb), re_, rl_, rn_,
                                       M["country_slug"], M["country_name"], spread=spread)
        elif M.get("type") == "head":
            result = atlaslib.predict(list(emb), M["W"], M["b"], M["classes"],
                                      STATE["centroids"], STATE["cfg"], STATE["priors"],
                                      STATE["script_vecs"], STATE["script_names"],
                                      STATE["country_scripts"])
        else:
            result = guesslib.guess(list(emb), M["index"], M["rows"],
                                    STATE["centroids"], M["count"], STATE["cfg"],
                                    STATE["priors"], M["text_vecs"], M["text_countries"])
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"guess failed: {type(e).__name__}")
    result["timing_ms"] = int((time.time() - t0) * 1000)
    result["model"] = key
    result["model_id"] = M["model_id"]
    result["index_version"] = M.get("index_version")
    print(f"[{key}] guess winner={result['country']} conf={result['confidence']} "
          f"ms={result['timing_ms']}")
    return result