#!/usr/bin/env python3
"""validate_dataset.py — trainpipe dataset preflight (stdlib only).

Checks that a dataset directory matches what the chosen model family expects,
WITHOUT importing lerobot/torch (fast, venv-independent). Last stdout line is
always:

    PIPELINE_RESULT {"ok": true|false, "message": ..., "episodes": N, ...}

    python validate_dataset.py --root <dir> --expect lerobot-v3 --policy smolvla
    python validate_dataset.py --root <file.jsonl> --expect jsonl-grounding
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

# Feature keys every lerobot policy trained by the pipeline needs. Image keys
# are prefix-matched (any observation.images.* camera set is fine).
LEROBOT_REQUIRED = ["observation.state", "action"]
JSONL_REQUIRED = ["image", "instruction", "pixel"]


def result(ok: bool, message: str, **extra):
    print("PIPELINE_RESULT " + json.dumps({"ok": ok, "message": message, **extra}))
    sys.exit(0 if ok else 1)


def validate_lerobot(root: Path, policy: str):
    info_path = root / "meta" / "info.json"
    if not info_path.exists():
        result(False, f"not a LeRobot v3 dataset (no {info_path}). "
               "Collect one with sim/scripts/collect_rx1_lerobot.py or pass --collect N",
               code="unrecognized_format")
    try:
        info = json.loads(info_path.read_text())
    except Exception as e:  # corrupted meta
        result(False, f"unreadable meta/info.json: {e}", code="corrupt_meta")

    version = str(info.get("codebase_version", ""))
    if not version.startswith("v3"):
        result(False, f"dataset codebase_version {version!r}; the vendored lerobot speaks v3.x",
               code="version_mismatch")

    features = info.get("features", {})
    missing = [k for k in LEROBOT_REQUIRED if k not in features]
    images = [k for k in features if k.startswith("observation.images.")]
    if missing or not images:
        result(False, f"missing features {missing or 'observation.images.*'} — found {sorted(features)}",
               code="missing_features")

    episodes = info.get("total_episodes", 0)
    if episodes < 5:
        result(False, f"only {episodes} episodes — too few to train on", code="too_small",
               episodes=episodes)

    result(True, "lerobot-v3 dataset OK",
           episodes=episodes, frames=info.get("total_frames"),
           robot=info.get("robot_type"), cameras=images,
           action_dim=(features.get("action", {}).get("shape") or [None])[0])


def validate_jsonl(root: Path):
    # accept either the jsonl file itself or a directory containing grounding.jsonl
    f = root if root.is_file() else root / "grounding.jsonl"
    if not f.exists():
        result(False, f"no grounding jsonl at {f}. Generate with "
               "sim/scripts/pipeline/qwen_dataset_gen.py or pass --collect N",
               code="unrecognized_format")
    n, bad = 0, 0
    with open(f) as fh:
        for i, line in enumerate(fh):
            if i >= 500:
                break
            line = line.strip()
            if not line:
                continue
            try:
                row = json.loads(line)
                if any(k not in row for k in JSONL_REQUIRED):
                    bad += 1
                else:
                    n += 1
            except json.JSONDecodeError:
                bad += 1
    if n < 20:
        result(False, f"only {n} valid grounding rows (and {bad} malformed) — need at least 20",
               code="too_small", rows=n)
    if bad > n // 4:
        result(False, f"{bad} malformed rows vs {n} valid — dataset looks corrupted", code="corrupt")
    result(True, "grounding jsonl OK", rows=n, episodes=n)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", required=True)
    ap.add_argument("--expect", required=True, choices=["lerobot-v3", "jsonl-grounding"])
    ap.add_argument("--policy", default="")
    a = ap.parse_args()

    root = Path(a.root)
    if not root.exists():
        result(False, f"dataset path does not exist: {root}", code="missing")
    if a.expect == "lerobot-v3":
        validate_lerobot(root, a.policy)
    else:
        validate_jsonl(root)


if __name__ == "__main__":
    main()