#!/usr/bin/env python3
"""Audit and canonicalize the pinned Vietnamese legal-document snapshot."""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from collections import Counter
from pathlib import Path

import pandas as pd
import pyarrow.parquet as pq

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

from telen.bootstrap import (  # noqa: E402
    RAW_RELATION_RULES,
    canonicalize_relationship_frame,
)


PRIMARY_RELATIONS = ("REPLACES", "REPEALS", "AMENDS")
TITLE_PATTERNS = {
    "REPLACES": r"\bthay thế\b",
    "REPEALS": r"\bbãi bỏ\b|\bhết hiệu lực\b",
    "AMENDS": r"\bsửa đổi\b|\bbổ sung\b",
}


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


def parse_vietnamese_date(series: pd.Series) -> pd.Series:
    return pd.to_datetime(series, format="%d/%m/%Y", errors="coerce")


def _json_counts(series: pd.Series) -> dict[str, int]:
    values = series.fillna("<NULL>").astype(str).value_counts()
    return {str(key): int(value) for key, value in values.items()}


def _document_url(doc_id: str) -> str:
    return (
        f"https://vbpl.vn/TW/Pages/vbpq-toanvan.aspx?ItemID={doc_id}"
        if str(doc_id).isdigit()
        else ""
    )


def select_discovery_candidates(
    metadata: pd.DataFrame,
    canonical: pd.DataFrame,
    content_ids: set[str],
    *,
    per_relation: int,
    as_of: pd.Timestamp,
) -> pd.DataFrame:
    """Select deterministic, diverse discovery pairs; these are not gold edges."""
    meta = metadata.copy()
    meta["document_date"] = parse_vietnamese_date(meta["ngay_co_hieu_luc"]).fillna(
        parse_vietnamese_date(meta["ngay_ban_hanh"])
    )
    source = meta.add_prefix("event_")
    target = meta.add_prefix("affected_")
    pairs = canonical[canonical["relation"].isin(PRIMARY_RELATIONS)].merge(
        source,
        left_on="event_doc_id",
        right_on="event_id",
        how="inner",
    )
    pairs = pairs.merge(
        target,
        left_on="affected_doc_id",
        right_on="affected_id",
        how="inner",
    )
    pairs["chronology_days"] = (
        pairs["event_document_date"] - pairs["affected_document_date"]
    ).dt.days
    pairs = pairs[
        pairs["event_doc_id"].isin(content_ids)
        & pairs["affected_doc_id"].isin(content_ids)
        & pairs["event_document_date"].notna()
        & pairs["affected_document_date"].notna()
        & pairs["chronology_days"].gt(0)
        & pairs["event_document_date"].le(as_of)
        & pairs["event_pham_vi"].isin({"Trung ương", "Toàn quốc"})
    ].copy()
    pairs["event_year"] = pairs["event_document_date"].dt.year
    pairs = pairs[pairs["event_year"].between(2015, 2026)]
    pairs["year_band"] = pd.cut(
        pairs["event_year"],
        bins=[2014, 2018, 2022, 2026],
        labels=["2015-2018", "2019-2022", "2023-2026"],
        include_lowest=True,
    ).astype("string")
    pairs["domain"] = (
        pairs["event_nganh"]
        .fillna(pairs["event_linh_vuc"])
        .fillna("Chưa phân loại")
        .astype(str)
    )
    pairs["title_signal"] = False
    for relation, pattern in TITLE_PATTERNS.items():
        mask = pairs["relation"].eq(relation)
        pairs.loc[mask, "title_signal"] = pairs.loc[mask, "event_title"].str.contains(
            pattern, case=False, regex=True, na=False
        )
    pairs["target_status_signal"] = pairs["affected_tinh_trang_hieu_luc"].isin(
        {"Hết hiệu lực toàn bộ", "Hết hiệu lực một phần", "Ngưng hiệu lực"}
    )
    pairs = pairs.sort_values(
        [
            "relation",
            "title_signal",
            "target_status_signal",
            "evidence_count",
            "event_document_date",
            "chronology_days",
            "event_doc_id",
        ],
        ascending=[True, False, False, False, False, True, True],
    )

    selected_rows = []
    for relation in PRIMARY_RELATIONS:
        group = pairs[pairs["relation"].eq(relation)]
        domain_counts: Counter[str] = Counter()
        event_ids: set[str] = set()
        base_quota = per_relation // 3
        quotas = {
            "2015-2018": base_quota,
            "2019-2022": base_quota,
            "2023-2026": base_quota + per_relation % 3,
        }

        def consider(row) -> bool:
            if row.event_doc_id in event_ids or domain_counts[row.domain] >= 2:
                return False
            selected_rows.append(row._asdict())
            event_ids.add(row.event_doc_id)
            domain_counts[row.domain] += 1
            return True

        for band, quota in quotas.items():
            added = 0
            for row in group[group["year_band"].eq(band)].itertuples(index=False):
                if consider(row):
                    added += 1
                if added >= quota:
                    break
        if len(event_ids) < per_relation:
            for row in group.itertuples(index=False):
                consider(row)
                if len(event_ids) >= per_relation:
                    break

    selected = pd.DataFrame(selected_rows)
    if selected.empty:
        return selected
    selected.insert(
        0,
        "candidate_id",
        [
            f"{relation.lower()}:{event_id}:{affected_id}"
            for relation, event_id, affected_id in zip(
                selected["relation"],
                selected["event_doc_id"],
                selected["affected_doc_id"],
                strict=True,
            )
        ],
    )
    selected["event_source_url"] = selected["event_doc_id"].map(_document_url)
    selected["affected_source_url"] = selected["affected_doc_id"].map(_document_url)
    selected["verification_status"] = "unverified_discovery"
    keep = [
        "candidate_id",
        "relation",
        "scope",
        "event_doc_id",
        "event_so_ky_hieu",
        "event_title",
        "event_loai_van_ban",
        "event_co_quan_ban_hanh",
        "event_ngay_ban_hanh",
        "event_ngay_co_hieu_luc",
        "year_band",
        "domain",
        "affected_doc_id",
        "affected_so_ky_hieu",
        "affected_title",
        "affected_loai_van_ban",
        "affected_ngay_ban_hanh",
        "affected_ngay_co_hieu_luc",
        "affected_ngay_het_hieu_luc",
        "affected_tinh_trang_hieu_luc",
        "chronology_days",
        "raw_labels",
        "direction_rules",
        "evidence_count",
        "title_signal",
        "target_status_signal",
        "event_source_url",
        "affected_source_url",
        "verification_status",
    ]
    return selected.loc[:, keep]


def audit(args: argparse.Namespace) -> dict:
    metadata_path = args.metadata.resolve()
    relationships_path = args.relationships.resolve()
    content_path = args.content.resolve()
    metadata = pd.read_parquet(metadata_path)
    relationships = pd.read_parquet(relationships_path)
    canonical = canonicalize_relationship_frame(relationships)
    content_ids = set(
        pq.read_table(content_path, columns=["id"])["id"].to_pylist()
    )

    metadata_dates = metadata.copy()
    metadata_dates["document_date"] = parse_vietnamese_date(
        metadata_dates["ngay_co_hieu_luc"]
    ).fillna(parse_vietnamese_date(metadata_dates["ngay_ban_hanh"]))
    event_dates = metadata_dates.loc[:, ["id", "document_date"]].rename(
        columns={"id": "event_doc_id", "document_date": "event_date"}
    )
    affected_dates = metadata_dates.loc[:, ["id", "document_date"]].rename(
        columns={"id": "affected_doc_id", "document_date": "affected_date"}
    )
    chronology = canonical.merge(event_dates, on="event_doc_id", how="left").merge(
        affected_dates, on="affected_doc_id", how="left"
    )
    chronology["delta_days"] = (
        chronology["event_date"] - chronology["affected_date"]
    ).dt.days
    chronology_stats = {}
    for relation, group in chronology.groupby("relation"):
        values = group["delta_days"].dropna()
        chronology_stats[str(relation)] = {
            "edges": int(len(group)),
            "dated_edges": int(len(values)),
            "event_later_fraction": float((values > 0).mean()) if len(values) else None,
            "same_day_fraction": float((values == 0).mean()) if len(values) else None,
            "median_delta_days": float(values.median()) if len(values) else None,
        }

    candidates = select_discovery_candidates(
        metadata,
        canonical,
        content_ids,
        per_relation=args.per_relation,
        as_of=pd.Timestamp(args.as_of),
    )
    args.output_dir.mkdir(parents=True, exist_ok=True)
    canonical_path = args.output_dir / "canonical_relationships.parquet"
    candidates_path = args.output_dir / "discovery_candidates.csv"
    canonical.to_parquet(canonical_path, index=False, compression="zstd")
    candidates.to_csv(candidates_path, index=False)

    mapped_raw = relationships["relationship"].isin(RAW_RELATION_RULES)
    result = {
        "schema_version": 1,
        "artifact_type": "bootstrap_dataset_audit",
        "snapshot": {
            "repository": "th1nhng0/vietnamese-legal-documents",
            "revision": args.revision,
            "metadata_sha256": sha256_file(metadata_path),
            "relationships_sha256": sha256_file(relationships_path),
            "content_sha256": sha256_file(content_path),
        },
        "metadata": {
            "rows": int(len(metadata)),
            "unique_ids": int(metadata["id"].nunique()),
            "content_ids": int(len(content_ids)),
            "metadata_without_content": int((~metadata["id"].isin(content_ids)).sum()),
            "effective_date_non_null": int(metadata["ngay_co_hieu_luc"].notna().sum()),
            "expiry_date_non_null": int(metadata["ngay_het_hieu_luc"].notna().sum()),
            "status_counts": _json_counts(metadata["tinh_trang_hieu_luc"]),
            "scope_counts": _json_counts(metadata["pham_vi"]),
        },
        "relationships": {
            "raw_rows": int(len(relationships)),
            "raw_label_counts": _json_counts(relationships["relationship"]),
            "mapped_raw_rows": int(mapped_raw.sum()),
            "unmapped_raw_rows": int((~mapped_raw).sum()),
            "canonical_rows": int(len(canonical)),
            "canonical_relation_counts": _json_counts(canonical["relation"]),
            "multi_evidence_rows": int(canonical["evidence_count"].gt(1).sum()),
            "chronology": chronology_stats,
            "warning": (
                "Canonical edges are discovery evidence only. Raw active/passive labels "
                "and chronology do not replace official/manual direction verification."
            ),
        },
        "discovery_candidates": {
            "rows": int(len(candidates)),
            "relation_counts": _json_counts(candidates["relation"])
            if not candidates.empty
            else {},
            "path": str(candidates_path),
            "verification_status": "unverified_discovery",
            "as_of": args.as_of,
        },
        "generated": {
            "canonical_relationships": str(canonical_path),
            "discovery_candidates": str(candidates_path),
        },
    }
    output_json = args.output_dir / "audit.json"
    output_json.write_text(
        json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return result


def parse_args() -> argparse.Namespace:
    default_root = REPO_ROOT / "tmp/deep-research/telen/vn-legal-data/data"
    parser = argparse.ArgumentParser()
    parser.add_argument("--metadata", type=Path, default=default_root / "metadata.parquet")
    parser.add_argument(
        "--relationships", type=Path, default=default_root / "relationships.parquet"
    )
    parser.add_argument("--content", type=Path, default=default_root / "content.parquet")
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=REPO_ROOT / "tmp/telen/bootstrap-audit",
    )
    parser.add_argument(
        "--revision", default="8977887f17be2defae4c5171d55562e1cde7d695"
    )
    parser.add_argument("--per-relation", type=int, default=10)
    parser.add_argument("--as-of", default="2026-08-29")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.per_relation <= 0:
        raise ValueError("--per-relation must be positive")
    result = audit(args)
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()