"""
grading.py — canonical Kalanos RDQ metadata-only grading core.

SINGLE SOURCE OF TRUTH for metadata-only ("0.1-meta") scoring. Imported by:
  - batch_grade.py   (bulk grading of the launch list)
  - app.py           (the Hugging Face Space on-demand grader)

so that a dataset graded on the Space and the same dataset graded in the batch
produce byte-identical *scores* — and, when HF returns the same metadata,
byte-identical *reports*. All scoring logic lives here and nowhere else; do not
re-implement any part of it in app.py or batch_grade.py. Changing a weight or a
check here changes both surfaces at once.

Metadata-only grade: `format` and `coverage` are MEASURED from repository
metadata; `sync` and `outliers` are NOT measured and are scored null. The
disclosure in `notes` must travel with the report — these grades are not a
substitute for a measured (full) grade and must not be published publicly
without that disclosure intact.
"""
import io
import json
import os
import time
from datetime import datetime, timezone

import requests

HF = "https://huggingface.co"
GRADER_VERSION = "0.1-meta"
REPORT_KIND = "metadata"

# Default hard cap for the repo-size tree walk. batch_grade.py keeps the full
# 45s (throughput over latency); the Space passes a smaller value so an
# on-demand grade of a huge repo can't hang the UI (size never affects scores).
SIZE_WALK_DEADLINE_S = 45


def UTC():
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _auth_headers():
    """Attach an HF token only if one is present in the environment.
    Public datasets resolve fine without it; this keeps the free public Space
    working with no secret set, while private/gated repos work when HF_TOKEN is."""
    tok = os.environ.get("HF_TOKEN")
    return {"Authorization": f"Bearer {tok}"} if tok else {}


def hf_api(path):
    r = requests.get(f"{HF}/api/{path}", timeout=30, headers=_auth_headers())
    if r.status_code == 404:
        return None
    r.raise_for_status()
    return r.json()


def hf_file(slug, path):
    r = requests.get(f"{HF}/datasets/{slug}/resolve/main/{path}", timeout=60, headers=_auth_headers())
    if r.status_code in (401, 403, 404):
        return None
    r.raise_for_status()
    return r


def fetch_repo_size_bytes(slug, deadline_s=SIZE_WALK_DEADLINE_S):
    """Repo size in bytes. Returns (size_bytes, complete):
      - complete=True  -> size_bytes is the real total (from usedStorage, or a tree
        walk that finished before the cap).
      - complete=False, size_bytes not None -> LOWER BOUND: the tree walk hit the
        page/time cap before finishing (huge repos, e.g. droid_lerobot's ~1.7TB /
        tens of thousands of files).
      - complete=False, size_bytes is None -> couldn't determine size at all.

    Tries the dataset info endpoint's usedStorage first: one request regardless of
    repo size. Falls back to summing the paginated tree API only if that's absent,
    with a hard page/time cap so one huge repo can't stall the whole batch (or the
    interactive Space, which passes a shorter deadline_s)."""
    headers = _auth_headers()

    try:
        r = requests.get(f"{HF}/api/datasets/{slug}?expand[]=usedStorage", headers=headers, timeout=30)
        if r.status_code == 200:
            used = r.json().get("usedStorage")
            if isinstance(used, int) and used > 0:
                return used, True
    except Exception:
        pass

    url = f"{HF}/api/datasets/{slug}/tree/main?recursive=true"
    total = 0
    pages = 0
    deadline = time.time() + deadline_s  # hard cap per dataset — large repos become a lower bound, not a hang
    try:
        while url and pages < 200 and time.time() < deadline:
            r = requests.get(url, headers=headers, timeout=20)
            if r.status_code == 404:
                return None, False
            r.raise_for_status()
            for item in r.json():
                if item.get("type") == "file" and isinstance(item.get("size"), int):
                    total += item["size"]
            url = r.links.get("next", {}).get("url")
            pages += 1
        return total, (url is None)  # url is None => walked the whole tree; else capped out
    except Exception:
        return (total or None), False


def fetch_metadata(slug, size_deadline_s=SIZE_WALK_DEADLINE_S):
    """Real facts from HF: license, gated, and LeRobot meta if present.
    Returns None if the dataset does not exist (or the API 404s)."""
    info = hf_api(f"datasets/{slug}")
    if info is None:
        return None
    card = info.get("cardData") or {}
    meta = {
        "license": (card.get("license") if isinstance(card.get("license"), str) else (card.get("license") or [""])[0]) or "unknown",
        "gated": bool(info.get("gated")),
        "downloads": info.get("downloads", 0),
        "last_modified": info.get("lastModified", ""),
        "sha": (info.get("sha") or "")[:7],
        "lerobot": None, "episodes_meta": None, "size_bytes": None, "size_complete": False,
    }
    meta["size_bytes"], meta["size_complete"] = fetch_repo_size_bytes(slug, deadline_s=size_deadline_s)
    for path in ("meta/info.json",):
        r = hf_file(slug, path)
        if r is not None:
            try:
                meta["lerobot"] = r.json()
            except Exception:
                pass
    # episode stats (v2 jsonl or v3 parquet listing is heavy; jsonl is cheap when present)
    r = hf_file(slug, "meta/episodes.jsonl")
    if r is not None:
        lengths = []
        for line in io.StringIO(r.text):
            try:
                lengths.append(json.loads(line).get("length", 0))
            except Exception:
                pass
        if lengths:
            meta["episodes_meta"] = {"count": len(lengths), "lengths": lengths}
    return meta


def hist(values, bins):
    counts = [0] * (len(bins) - 1)
    for v in values:
        for i in range(len(bins) - 1):
            if bins[i] <= v < bins[i + 1]:
                counts[i] += 1
                break
        else:
            counts[-1] += 1
    return {"bins": bins, "counts": counts}


def gini(xs):
    xs = sorted(x for x in xs if x >= 0)
    n = len(xs)
    s = sum(xs)
    if n == 0 or s == 0:
        return 0.0
    g = 0.0
    for i, x in enumerate(xs, 1):
        g += (2 * i - n - 1) * x
    return round(g / (n * s), 3)


def meta_grade(slug, m):
    """Metadata-only v0 grade. Format & coverage are measured; sync & outliers are
    provisional placeholders (null) disclosed in the notes. Do NOT publish these
    publicly without the disclosure intact.

    `m` is the dict returned by fetch_metadata(). This is the ONE function whose
    output defines a metadata grade; both the batch and the Space call it."""
    lr = m.get("lerobot") or {}
    fps = lr.get("fps", 0)
    total_eps = lr.get("total_episodes") or (m.get("episodes_meta") or {}).get("count") or 0
    total_frames = lr.get("total_frames", 0)
    tasks = lr.get("total_tasks", 0)
    features = lr.get("features", {}) or {}
    cams = [k for k in features if "image" in k or features[k].get("dtype") in ("video", "image")]
    fmt_checks = {
        "has_lerobot_meta": lr != {},
        "fps_declared": bool(fps),
        "features_declared": bool(features),
        "chunked_layout": bool(lr.get("data_path")),
    }
    fmt_score = 100.0 * sum(fmt_checks.values()) / len(fmt_checks) if lr else 40.0
    lens = (m.get("episodes_meta") or {}).get("lengths") or []
    secs = [l / fps for l in lens] if fps and lens else []
    ep_hist = hist(secs, [0, 10, 20, 30, 45, 60, 90, 120]) if secs else {"bins": [0, 1], "counts": [total_eps]}
    task_balance = 0.0  # per-task counts not in cheap metadata; neutral
    cov_score = min(100.0, 30 + 8 * tasks + 6 * len(cams)) if lr else 50.0
    hours = round(total_frames / fps / 3600, 1) if fps and total_frames else 0
    # meta/info.json's codebase_version already includes a leading "v" (e.g. "v2.0");
    # strip it before prepending our own so we never emit "LeRobot vv2.0".
    raw_version = str(lr.get("codebase_version", "")).strip()
    version = raw_version.lstrip("vV")
    fmt_label = f"LeRobot v{version}" if version else ("LeRobot" if lr else "unknown")

    size_bytes = m.get("size_bytes")
    size_complete = m.get("size_complete")
    size_gb = round(size_bytes / 1e9, 2) if isinstance(size_bytes, int) else 0
    if not isinstance(size_bytes, int):
        size_note = ["Repo size could not be determined; size_gb=0 is a placeholder, not a measurement."]
    elif not size_complete:
        size_note = [f"Repo size is a LOWER BOUND ({size_gb} GB): the file listing was capped before finishing "
                     f"(large repo). Re-run with a longer cap or rely on usedStorage for an exact total."]
    else:
        size_note = []

    return {
        "report_id": "",
        "report_kind": REPORT_KIND,
        "grader_version": GRADER_VERSION,
        "graded_at": UTC(),
        "source": {"kind": "hf", "ref": slug, "url": f"{HF}/datasets/{slug}",
                   "commit": m["sha"], "license": m["license"].lower(), "gated": m["gated"]},
        "dataset": {"format": fmt_label,
                    "episodes": total_eps, "hours": hours,
                    "embodiment": lr.get("robot_type", "unknown"),
                    "sensors": cams or ["unknown"], "size_gb": size_gb},
        "scores": {"total": round((fmt_score + cov_score) / 2, 1),
                   "sync": None, "coverage": round(cov_score, 1),
                   "outliers": None, "format": round(fmt_score, 1)},
        "metrics": {
            "sync": {"median_drift_ms": 0, "p95_drift_ms": 0, "max_drift_ms": 0, "episodes_over_10ms": 0, "dropped_frames_pct": 0},
            "coverage": {"task_families": tasks, "environments": 0, "episodes_per_task_min": 0,
                         "episodes_per_task_max": 0, "gini_task_balance": task_balance, "lighting_variation_score": 0},
            "outliers": {"flagged_episodes": 0, "flagged_pct": 0, "categories": {}},
            "format": {"schema_valid_pct": 100 if lr else 0, "calibration_present_pct": 0,
                       "units_consistent": bool(lr), "timestamps_monotonic_pct": 0},
        },
        "histograms": {"drift_ms": {"bins": [0, 1], "counts": [0]}, "episode_length_s": ep_hist},
        "flagged_sample": [],
        "notes": [
            f"Metadata-level report: file structure, coverage and episode statistics are measured from repository metadata. "
            f"Timestamp synchronization and trajectory outliers are not yet measured and are not scored.",
            f"{total_eps:,} episodes, {tasks} task(s), {len(cams)} camera stream(s), fps={fps}.",
            f"HF downloads: {m['downloads']:,}; last modified {m['last_modified'][:10]}.",
            *size_note,
        ],
    }