#!/usr/bin/env python3
"""Exhaustively validate full Grasp-Anything++ streaming shards."""

import argparse
import gzip
import json
import pickle
import tarfile
import time
import zlib
from pathlib import Path


SPLITS = ("train", "seen", "unseen")
SUFFIX_ORDER = ("jpg", "ptz", "pklz", "npyz", "json")


def log(message):
    print(time.strftime("[%Y-%m-%d %H:%M:%S]"), message, flush=True)


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--root", type=Path,
        default=Path("/anvme/workspace/b305bb10-vlagrasp/graspanything_full"))
    parser.add_argument("--report", type=Path)
    parser.add_argument("--progress-every", type=int, default=100000)
    parser.add_argument("--worker-count", type=int, default=1)
    parser.add_argument("--worker-id", type=int, default=0)
    parser.add_argument("--manifest-only", action="store_true")
    parser.add_argument("--aggregate-workers", type=int, default=0)
    args = parser.parse_args()
    if args.worker_count <= 0 or not 0 <= args.worker_id < args.worker_count:
        parser.error("require worker-count > 0 and 0 <= worker-id < worker-count")
    return args


def atomic_json(path, value):
    partial = path.with_name(path.name + ".partial")
    partial.write_text(json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
                       encoding="utf-8")
    partial.replace(path)


def validate_object_manifest(root, expected, represented_objects):
    path = root / "manifests/objects.jsonl.gz"
    counts = {split: 0 for split in SPLITS}
    previous = None
    manifest_objects = {split: set() for split in SPLITS}
    with gzip.open(str(path), "rt", encoding="utf-8") as fp:
        for line_number, line in enumerate(fp, 1):
            item = json.loads(line)
            required = {"object_id", "scene_id", "obj_id", "category", "split"}
            if set(item) != required:
                raise ValueError("Object manifest line {} fields differ: {}".format(
                    line_number, sorted(item)))
            object_id = "{}_{}".format(item["scene_id"], item["obj_id"])
            if item["object_id"] != object_id or item["split"] not in SPLITS:
                raise ValueError("Invalid object manifest line {}".format(line_number))
            order_key = (item["scene_id"], item["obj_id"])
            if previous is not None and order_key <= previous:
                raise ValueError("Object manifest is not strictly ordered at line {}".format(line_number))
            previous = order_key
            split = item["split"]
            if object_id in manifest_objects[split]:
                raise ValueError("Duplicate object {}".format(object_id))
            manifest_objects[split].add(object_id)
            counts[split] += 1
    if counts != expected:
        raise ValueError("Object manifest counts {} != {}".format(counts, expected))
    for split in SPLITS:
        if manifest_objects[split] != represented_objects[split]:
            raise ValueError("{} object manifest/tar mismatch".format(split))
    return counts


def validate_manifests_only(args, summary):
    root = args.root
    manifest_root = root / "manifests"
    samples = {split: set() for split in SPLITS}
    objects = {split: set() for split in SPLITS}
    categories = {split: set() for split in SPLITS}
    report = {
        "format": "graspanything-pp-full-validation-v1",
        "dataset_format": summary.get("format"),
        "root": str(root),
        "splits": {},
        "checks": {},
    }
    for split in SPLITS:
        expected = summary["splits"][split]
        path = manifest_root / (split + ".jsonl")
        with path.open("rb") as fp:
            for line_number, line in enumerate(fp, 1):
                metadata = json.loads(line)
                sample_id = metadata.get("sample_id")
                if not isinstance(sample_id, str) or metadata.get("split") != split:
                    raise ValueError("Invalid {} manifest line {}".format(split, line_number))
                object_id = sample_id.rsplit("_", 1)[0]
                if metadata.get("object_id") != object_id:
                    raise ValueError("Object mismatch on {} line {}".format(split, line_number))
                category = metadata.get("category")
                if not isinstance(category, str) or not category:
                    raise ValueError("Empty category on {} line {}".format(split, line_number))
                if sample_id in samples[split]:
                    raise ValueError("Duplicate sample {}".format(sample_id))
                samples[split].add(sample_id)
                objects[split].add(object_id)
                categories[split].add(category)
                if args.progress_every and line_number % args.progress_every == 0:
                    log("{}: audited {:,}/{:,} manifest rows".format(
                        split, line_number, expected["samples"]))
        if len(samples[split]) != expected["samples"]:
            raise ValueError("{} manifest sample count mismatch".format(split))
        if len(objects[split]) != expected["objects"]:
            raise ValueError("{} manifest object count mismatch".format(split))
        if len(categories[split]) != expected["categories"]:
            raise ValueError("{} manifest category count mismatch".format(split))
        report["splits"][split] = {
            "samples": len(samples[split]), "objects": len(objects[split]),
            "categories": len(categories[split]), "shards": expected["shards"],
            "bytes": expected["bytes"],
        }
    for index, left in enumerate(SPLITS):
        for right in SPLITS[index + 1:]:
            if samples[left] & samples[right] or objects[left] & objects[right]:
                raise ValueError("Global sample/object overlap between {} and {}".format(left, right))
    if (categories["train"] | categories["seen"]) & categories["unseen"]:
        raise ValueError("Global seen/unseen category overlap")
    expected_objects = {split: summary["splits"][split]["objects"] for split in SPLITS}
    validate_object_manifest(root, expected_objects, objects)
    if args.aggregate_workers:
        worker_root = manifest_root / "validation_workers"
        for split in SPLITS:
            shard_ids = set()
            worker_samples = 0
            worker_bytes = 0
            for worker_id in range(args.aggregate_workers):
                worker = json.loads((worker_root / "worker-{:02d}.json".format(worker_id)).read_text(
                    encoding="utf-8"))
                if worker.get("worker_count") != args.aggregate_workers:
                    raise ValueError("Validation worker-count mismatch")
                item = worker["splits"][split]
                overlap = shard_ids & set(item["shard_indices"])
                if overlap:
                    raise ValueError("Validation workers overlap on {} shards".format(split))
                shard_ids.update(item["shard_indices"])
                worker_samples += item["samples"]
                worker_bytes += item["bytes"]
            if shard_ids != set(range(summary["splits"][split]["shards"])):
                raise ValueError("Validation workers did not cover all {} shards".format(split))
            if worker_samples != summary["splits"][split]["samples"]:
                raise ValueError("Validation worker {} sample sum mismatch".format(split))
            if worker_bytes != summary["splits"][split]["bytes"]:
                raise ValueError("Validation worker {} byte sum mismatch".format(split))
    report["checks"] = {
        "all_zlib_payloads_decoded": bool(args.aggregate_workers),
        "all_jpeg_markers_valid": bool(args.aggregate_workers),
        "all_instructions_nonempty_strings": bool(args.aggregate_workers),
        "all_tar_member_groups_exact": bool(args.aggregate_workers),
        "all_tar_json_matches_combined_manifest": bool(args.aggregate_workers),
        "sample_ids_pairwise_disjoint": True,
        "object_ids_pairwise_disjoint": True,
        "seen_unseen_categories_disjoint": True,
        "object_manifest_matches_represented_objects": True,
        "partial_files": 0,
        "payload_worker_count": args.aggregate_workers,
    }
    report["completed_at_unix"] = int(time.time())
    report_path = args.report or manifest_root / "validation_report.json"
    atomic_json(report_path, report)
    log("Global manifest and worker-coverage validation OK; report: {}".format(report_path))


def main():
    args = parse_args()
    root = args.root
    manifest_root = root / "manifests"
    if not (manifest_root / "_SUCCESS").is_file():
        raise ValueError("Missing manifests/_SUCCESS")
    partials = list(root.rglob("*.partial"))
    if partials:
        raise ValueError("Found partial outputs, first entries: {}".format(partials[:10]))
    summary = json.loads((manifest_root / "dataset_summary.json").read_text(encoding="utf-8"))
    if summary.get("uses_official_object_id_manifests") is not False:
        raise ValueError("Summary does not affirm full non-manifest selection")
    if args.manifest_only:
        validate_manifests_only(args, summary)
        return
    report = {
        "format": "graspanything-pp-full-validation-v1",
        "dataset_format": summary.get("format"),
        "root": str(root),
        "splits": {},
        "checks": {},
    }
    split_samples = {split: set() for split in SPLITS}
    split_objects = {split: set() for split in SPLITS}
    split_categories = {split: set() for split in SPLITS}

    for split in SPLITS:
        expected = summary["splits"][split]
        shard_paths = sorted((root / "shards" / split).glob("ga-pp-{}-*.tar".format(split)))
        if len(shard_paths) != expected["shards"]:
            raise ValueError("{} shard count {} != {}".format(
                split, len(shard_paths), expected["shards"]))
        expected_names = ["ga-pp-{}-{:05d}.tar".format(split, index)
                          for index in range(len(shard_paths))]
        if [path.name for path in shard_paths] != expected_names:
            raise ValueError("{} shard sequence has gaps or extra files".format(split))
        actual_bytes = sum(path.stat().st_size for path in shard_paths)
        if actual_bytes != expected["bytes"]:
            raise ValueError("{} bytes {} != {}".format(split, actual_bytes, expected["bytes"]))
        count = 0
        assigned_expected = 0
        assigned_bytes = 0
        assigned_shards = []
        manifest_path = manifest_root / (split + ".jsonl")
        with manifest_path.open("rb") as manifest_fp:
            for shard_number, shard_path in enumerate(shard_paths):
                shard_expected = min(expected["samples"] - shard_number * expected["samples_per_shard"],
                                     expected["samples_per_shard"])
                if args.worker_count > 1 and shard_number % args.worker_count != args.worker_id:
                    for _ in range(shard_expected):
                        if not manifest_fp.readline():
                            raise ValueError("{} manifest ended while worker skipped shard".format(split))
                    continue
                assigned_expected += shard_expected
                assigned_bytes += shard_path.stat().st_size
                assigned_shards.append(shard_number)
                per_shard = 0
                pending_id = None
                pending_suffixes = []
                pending_json = None
                with tarfile.open(str(shard_path), "r:") as archive:
                    for member in archive:
                        if not member.isfile() or "." not in member.name:
                            raise ValueError("Unexpected member {} in {}".format(
                                member.name, shard_path))
                        sample_id, suffix = member.name.rsplit(".", 1)
                        expected_suffix = SUFFIX_ORDER[len(pending_suffixes)]
                        if pending_id is None:
                            pending_id = sample_id
                        if sample_id != pending_id or suffix != expected_suffix:
                            raise ValueError("Member order/group mismatch in {} at {}".format(
                                shard_path, member.name))
                        extracted = archive.extractfile(member)
                        if extracted is None:
                            raise ValueError("Cannot extract {}".format(member.name))
                        payload = extracted.read()
                        if suffix == "jpg":
                            if not (payload.startswith(b"\xff\xd8") and payload.endswith(b"\xff\xd9")):
                                raise ValueError("Invalid JPEG {}".format(member.name))
                        elif suffix in ("ptz", "pklz", "npyz"):
                            decoded = zlib.decompress(payload)
                            if not decoded:
                                raise ValueError("Empty decoded payload {}".format(member.name))
                            if suffix == "pklz":
                                instruction = pickle.loads(decoded)
                                if not isinstance(instruction, str) or not instruction.strip():
                                    raise ValueError("Invalid instruction {}".format(member.name))
                        else:
                            pending_json = json.loads(payload)
                        pending_suffixes.append(suffix)
                        if len(pending_suffixes) != len(SUFFIX_ORDER):
                            continue
                        metadata = pending_json
                        if (metadata is None or metadata.get("sample_id") != pending_id or
                                metadata.get("split") != split):
                            raise ValueError("Metadata mismatch for {}".format(pending_id))
                        manifest_line = manifest_fp.readline()
                        if not manifest_line:
                            raise ValueError("{} combined manifest ended early".format(split))
                        if json.loads(manifest_line) != metadata:
                            raise ValueError("{} manifest/tar JSON mismatch for {}".format(
                                split, pending_id))
                        if pending_id in split_samples[split]:
                            raise ValueError("Duplicate sample {}".format(pending_id))
                        split_samples[split].add(pending_id)
                        object_id = pending_id.rsplit("_", 1)[0]
                        if metadata.get("object_id") != object_id:
                            raise ValueError("Object ID mismatch for {}".format(pending_id))
                        split_objects[split].add(object_id)
                        split_categories[split].add(metadata.get("category"))
                        count += 1
                        per_shard += 1
                        pending_id = None
                        pending_suffixes = []
                        pending_json = None
                        if args.progress_every and count % args.progress_every == 0:
                            log("{}: validated {:,}/{:,} samples".format(
                                split, count, expected["samples"]))
                if pending_id is not None:
                    raise ValueError("Incomplete last sample in {}".format(shard_path))
                if per_shard != shard_expected:
                    raise ValueError("{} samples {} != {}".format(
                        shard_path, per_shard, shard_expected))
            if manifest_fp.readline():
                raise ValueError("{} combined manifest has extra rows".format(split))
        target_count = expected["samples"] if args.worker_count == 1 else assigned_expected
        if count != target_count:
            raise ValueError("{} sample count {} != {}".format(split, count, target_count))
        if args.worker_count == 1 and len(split_objects[split]) != expected["objects"]:
            raise ValueError("{} object count {} != {}".format(
                split, len(split_objects[split]), expected["objects"]))
        if args.worker_count == 1 and len(split_categories[split]) != expected["categories"]:
            raise ValueError("{} category count {} != {}".format(
                split, len(split_categories[split]), expected["categories"]))
        report["splits"][split] = {
            "samples": count,
            "objects": len(split_objects[split]),
            "categories": len(split_categories[split]),
            "shards": len(assigned_shards) if args.worker_count > 1 else len(shard_paths),
            "shard_indices": assigned_shards if args.worker_count > 1 else list(range(len(shard_paths))),
            "bytes": assigned_bytes if args.worker_count > 1 else actual_bytes,
        }
        log("{}: exhaustive shard validation OK".format(split))

    if args.worker_count > 1:
        report["worker_id"] = args.worker_id
        report["worker_count"] = args.worker_count
        report["checks"] = {
            "assigned_zlib_payloads_decoded": True,
            "assigned_jpeg_markers_valid": True,
            "assigned_instructions_nonempty_strings": True,
            "assigned_tar_member_groups_exact": True,
            "assigned_tar_json_matches_manifest": True,
        }
        worker_root = manifest_root / "validation_workers"
        worker_root.mkdir(parents=True, exist_ok=True)
        atomic_json(worker_root / "worker-{:02d}.json".format(args.worker_id), report)
        log("Payload worker {}/{} complete".format(args.worker_id, args.worker_count))
        return

    for index, left in enumerate(SPLITS):
        for right in SPLITS[index + 1:]:
            sample_overlap = split_samples[left] & split_samples[right]
            object_overlap = split_objects[left] & split_objects[right]
            if sample_overlap or object_overlap:
                raise ValueError("Overlap {} / {}: samples={}, objects={}".format(
                    left, right, len(sample_overlap), len(object_overlap)))
    base_categories = split_categories["train"] | split_categories["seen"]
    category_overlap = base_categories & split_categories["unseen"]
    if category_overlap:
        raise ValueError("Seen/unseen category overlap: {}".format(sorted(category_overlap)))
    object_counts = {split: summary["splits"][split]["objects"] for split in SPLITS}
    validate_object_manifest(root, object_counts, split_objects)
    report["checks"] = {
        "all_zlib_payloads_decoded": True,
        "all_jpeg_markers_valid": True,
        "all_instructions_nonempty_strings": True,
        "all_tar_member_groups_exact": True,
        "combined_manifests_match_tar_json": True,
        "sample_ids_pairwise_disjoint": True,
        "object_ids_pairwise_disjoint": True,
        "seen_unseen_categories_disjoint": True,
        "object_manifest_matches_represented_objects": True,
        "partial_files": 0,
    }
    report["completed_at_unix"] = int(time.time())
    report_path = args.report or manifest_root / "validation_report.json"
    atomic_json(report_path, report)
    log("All splits passed exhaustive validation; report: {}".format(report_path))


if __name__ == "__main__":
    main()