#!/usr/bin/env python3
import os, sys, json, time, argparse, traceback
import requests

REPO_ID = "morningstarxcdcode/base-llm-400m"
HF_TOKEN = os.environ.get("HF_TOKEN", "")
SIDECAR_KEYS = {
    "shard_idx", "filename", "num_tokens", "dtype", "size_bytes",
    "created_at", "tokens", "avg_score", "min_score", "max_score",
    "std_score", "n_above_3", "n_above_5", "score_hist",
    "modality_comp", "pillar_comp", "ts", "modality"
}

HIST_BINS = {f"{i*0.5:.1f}" for i in range(21)}

def check_sidecar(path):
    with open(path) as f:
        meta = json.load(f)
    keys = set(meta.keys())
    missing = SIDECAR_KEYS - keys
    extra = keys - SIDECAR_KEYS
    errors = []
    if missing:
        errors.append(f"  MISSING keys: {missing}")
    if extra:
        errors.append(f"  EXTRA keys: {extra}")
    hist = meta.get("score_hist", {})
    hist_keys = set(hist.keys())
    if hist_keys != HIST_BINS:
        errors.append(f"  score_hist wrong bins: missing={HIST_BINS-hist_keys} extra={hist_keys-HIST_BINS}")
    mod = meta.get("modality_comp", {})
    for k, v in mod.items():
        if v is not None and not isinstance(v, (int, float)):
            errors.append(f"  modality_comp[{k}] is {type(v).__name__}, expected int or null")
    return errors

def try_load_dataset(retries=3, backoff=30):
    from datasets import load_dataset
    for attempt in range(1, retries + 1):
        try:
            ds = load_dataset(REPO_ID, split="train", streaming=True, trust_remote_code=True)
            sample = next(iter(ds))
            print(f"  OK — dataset loads. Sample keys: {list(sample.keys())}")
            return True
        except Exception as e:
            if attempt < retries:
                print(f"  Attempt {attempt} failed: {e} — retrying in {backoff}s...")
                time.sleep(backoff)
            else:
                print(f"  FAILED after {retries} attempts: {e}")
                return False

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--retries", type=int, default=3)
    p.add_argument("--backoff", type=int, default=30)
    args = p.parse_args()

    print(f"Verifying {REPO_ID}...")
    ok = try_load_dataset(retries=args.retries, backoff=args.backoff)

    api = __import__("huggingface_hub", fromlist=["HfApi"]).HfApi()
    try:
        files = api.list_repo_files(REPO_ID, token=HF_TOKEN)
    except:
        files = []

    sidecars = [f for f in files if f.endswith(".meta.json")]
    if not sidecars:
        print("  No sidecars found — skipping schema check")
    else:
        print(f"  Checking {len(sidecars)} sidecar(s):")
        all_ok = True
        for sc in sidecars:
            local = api.hf_hub_download(REPO_ID, sc, token=HF_TOKEN)
            errs = check_sidecar(local)
            if errs:
                all_ok = False
                print(f"  {sc}:")
                for e in errs:
                    print(f"    {e}")
        if all_ok:
            print(f"  All {len(sidecars)} sidecar(s) pass schema check")

    if ok:
        print(f"\n  RESULT: PASS — viewer is healthy")
        return 0
    else:
        print(f"\n  RESULT: FAIL — viewer broken or unreachable")
        return 1

if __name__ == "__main__":
    sys.exit(main())