#!/usr/bin/env python3
"""Freeze a chronological, law-family-disjoint L3 train/development split."""

from __future__ import annotations

import argparse
import hashlib
import json
from datetime import datetime, timezone
from pathlib import Path

import pandas as pd

REPO_ROOT = Path(__file__).resolve().parents[1]


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def resolve_recorded_path(value: str, receipt_path: Path) -> Path:
    candidate = Path(value)
    if candidate.is_absolute():
        return candidate
    repo_candidate = REPO_ROOT / candidate
    return repo_candidate if repo_candidate.exists() else receipt_path.parent / candidate


def audit(args: argparse.Namespace) -> dict[str, object]:
    annotation = json.loads(args.annotation_audit.read_text(encoding="utf-8"))
    inputs = annotation.get("inputs", {})
    events_path = resolve_recorded_path(str(inputs.get("events", "")), args.annotation_audit)
    events_hash_matches = bool(
        events_path.is_file()
        and sha256_file(events_path) == inputs.get("events_sha256")
    )
    annotation_authorized = bool(
        annotation.get("artifact_type") == "expert_annotation_gate_audit"
        and annotation.get("learned_residual_gate", {}).get("authorized") is True
        and annotation.get("event_gate", {}).get("authorized") is True
        and events_hash_matches
    )

    assignments = pd.read_csv(args.assignments, dtype=str).fillna("")
    required = {"event_group_id", "split", "effective_date", "law_family"}
    missing = sorted(required - set(assignments.columns))
    if missing:
        raise ValueError(f"split assignments are missing columns: {missing}")
    duplicate_groups = bool(assignments["event_group_id"].duplicated().any())
    allowed_splits = bool(assignments["split"].isin({"train", "dev"}).all())
    dates = pd.to_datetime(
        assignments["effective_date"], format="%Y-%m-%d", errors="coerce"
    )
    valid_dates = bool(dates.notna().all())
    nonempty_family = bool(assignments["law_family"].str.strip().ne("").all())

    gold_event_groups: set[str] = set()
    if events_path.is_file():
        event_frame = pd.read_csv(events_path, dtype=str).fillna("")
        if {"event_group_id", "verification_status"}.issubset(event_frame.columns):
            gold_event_groups = set(
                event_frame.loc[
                    event_frame["verification_status"].eq("adjudicated_gold"),
                    "event_group_id",
                ]
            )
    assigned_groups = set(assignments["event_group_id"])
    exact_gold_partition = bool(gold_event_groups and assigned_groups == gold_event_groups)

    train = assignments.loc[assignments["split"].eq("train")].copy()
    dev = assignments.loc[assignments["split"].eq("dev")].copy()
    minimum_counts_pass = bool(
        train["event_group_id"].nunique() >= args.minimum_train_events
        and dev["event_group_id"].nunique() >= args.minimum_dev_events
    )
    chronological = False
    if valid_dates and not train.empty and not dev.empty:
        train_dates = pd.to_datetime(train["effective_date"], format="%Y-%m-%d")
        dev_dates = pd.to_datetime(dev["effective_date"], format="%Y-%m-%d")
        chronological = bool(train_dates.max() < dev_dates.min())
    train_families = set(train["law_family"])
    dev_families = set(dev["law_family"])
    family_overlap = sorted(train_families & dev_families)
    family_disjoint = not family_overlap

    authorized = bool(
        annotation_authorized
        and not duplicate_groups
        and allowed_splits
        and valid_dates
        and nonempty_family
        and exact_gold_partition
        and minimum_counts_pass
        and chronological
        and family_disjoint
    )
    receipt: dict[str, object] = {
        "schema_version": 1,
        "artifact_type": "telen_l3_event_split_audit",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "authorized": authorized,
        "inputs": {
            "annotation_audit": str(args.annotation_audit),
            "annotation_audit_sha256": sha256_file(args.annotation_audit),
            "annotation_events": str(events_path),
            "annotation_events_sha256": sha256_file(events_path)
            if events_path.is_file()
            else None,
            "assignments": str(args.assignments),
            "assignments_sha256": sha256_file(args.assignments),
        },
        "configuration": {
            "minimum_train_events": args.minimum_train_events,
            "minimum_dev_events": args.minimum_dev_events,
            "final_test_included": False,
        },
        "counts": {
            "gold_event_groups": len(gold_event_groups),
            "assigned_event_groups": len(assigned_groups),
            "train_event_groups": int(train["event_group_id"].nunique()),
            "development_event_groups": int(dev["event_group_id"].nunique()),
        },
        "gates": {
            "expert_annotation_authorized": annotation_authorized,
            "unique_event_assignment": not duplicate_groups,
            "train_dev_only": allowed_splits,
            "valid_effective_dates": valid_dates,
            "nonempty_law_family": nonempty_family,
            "exact_gold_event_partition": exact_gold_partition,
            "minimum_event_counts": minimum_counts_pass,
            "strict_chronology": chronological,
            "law_family_disjoint": family_disjoint,
            "final_test_sealed": True,
        },
        "law_family_overlap": family_overlap,
        "rule": (
            "authorized expert events AND every gold event assigned exactly once "
            "to train/dev AND max(train effective_date) < min(dev effective_date) "
            "AND no law_family overlap AND minimum event counts pass"
        ),
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(receipt, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return receipt


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--annotation-audit", type=Path, required=True)
    parser.add_argument("--assignments", type=Path, required=True)
    parser.add_argument("--minimum-train-events", type=int, default=5)
    parser.add_argument("--minimum-dev-events", type=int, default=5)
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    result = audit(parse_args())
    print(json.dumps(result, ensure_ascii=False, indent=2))
    if not result["authorized"]:
        raise SystemExit(2)


if __name__ == "__main__":
    main()