#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Audit CIL chart indexes for train/eval outcome and same-state leakage."
    )
    parser.add_argument(
        "paths",
        nargs="+",
        type=Path,
        help="Index JSON files or roots containing */index.json.",
    )
    parser.add_argument("--out", type=Path, default=None)
    parser.add_argument(
        "--allow-evaluator-outcomes",
        action="store_true",
        help="Allow non-train indexes with outcomes only when audience=evaluator_only.",
    )
    args = parser.parse_args(argv)

    indexes = _load_indexes(args.paths)
    if not indexes:
        raise SystemExit("no index.json files found")

    violations: list[str] = []
    warnings: list[str] = []
    by_split: dict[str, list[dict[str, Any]]] = {}
    for item in indexes:
        split = str(item["payload"].get("split", "unknown"))
        by_split.setdefault(split, []).append(item)
        _audit_single_index(item, violations, warnings, allow_evaluator=args.allow_evaluator_outcomes)
    _audit_cross_split_overlap(by_split, violations)

    report = {
        "status": "pass" if not violations else "fail",
        "num_indexes": len(indexes),
        "indexes": [
            {
                "path": str(item["path"]),
                "split": item["payload"].get("split"),
                "num_groups": item["payload"].get("num_groups_exported"),
                "num_rows": item["payload"].get("num_rows"),
                "include_outcomes": item["payload"].get("include_outcomes"),
                "retrieval_index_allowed": item["payload"].get("retrieval_index_allowed"),
                "audience": item["payload"].get("audience"),
            }
            for item in indexes
        ],
        "violations": violations,
        "warnings": warnings,
    }
    text = json.dumps(report, indent=2, sort_keys=True) + "\n"
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(text)
    print(text, end="")
    return 1 if violations else 0


def _load_indexes(paths: list[Path]) -> list[dict[str, Any]]:
    indexes: list[dict[str, Any]] = []
    for path in paths:
        if path.is_dir():
            candidates = sorted(path.glob("**/index.json"))
        else:
            candidates = [path]
        for candidate in candidates:
            if candidate.exists() and candidate.name == "index.json":
                indexes.append({"path": candidate, "payload": json.loads(candidate.read_text())})
    return indexes


def _audit_single_index(
    item: dict[str, Any],
    violations: list[str],
    warnings: list[str],
    *,
    allow_evaluator: bool,
) -> None:
    path = item["path"]
    payload = item["payload"]
    split = str(payload.get("split", "unknown"))
    include_outcomes = bool(payload.get("include_outcomes", False))
    retrieval_allowed = bool(payload.get("retrieval_index_allowed", False))
    audience = str(payload.get("audience", "unknown"))
    if split != "train" and retrieval_allowed:
        violations.append(f"{path}: non-train split is marked retrieval_index_allowed")
    if split == "train" and not retrieval_allowed:
        warnings.append(f"{path}: train split is not marked retrieval_index_allowed")
    if split != "train" and include_outcomes:
        if not (allow_evaluator and audience == "evaluator_only"):
            violations.append(
                f"{path}: non-train outcomes are included outside explicit evaluator-only audit mode"
            )
    if retrieval_allowed and include_outcomes and split != "train":
        violations.append(f"{path}: retrieval index includes non-train outcomes")
    if not payload.get("content_hash"):
        warnings.append(f"{path}: missing content_hash")
    for shard in payload.get("shards", []):
        shard_path = path.parent / str(shard.get("path", ""))
        if not shard_path.exists():
            violations.append(f"{path}: missing shard {shard_path}")


def _audit_cross_split_overlap(
    by_split: dict[str, list[dict[str, Any]]],
    violations: list[str],
) -> None:
    train_group_hashes = _hashes_for_split(by_split.get("train", []), "group_hashes")
    train_state_hashes = _hashes_for_split(by_split.get("train", []), "state_hashes")
    for split, items in sorted(by_split.items()):
        if split == "train":
            continue
        group_overlap = train_group_hashes & _hashes_for_split(items, "group_hashes")
        state_overlap = train_state_hashes & _hashes_for_split(items, "state_hashes")
        if group_overlap:
            violations.append(
                f"train/{split}: {len(group_overlap)} overlapping group hashes"
            )
        if state_overlap:
            violations.append(
                f"train/{split}: {len(state_overlap)} overlapping simulator-state hashes"
            )


def _hashes_for_split(items: list[dict[str, Any]], key: str) -> set[str]:
    output: set[str] = set()
    for item in items:
        output.update(str(value) for value in item["payload"].get(key, []))
    return output


if __name__ == "__main__":
    raise SystemExit(main())