#!/usr/bin/env python3
"""Create / update the CellQuant Space on Hugging Face and wait until it runs.

    python scripts/deploy_hf.py --repo <user>/cellquant [--public] [--no-wait]

Uses the token from `hf auth login` (or HF_TOKEN). Large files (the .pt checkpoint)
are uploaded through LFS automatically. Then polls the Space runtime and, once it is
RUNNING, checks that the Space serves the bundled example image as a real PNG (the Hub
stores binary files as LFS objects, and a checkout only resolves them for extensions
listed in .gitattributes) and calls the API with it to prove the model works.
"""
from __future__ import annotations

import argparse
import os
import sys
import time

from huggingface_hub import HfApi, get_token, whoami

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
IGNORE = ["runs/**", "runs/*", ".pytest_cache/**", "**/__pycache__/**", "__pycache__/**", "*.pyc",
          "reports/final/*.png", "reports/cv/*.png", "*.log", ".gitignore", ".DS_Store"]


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default=None, help="<user>/<space-name>; default <me>/cellquant")
    ap.add_argument("--public", action="store_true")
    ap.add_argument("--no-wait", action="store_true")
    ap.add_argument("--timeout", type=int, default=1800)
    ap.add_argument("--hardware", default="", help="Space hardware to request after the upload, e.g. zero-a10g (ZeroGPU) or cpu-basic")
    ap.add_argument("--bench", action="store_true", help="after the smoke test, time a full validation image with no TTA and with 8x TTA")
    a = ap.parse_args(argv)
    api = HfApi()
    me = whoami()["name"]
    repo = a.repo or f"{me}/cellquant"
    if not os.path.exists(os.path.join(ROOT, "weights", "cellquant.pt")):
        raise SystemExit("weights/cellquant.pt missing: train first")
    n_val = len([f for f in os.listdir(os.path.join(ROOT, "validation")) if f.lower().endswith(".png")]) if os.path.isdir(os.path.join(ROOT, "validation")) else 0
    if n_val == 0:
        print("warning: validation/ holds no images, the Space will have no 'Check against the annotations' tab "
              "(build it with scripts/make_validation_set.py)")
    else:
        print(f"validation bundle: {n_val} annotated images")
    url = api.create_repo(repo, repo_type="space", space_sdk="gradio", private=not a.public, exist_ok=True)
    print("space:", url)
    api.upload_folder(folder_path=ROOT, repo_id=repo, repo_type="space", ignore_patterns=IGNORE,
                      commit_message="Deploy CellQuant")
    print("uploaded")
    if a.hardware:
        rt = api.get_space_runtime(repo)
        if str(rt.hardware) != a.hardware or (rt.requested_hardware and str(rt.requested_hardware) != a.hardware):
            api.request_space_hardware(repo, a.hardware)
            print(f"requested hardware {a.hardware} (was {rt.hardware})")
        else:
            print(f"hardware already {a.hardware}")
    if a.no_wait:
        return
    t0 = time.time()
    last = None
    while time.time() - t0 < a.timeout:
        rt = api.get_space_runtime(repo)
        stage = str(rt.stage)
        if stage != last:
            print(f"[{time.time()-t0:5.0f}s] stage {stage} hardware {rt.hardware}", flush=True)
            last = stage
        if stage.endswith("RUNNING"):
            break
        if "ERROR" in stage or stage.endswith("STOPPED"):
            print("build failed; last log lines:")
            try:
                print(api.get_space_runtime(repo))
            except Exception as e:  # noqa: BLE001
                print(e)
            sys.exit(1)
        time.sleep(15)
    else:
        raise SystemExit("timed out waiting for the Space")
    # the example image as the Space itself serves it must be a PNG, not an LFS pointer
    import requests
    host = api.space_info(repo).host or f"https://{repo.replace('/', '-')}.hf.space"
    served = f"{host.rstrip('/')}/gradio_api/file=/home/user/app/examples/F100_crop_512.png"
    r = requests.get(served, timeout=60)
    if r.status_code != 200 or not r.content.startswith(b"\x89PNG"):
        raise SystemExit(f"Space serves {served} as {r.status_code} / {len(r.content)} bytes starting {r.content[:40]!r}: "
                         "not a PNG (unresolved LFS pointer? check .gitattributes)")
    print(f"served example ok: {len(r.content)} bytes")
    # smoke test through the Gradio API with that served copy
    from gradio_client import Client, handle_file
    client = Client(repo, hf_token=get_token())
    t = time.time()
    out = client.predict([handle_file(served)], 0.5, 0.5, 50, "head", "contains", api_name="/analyze_files")
    print(f"API call ok in {time.time()-t:.1f}s; counts row: {out[0]['data'][0] if isinstance(out[0], dict) else out[0]}")
    lines = [l for l in str(out[3]).splitlines() if l.strip()]
    print(f"  app reports: {lines[0] if lines else '(no summary line)'} | summary: {lines[-1][:90] if len(lines) > 1 else '(none)'}")
    if a.bench:
        val_dir = os.path.join(ROOT, "validation")
        pngs = sorted(f for f in os.listdir(val_dir) if f.lower().endswith(".png")) if os.path.isdir(val_dir) else []
        if not pngs:
            print("bench: no validation image to time")
            return
        full = os.path.join(val_dir, pngs[0])
        t = time.time()
        out = client.predict([handle_file(full)], 0.5, 0.5, 50, "head", "contains", api_name="/analyze_files")
        row = out[0]["data"][0] if isinstance(out[0], dict) else out[0]
        lines = [l for l in str(out[3]).splitlines() if l.strip()]
        print(f"bench section 1, {os.path.basename(full)}: round trip {time.time()-t:.1f}s; {lines[0] if lines else ''}; nuclei {sum(row[1:5])}")
        sys.path.insert(0, ROOT)
        from cellquant.validation import ValidationSet
        key = ValidationSet(val_dir).keys[0]
        t = time.time()
        out = client.predict(key, api_name="/validate_image")
        lines = [l for l in str(out[0]).splitlines() if l.strip()]
        print(f"bench section 2, {key}: round trip {time.time()-t:.1f}s; {lines[0] if lines else ''}")
        print(f"  {[l for l in lines if 'Run time' in l] or lines[-1:]}")


if __name__ == "__main__":
    main()