"""Math Ink 0.6 P-track trajectory federation의 권리·중복·coverage를 감사한다."""

from __future__ import annotations

import argparse
from hashlib import sha256
import json
import sys
from pathlib import Path

import torch

PROJECT_ROOT = Path(__file__).parents[1]
SOURCE_ROOT = PROJECT_ROOT / "src"
if str(SOURCE_ROOT) not in sys.path:
    sys.path.insert(0, str(SOURCE_ROOT))

from math_grid_drawer.research.ink06_federation import federation_audit06, load_product_federation06
from math_grid_drawer.research.ink06_source_registry import (
    SourceRegistryEntry06,
    load_source_registry06,
    source_registry_audit06,
)


def _file_sha25606(path: Path) -> str:
    """필요 변수: 감사 대상 checkpoint. 작동 원리: 최종 model bundle과 연결할 byte-level SHA-256을 계산한다."""

    digest = sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def build_federation_release_gate06(
    *,
    checkpoint: dict,
    federation: dict,
    registry_entries: tuple[SourceRegistryEntry06, ...],
    registry_audit: dict,
) -> dict:
    """필요 변수: checkpoint provenance·실데이터 audit·registry. 작동 원리: 200/30·중복·split 조건을 AND로 검증한다."""

    declared = tuple(
        sorted(set(str(value) for value in checkpoint.get("training_source_ids", [])))
    )
    by_id = {entry.source_id: entry for entry in registry_entries}
    unknown = sorted(source_id for source_id in declared if source_id not in by_id)
    ineligible = sorted(
        source_id
        for source_id in declared
        if source_id in by_id
        and (
            by_id[source_id].stage != "approved"
            or by_id[source_id].deployment_role == "evaluation_only"
        )
    )
    groups = {
        by_id[source_id].independent_source_group
        for source_id in declared
        if source_id in by_id
        and by_id[source_id].stage == "approved"
        and by_id[source_id].deployment_role != "evaluation_only"
        and by_id[source_id].independent_source_group
    }
    checks = {
        "discovered_200": bool(registry_audit["discovery_gate_passed"]),
        "approved_independent_30": bool(registry_audit["approved_group_gate_passed"]),
        "checkpoint_declares_30_groups": len(groups) >= 30,
        "checkpoint_sources_approved": not unknown and not ineligible,
        "origin_overlap_zero": int(federation["origin_overlap_total"]) == 0,
        "trajectory_signature_overlap_zero": int(
            federation["trajectory_signature_overlap_total"],
        ) == 0,
        "writer_device_origin_split_leakage_zero": int(
            federation["split_identity_leakage_total"],
        ) == 0,
    }
    return {
        "required_discovered_sources": 200,
        "required_independent_groups": 30,
        "checkpoint_training_source_ids": list(declared),
        "checkpoint_training_source_count": len(declared),
        "checkpoint_independent_groups": len(groups),
        "unknown_checkpoint_sources": unknown,
        "ineligible_checkpoint_sources": ineligible,
        "checks": checks,
        "passed": all(checks.values()),
    }


def main() -> None:
    """필요 변수: 0.6 checkpoint·registry·shard. 작동 원리: fail-closed source audit를 UTF-8 JSON으로 고정한다."""

    parser = argparse.ArgumentParser(description="Audit Math Ink 0.6 product trajectory federation")
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--registry", type=Path, default=PROJECT_ROOT / "research/dataset_registry.json")
    parser.add_argument("--source-registry", type=Path, default=PROJECT_ROOT / "research/math_ink_06_source_registry.json")
    parser.add_argument("--commercial", type=Path, default=PROJECT_ROOT / "research/data/external_trajectory_v1/commercial_ccby4.jsonl.gz")
    parser.add_argument("--hwrt", type=Path, default=PROJECT_ROOT / "research/data/open_pretrain/hwrt_expanded_v2/hwrt_expanded.jsonl.gz")
    parser.add_argument("--approval", type=Path, default=PROJECT_ROOT / "research/approvals/HWRT-ODBL-USE-APPROVAL-v1.json")
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    sources = load_product_federation06(
        registry_path=args.registry, commercial_path=args.commercial, hwrt_path=args.hwrt,
        approval_path=args.approval, allowed_labels=checkpoint["exact_labels"], source_registry_path=args.source_registry,
    )
    report = federation_audit06(sources)
    registry_entries = load_source_registry06(args.source_registry)
    discovery = source_registry_audit06(
        registry_entries,
        required_discovered_sources=200,
        required_approved_groups=30,
    )
    release_gate = build_federation_release_gate06(
        checkpoint=checkpoint,
        federation=report,
        registry_entries=registry_entries,
        registry_audit=discovery,
    )
    report.update({
        "schema": "aiflow-math-ink-06-federation-audit-v1",
        "checkpoint_sha256": _file_sha25606(args.checkpoint),
        "vocabulary": len(checkpoint["exact_labels"]),
        "source_registry": discovery,
        "release_source_gate": release_gate,
        "product_validation": False,
        "note": "HWRT official test는 모델 선택·writer-independent 제품 gate에 사용하지 않는다.",
    })
    args.output.mkdir(parents=True, exist_ok=True)
    (args.output / "federation_audit.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8",
    )
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()