#!/usr/bin/env python3
"""Summarize an immutable dataset manifest and its selected training prefix."""

from __future__ import annotations

import argparse
import hashlib
import json
from collections import Counter, defaultdict
from pathlib import Path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", required=True, type=Path)
    parser.add_argument("--training-prefix", type=int, default=2500)
    parser.add_argument(
        "--class-component",
        type=int,
        default=2,
        help="zero-based path component containing a class ID",
    )
    parser.add_argument("--output-json", required=True, type=Path)
    parser.add_argument("--output-markdown", required=True, type=Path)
    args = parser.parse_args()
    if args.training_prefix < 1 or args.class_component < 0:
        raise SystemExit("prefix must be positive and class component non-negative")

    manifest_bytes = args.manifest.read_bytes()
    manifest = json.loads(manifest_bytes)
    if manifest.get("schema_version") != 1 or not manifest.get("images"):
        raise SystemExit("unsupported or empty dataset manifest")

    split_counts: Counter[str] = Counter()
    source_counts: dict[str, Counter[str]] = defaultdict(Counter)
    group_splits: dict[tuple[str, str], str] = {}
    hashes: set[str] = set()
    paths: set[str] = set()
    class_counts: dict[str, Counter[str]] = defaultdict(Counter)
    train_entries: list[dict[str, object]] = []
    for image in manifest["images"]:
        split = str(image["split"])
        source = str(image["source"])
        group = str(image["group"])
        path = str(image["path"])
        digest = str(image["sha256"])
        if path in paths or digest in hashes:
            raise SystemExit(f"duplicate path or content in manifest: {path}")
        paths.add(path)
        hashes.add(digest)
        key = (source, group)
        previous = group_splits.setdefault(key, split)
        if previous != split:
            raise SystemExit(f"group {source}/{group} crosses {previous} and {split}")
        parts = Path(path).parts
        if len(parts) <= args.class_component:
            raise SystemExit(f"{path} has no class component {args.class_component}")
        class_id = parts[args.class_component]
        split_counts[split] += 1
        source_counts[split][source] += 1
        class_counts[split][class_id] += 1
        if split == "train":
            train_entries.append(image)

    if len(train_entries) < args.training_prefix:
        raise SystemExit(
            f"training split has {len(train_entries)} entries, fewer than prefix "
            f"{args.training_prefix}"
        )
    prefix_classes = Counter(
        Path(str(image["path"])).parts[args.class_component]
        for image in train_entries[: args.training_prefix]
    )
    artifact = {
        "schema_version": 1,
        "dataset_name": manifest.get("name"),
        "manifest_sha256": hashlib.sha256(manifest_bytes).hexdigest(),
        "provenance": manifest.get("provenance", []),
        "images": len(manifest["images"]),
        "unique_content_hashes": len(hashes),
        "groups": len(group_splits),
        "split_counts": dict(sorted(split_counts.items())),
        "source_counts_by_split": {
            split: dict(sorted(counts.items()))
            for split, counts in sorted(source_counts.items())
        },
        "class_counts_by_split": {
            split: dict(sorted(counts.items()))
            for split, counts in sorted(class_counts.items())
        },
        "selected_training_prefix": args.training_prefix,
        "selected_training_prefix_class_counts": dict(sorted(prefix_classes.items())),
    }
    args.output_json.parent.mkdir(parents=True, exist_ok=True)
    args.output_markdown.parent.mkdir(parents=True, exist_ok=True)
    args.output_json.write_text(json.dumps(artifact, indent=2) + "\n", encoding="utf-8")

    lines = [
        "# Dataset summary",
        "",
        f"Manifest SHA-256: `{artifact['manifest_sha256']}`.",
        "",
        "| Split | Images | Classes | Sources |",
        "|---|---:|---:|---|",
    ]
    for split, count in sorted(split_counts.items()):
        lines.append(
            f"| {split} | {count} | {len(class_counts[split])} "
            f"| {', '.join(f'{source}: {n}' for source, n in sorted(source_counts[split].items()))} |"
        )
    lines.extend(
        [
            "",
            f"The selected training prefix contains {args.training_prefix} images:",
            "",
            "| Class ID | Images |",
            "|---|---:|",
        ]
    )
    lines.extend(f"| {name} | {count} |" for name, count in sorted(prefix_classes.items()))
    args.output_markdown.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"wrote {args.output_json} and {args.output_markdown}")


if __name__ == "__main__":
    main()