#!/usr/bin/env python3
"""
End-to-end API checks for the HF/Docker image (same stack as Spaces).

Usage:
  python scripts/verify_hf_stack.py [--base-url http://127.0.0.1:7860]
  HF_VERIFY_BASE_URL=https://your-space.hf.space python scripts/verify_hf_stack.py

Or run the full Docker probe (build, run container, hit API):
  python scripts/verify_hf_stack.py --docker

Deploy to the Hub (requires HF_TOKEN + HF_SPACE_REPO), then verify live:
  HF_TOKEN=... HF_SPACE_REPO=user/slug python scripts/publish_hf_space.py --verify

Prove catalog + kinks + recommendations on a deployed Space (no partner flow):
  HF_VERIFY_BASE_URL=https://owner-slug.hf.space HF_VERIFY_MIN_KINKS=1000 \\
    python scripts/verify_hf_stack.py --smoke

Auth login contract (unknown user 404 vs wrong password 401, whitespace body) is asserted in
``--smoke`` and full checks. Deploy + prove in one step (needs write token + repo resolution):
  HF_TOKEN=... HF_SPACE_REPO=owner/slug python scripts/publish_hf_space.py --verify

Smoke/full checks also **fetch one catalog ``asset_url``** (HTTPS image) when ``stats.assets >= 1``,
unless ``HF_VERIFY_SKIP_PICTURE_PROBE=1``. Tunable: ``HF_VERIFY_IMAGE_TIMEOUT_S`` (seconds).
"""
from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
import time
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


def _http_timeout_s() -> float:
    return float(os.environ.get("HF_VERIFY_HTTP_TIMEOUT_S", "120") or "120")


def _get(url: str, headers: dict[str, str] | None = None) -> dict[str, Any]:
    req = Request(url, headers=headers or {})
    with urlopen(req, timeout=_http_timeout_s()) as resp:
        return json.loads(resp.read().decode())


def _get_text(url: str, headers: dict[str, str] | None = None) -> str:
    req = Request(url, headers=headers or {})
    with urlopen(req, timeout=_http_timeout_s()) as resp:
        return resp.read().decode()


def _assert_live_frontend_recs_error_shaping_once(root: str) -> tuple[str, str]:
    """Single attempt: follow index to the built Vite JS bundle and assert error shaping is deployed."""
    base = root.rstrip("/")
    html = _get_text(f"{base}/")
    for stale_origin in ("https://esm.sh", "https://unpkg.com"):
        if stale_origin in html:
            raise RuntimeError(f"GET /: found stale CDN frontend origin {stale_origin!r}")
    m = re.search(r"""src=["'](/frontend/assets/[^"']+\.js)["']""", html)
    if not m:
        raise RuntimeError("GET /: expected a built Vite script under /frontend/assets/")
    bundle_path = m.group(1)
    bundle_js = _get_text(f"{base}{bundle_path}")
    for label, needle in {
        "HTML API error shaping": "the response was a web page, not JSON",
        "generic API failure shaping": "Request failed (",
        "recommendations retry control": "discover-recs-retry",
    }.items():
        if needle not in bundle_js:
            raise RuntimeError(
                f"GET {bundle_path}: missing {label} marker {needle!r} — stale UI bundle"
            )
    return bundle_path, bundle_path


def _assert_live_frontend_recs_error_shaping(root: str) -> None:
    """Follow index → Vite bundle and assert error-UI hardening is deployed."""
    if os.environ.get("HF_VERIFY_SKIP_FRONTEND", "").strip().lower() in ("1", "true", "yes", "on"):
        print("verify_hf_stack: skipping live frontend bundle check (HF_VERIFY_SKIP_FRONTEND=1)", file=sys.stderr)
        return
    r = root.rstrip("/")
    if "hf.space" in r:
        retries = max(1, int(os.environ.get("HF_VERIFY_FRONTEND_RETRIES", "120") or "120"))
        interval = float(os.environ.get("HF_VERIFY_FRONTEND_INTERVAL_S", "15") or "15")
    else:
        retries, interval = 1, 0.0

    last_err: BaseException | None = None
    for attempt in range(retries):
        try:
            api_path, disc_path = _assert_live_frontend_recs_error_shaping_once(root)
            if attempt:
                print(f"verify_hf_stack: frontend bundle ok after {attempt} wait(s)", file=sys.stderr)
            print(
                f"verify_hf_stack: frontend recs-error bundle ok ({api_path}, {disc_path})",
                file=sys.stderr,
            )
            return
        except (RuntimeError, URLError, HTTPError, OSError) as e:
            last_err = e
            if retries > 1 and attempt + 1 < retries:
                if attempt == 0 or (attempt + 1) % 6 == 0:
                    print(
                        f"verify_hf_stack: frontend bundle not ready ({e}); "
                        f"retry {attempt + 1}/{retries} ({interval:.0f}s)…",
                        file=sys.stderr,
                    )
                time.sleep(interval)
    if last_err is None:
        raise RuntimeError("frontend bundle check failed (no error recorded)")
    raise RuntimeError(f"frontend bundle check failed after {retries} attempt(s): {last_err}") from last_err


def _peek_coep(base: str) -> str:
    url = f"{base.rstrip('/')}/health"
    req = Request(url)
    with urlopen(req, timeout=_http_timeout_s()) as resp:
        return (resp.headers.get("Cross-Origin-Embedder-Policy") or "").strip()


def _assert_coep_credentialless(base: str) -> None:
    """COEP must allow cross-origin catalog images (see api middleware + tests/test_coep_headers.py).

    On ``*.hf.space``, retry for a while: after ``publish_hf_space`` pushes, the edge may still serve the
    previous image (``require-corp``) until the new container and README ``custom_headers`` propagate.
    """
    root = base.rstrip("/")
    if os.environ.get("HF_VERIFY_SKIP_COEP", "").strip().lower() in ("1", "true", "yes", "on"):
        print("verify_hf_stack: skipping COEP check (HF_VERIFY_SKIP_COEP=1)", file=sys.stderr)
        return
    if "hf.space" not in root:
        coep = _peek_coep(root)
        assert coep == "credentialless", (
            f"Cross-Origin-Embedder-Policy is {coep!r}; expected credentialless for catalog <img> hosts without CORP"
        )
        return

    retries = int(os.environ.get("HF_VERIFY_COEP_RETRIES", "90") or "90")
    interval = float(os.environ.get("HF_VERIFY_COEP_INTERVAL_S", "10") or "10")
    last = ""
    for attempt in range(retries):
        last = _peek_coep(root)
        if last == "credentialless":
            if attempt:
                print(f"verify_hf_stack: COEP credentialless after {attempt} wait(s)", file=sys.stderr)
            return
        if attempt == 0 or (attempt + 1) % 6 == 0:
            print(
                f"verify_hf_stack: COEP is {last!r} (want credentialless); "
                f"waiting for deploy ({attempt + 1}/{retries}, {interval:.0f}s)…",
                file=sys.stderr,
            )
        time.sleep(interval)
    assert last == "credentialless", (
        f"Cross-Origin-Embedder-Policy stayed {last!r} after {retries * interval:.0f}s; "
        "expected credentialless (see README custom_headers + api.py). "
        "Set HF_VERIFY_SKIP_COEP=1 only to bypass."
    )


def _post(url: str, body: dict[str, Any], headers: dict[str, str] | None = None) -> dict[str, Any]:
    data = json.dumps(body).encode()
    h = {"Content-Type": "application/json", **(headers or {})}
    req = Request(url, data=data, headers=h, method="POST")
    with urlopen(req, timeout=_http_timeout_s()) as resp:
        return json.loads(resp.read().decode())


def _post_status(
    url: str, body: dict[str, Any], headers: dict[str, str] | None = None
) -> tuple[int, dict[str, Any]]:
    """POST and return (status_code, json_body) without raising on 4xx/5xx."""
    data = json.dumps(body).encode()
    h = {"Content-Type": "application/json", **(headers or {})}
    req = Request(url, data=data, headers=h, method="POST")
    try:
        with urlopen(req, timeout=_http_timeout_s()) as resp:
            raw = resp.read().decode()
            code = int(getattr(resp, "status", None) or resp.getcode())
            return code, (json.loads(raw) if raw.strip() else {})
    except HTTPError as e:
        raw = e.read().decode()
        try:
            payload: dict[str, Any] = json.loads(raw) if raw.strip() else {}
        except json.JSONDecodeError:
            payload = {}
        return int(e.code), payload


def _wait_catalog_sample_ready(root: str) -> dict[str, Any]:
    """After cold start or redeploy, ``/health/catalog-sample`` may return ``{ok: false, starting: true}`` until backend binds."""
    url = f"{root.rstrip('/')}/health/catalog-sample"
    timeout_s = float(os.environ.get("HF_VERIFY_CATALOG_SAMPLE_TIMEOUT_S", "3600") or "3600")
    deadline = time.time() + timeout_s
    last: dict[str, Any] = {}
    while time.time() < deadline:
        last = _get(url)
        if last.get("ok") is True:
            return last
        if last.get("starting") is True:
            time.sleep(3.0)
            continue
        raise RuntimeError(f"catalog-sample not ok: {last!r}")
    raise TimeoutError(f"/health/catalog-sample did not become ok within {timeout_s:.0f}s (last={last!r})")


def _wait_for_health(base: str, timeout_s: float) -> None:
    """Hugging Face Spaces may cold-start; retry /health before failing."""
    url = f"{base.rstrip('/')}/health"
    deadline = time.time() + timeout_s
    last_err: Exception | None = None
    while time.time() < deadline:
        try:
            h = _get(url)
            if h.get("ok") is True:
                return
        except (URLError, HTTPError, TimeoutError, OSError) as e:
            last_err = e
        time.sleep(3.0)
    raise last_err or TimeoutError(f"/health did not return ok within {timeout_s:.0f}s: {url}")


def _require(cond: bool, msg: str) -> None:
    if not cond:
        raise RuntimeError(msg)


def _assert_auth_login_http_contract(root: str, u: dict[str, Any]) -> None:
    """POST /auth/login: unknown profile -> 404 (ephemeral DB UX), wrong token -> 401, strip whitespace -> 200."""
    r = root.rstrip("/")
    st, body = _post_status(
        f"{r}/auth/login",
        {"user_id": "no-verify-hf-profile-00", "private_token": "x"},
    )
    _require(st == 404, f"unknown user login expected 404, got {st}: {body!r}")
    detail = str(body.get("detail", ""))
    _require("Unknown profile" in detail, f"404 detail should mention Unknown profile: {body!r}")

    st, body = _post_status(
        f"{r}/auth/login",
        {"user_id": u["id"], "private_token": "not-the-token-for-verify-hf-stack"},
    )
    _require(st == 401, f"wrong token expected 401, got {st}: {body!r}")
    _require(body.get("detail") == "Invalid credentials", f"401 detail: {body!r}")

    st, body = _post_status(
        f"{r}/auth/login",
        {"user_id": f"  {u['id']} \n", "private_token": f" {u['private_token']} "},
    )
    _require(st == 200, f"padded login expected 200, got {st}: {body!r}")
    _require(body.get("user", {}).get("id") == u["id"], f"login user id mismatch: {body!r}")
    print("verify_hf_stack: /auth/login HTTP contract ok", file=sys.stderr)


def _assert_catalog_https_image_fetchable(root: str) -> None:
    """Prove same-origin ``/media/...`` images are served (bundled seed bytes + optional catalog hit)."""
    if os.environ.get("HF_VERIFY_SKIP_PICTURE_PROBE", "").strip().lower() in ("1", "true", "yes", "on"):
        print("verify_hf_stack: skipping catalog image probe (HF_VERIFY_SKIP_PICTURE_PROBE=1)", file=sys.stderr)
        return
    from backend.catalog_image_probe import (
        assert_catalog_image_url_works,
        first_catalog_image_probe_url,
        first_catalog_image_probe_url_from_list,
    )

    base = root.rstrip("/")
    t_img = float(os.environ.get("HF_VERIFY_IMAGE_TIMEOUT_S", "45") or "45")

    # Bundled demo tiles (always shipped under deploy/hf/seed/cached_assets/hf_seed; see deploy/hf/entrypoint.sh).
    demo = f"{base}/media/hf_seed/demo_kink_1.jpg"
    assert_catalog_image_url_works(demo, timeout_s=t_img)
    print("verify_hf_stack: bundled hf_seed demo image ok", file=sys.stderr)

    # Optional placeholder for a common FetLife ``/media/...`` path (shipped under deploy/hf/seed on newer builds).
    fl_placeholder = f"{base}/media/fetlife_fetishes/1/146853406.jpg"
    try:
        assert_catalog_image_url_works(fl_placeholder, timeout_s=t_img)
    except HTTPError as e:
        if int(e.code) != 404:
            raise
        print(
            "verify_hf_stack: fetlife_fetishes seed placeholder not on edge yet (404); "
            "hf_seed check above still proves /media is live.",
            file=sys.stderr,
        )
    else:
        print("verify_hf_stack: bundled fetlife_fetishes placeholder image ok", file=sys.stderr)

    health = _get(f"{base}/health")
    stats = health.get("stats") or {}
    media = health.get("media") or {}
    n_assets = int(stats.get("assets") or 0)
    if n_assets < 1:
        print("verify_hf_stack: skipping live catalog row probe (stats.assets is 0)", file=sys.stderr)
        return
    require_catalog_image = (
        os.environ.get("HF_VERIFY_REQUIRE_CATALOG_IMAGE", "").strip().lower() in ("1", "true", "yes", "on")
        or (bool(media.get("strict_missing_assets")) and int(stats.get("kinks") or 0) >= 1000)
    )

    lst = _get(f"{base}/kinks?limit=500")
    items = lst.get("items") or []
    kid, url = first_catalog_image_probe_url_from_list(base, items)
    if not url:
        probe = _get(f"{base}/health/catalog-sample")
        sid = probe.get("sample_kink_id") if probe.get("ok") else None
        if sid:
            det = _get(f"{base}/kinks/{sid}")
            url = first_catalog_image_probe_url(base, det)
            kid = sid
    if not url:
        print(
            "verify_hf_stack: no resolvable asset_url on catalog-sample / first 500 kinks (unexpected)",
            file=sys.stderr,
        )
        return
    try:
        assert_catalog_image_url_works(url, timeout_s=t_img)
    except HTTPError as e:
        if e.code == 404:
            if require_catalog_image:
                raise RuntimeError(
                    f"catalog image GET 404 for {kid!r}; full-catalog media is required but bytes are missing"
                ) from e
            print(
                f"verify_hf_stack: catalog image GET 404 for {kid!r} (missing on-disk bytes for that path; "
                "bundled seed images above still prove /media is wired).",
                file=sys.stderr,
            )
            return
        raise
    print(
        f"verify_hf_stack: live catalog image probe ok (kink_id={kid!r}, url={url[:72]!r}…)",
        file=sys.stderr,
    )


def run_smoke_checks(base: str) -> None:
    """Public + authenticated smoke: health stats, ``/kinks``, search, new user, one play, recommendations."""
    root = base.rstrip("/")
    if "hf.space" in root:
        _wait_for_health(root, timeout_s=float(os.environ.get("HF_VERIFY_HEALTH_TIMEOUT_S", "3600")))
    else:
        _wait_for_health(root, timeout_s=120.0)

    health = _get(f"{root}/health")
    _require(health.get("ok") is True, f"health not ok: {health!r}")
    stats = health.get("stats") or {}
    nk = int(stats.get("kinks") or 0)
    min_k = int(os.environ.get("HF_VERIFY_MIN_KINKS", "1") or "1")
    _require(nk >= min_k, f"catalog too small for smoke: kinks={nk}, need>={min_k}")
    print(f"verify_hf_stack smoke: health ok, kinks_in_catalog={nk}", file=sys.stderr)

    _assert_coep_credentialless(root)
    _assert_live_frontend_recs_error_shaping(root)

    probe = _wait_catalog_sample_ready(root)
    _require(probe.get("ok") is True, f"catalog-sample not ok: {probe!r}")
    sample_id = probe.get("sample_kink_id")
    _require(bool(sample_id), f"catalog-sample missing kink id: {probe!r}")
    detail = _get(f"{root}/kinks/{sample_id}")
    _require(detail.get("id") == sample_id, f"kink detail mismatch: {detail!r}")
    print(f"verify_hf_stack smoke: catalog-sample id={sample_id!r}", file=sys.stderr)

    u = _post(f"{root}/users", {})
    _require(bool(u.get("id") and u.get("private_token")), f"create user bad payload: {u!r}")
    _assert_auth_login_http_contract(root, u)
    tok = {"x-private-token": u["private_token"]}
    play = _post(
        f"{root}/users/{u['id']}/plays",
        {"kink_id": sample_id, "interest_state": "love", "directions": ["together"]},
        headers=tok,
    )
    _require(bool(play.get("plays", {}).get(sample_id)), f"play not saved: {play!r}")
    rec = _get(f"{root}/users/{u['id']}/recommendations?limit=8", headers=tok)
    nrec = len(rec.get("items", []))
    _require(nrec >= 1, f"recommendations empty: {rec!r}")
    print(f"verify_hf_stack smoke: recommendations count={nrec} for new user", file=sys.stderr)
    _assert_catalog_https_image_fetchable(root)
    print("verify_hf_stack: SMOKE OK", file=sys.stderr)


def run_checks(base: str) -> None:
    root = base.rstrip("/")
    # Direct app host (proxied Gradio/FastAPI). The hub page huggingface.co/spaces/... is not the API origin.
    if "hf.space" in root:
        # First boot may download a multi-GB catalog from the Hub before /health responds.
        _wait_for_health(root, timeout_s=float(os.environ.get("HF_VERIFY_HEALTH_TIMEOUT_S", "3600")))
        expect_k = int(os.environ.get("HF_VERIFY_EXPECT_KINKS", "0") or "0")
        min_k = int(os.environ.get("HF_VERIFY_MIN_KINKS", "0") or "0")
        if expect_k > 0 or min_k > 0:
            h = _get(f"{root}/health")
            nk = int((h.get("stats") or {}).get("kinks") or 0)
            if expect_k > 0:
                assert nk == expect_k, f"catalog count mismatch: kinks={nk}, expected={expect_k}"
            assert nk >= min_k, f"catalog too small: kinks={nk}, need>={min_k}"
    else:
        # Local/docker: catalog bootstrap may run in a background thread after bind.
        _wait_for_health(root, timeout_s=120.0)

    _assert_coep_credentialless(root)
    _assert_live_frontend_recs_error_shaping(root)
    _assert_catalog_https_image_fetchable(root)

    a = _post(f"{root}/users", {})
    b = _post(f"{root}/users", {})
    assert a["id"] and a["private_token"]
    assert b["id"] and b["private_token"]
    _assert_auth_login_http_contract(root, a)

    tok_a = {"x-private-token": a["private_token"]}
    tok_b = {"x-private-token": b["private_token"]}

    _post(
        f"{root}/users/{a['id']}/partners",
        {"partner_id": b["id"]},
        headers=tok_a,
    )
    u_b = _get(f"{root}/users/{b['id']}", headers=tok_b)
    assert a["id"] in (u_b.get("incoming_partner_requests") or []), u_b

    u_b2 = _post(
        f"{root}/users/{b['id']}/partner-requests/accept",
        {"from_user_id": a["id"]},
        headers=tok_b,
    )
    assert a["id"] in u_b2["partners"], u_b2

    grp = _post(
        f"{root}/users/{a['id']}/partner-groups",
        {"name": "Verify", "member_ids": [b["id"]]},
        headers=tok_a,
    )
    assert grp.get("id")

    klist = _get(f"{root}/kinks?limit=1")
    kitems = klist.get("items") or []
    assert kitems, ("no kinks in catalog", klist)
    sample_id = kitems[0].get("id")
    assert sample_id, kitems[0]

    play = _post(
        f"{root}/users/{a['id']}/plays",
        {"kink_id": sample_id, "interest_state": "love", "directions": ["together"]},
        headers=tok_a,
    )
    assert play["plays"].get(sample_id), play

    rec = _get(
        f"{root}/users/{a['id']}/recommendations?limit=5",
        headers=tok_a,
    )
    assert len(rec.get("items", [])) >= 1, rec

    sr = _get(f"{root}/search?q=massage&limit=5")
    assert len(sr.get("items", [])) >= 1, sr

    print("verify_hf_stack: OK", file=sys.stderr)


def docker_probe(image: str) -> int:
    port = "7860"
    cmd = [
        "docker",
        "run",
        "--rm",
        "-d",
        "-p",
        f"127.0.0.1:0:{port}",
        "-e",
        "KINK_STORE_PATH=/tmp/hf_verify_store.db",
        # Image defaults require Hub dataset; use bundled 5-kink seed for a fast local probe.
        "-e",
        "KINK_HF_REQUIRE_FULL_CATALOG=0",
        "-e",
        "KINK_HF_DATASET_REPO=",
        image,
    ]
    cid = subprocess.check_output(cmd, text=True).strip()
    try:
        port_line = subprocess.check_output(["docker", "port", cid, f"{port}/tcp"], text=True).strip()
        # e.g. 0.0.0.0:32768
        host_port = port_line.rsplit(":", 1)[-1]
        base = f"http://127.0.0.1:{host_port}"
        deadline = time.time() + 180
        last_err: Exception | None = None
        while time.time() < deadline:
            try:
                _get(f"{base}/health")
                last_err = None
                break
            except (URLError, HTTPError, TimeoutError, OSError) as e:
                last_err = e
                time.sleep(1)
        if last_err is not None:
            raise last_err
        run_checks(base)
        return 0
    finally:
        subprocess.run(["docker", "stop", cid], capture_output=True)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base-url", default=os.environ.get("HF_VERIFY_BASE_URL", ""))
    ap.add_argument(
        "--require-live",
        action="store_true",
        help="Require HF_LIVE_URL (or --base-url) for a deployed Space; exit 2 if missing",
    )
    ap.add_argument("--docker", action="store_true", help="Build image and run docker probe")
    ap.add_argument(
        "--smoke",
        action="store_true",
        help="Shorter live checks: health + catalog size + /kinks + /search + new user + recommendations (no partners)",
    )
    ap.add_argument("--image", default="kink-cli-hf:latest")
    args = ap.parse_args()
    if args.require_live:
        live = (args.base_url or os.environ.get("HF_LIVE_URL", "")).strip()
        if not live:
            print(
                "Set HF_LIVE_URL or pass --base-url to the deployed Space (*.hf.space).",
                file=sys.stderr,
            )
            return 2
        args.base_url = live
    if args.docker:
        subprocess.run(
            ["docker", "build", "-t", args.image, "."],
            cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            check=True,
        )
        return docker_probe(args.image)
    base = args.base_url or "http://127.0.0.1:7860"
    if args.smoke:
        run_smoke_checks(base)
    else:
        run_checks(base)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())