"""审计固定双人名义分类标注的一致率、Cohen's kappa 与冲突。

这是一个教学用的离线合同检查器：它审计一批已经分派给两名标注者的
初标记录，不会替代裁决、gold 审查、抽样代表性或隐私/许可审查。
"""

from __future__ import annotations

import argparse
from collections import Counter, defaultdict
from dataclasses import dataclass
from datetime import datetime
import json
from pathlib import Path
import re
from typing import Any, Sequence


LABELS = (
    "helpful",
    "not_helpful",
    "unsafe",
    "cannot_judge",
    "exclude",
)
REQUIRED_FIELDS = {
    "annotation_id",
    "sample_id",
    "source_revision",
    "data_version",
    "guideline_version",
    "label_set_version",
    "task_config_version",
    "annotator",
    "label",
    "evidence",
    "created_at",
}
UTC_SECOND_TIMESTAMP = re.compile(
    r"^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}Z$"
)


@dataclass(frozen=True)
class AnnotationPair:
    sample_id: str
    label_a: str
    label_b: str
    source_revision: str = "manual"


@dataclass(frozen=True)
class AnnotationBatch:
    data_version: str
    guideline_version: str
    label_set_version: str
    task_config_version: str
    annotators: tuple[str, str]
    pairs: tuple[AnnotationPair, ...]


@dataclass(frozen=True)
class AgreementResult:
    observed: float
    expected: float
    kappa: float | None
    confusion: tuple[tuple[str, str, int], ...]
    conflicts: tuple[AnnotationPair, ...]


def _object_without_duplicate_keys(
    pairs: list[tuple[str, Any]],
) -> dict[str, Any]:
    record: dict[str, Any] = {}
    for key, value in pairs:
        if key in record:
            raise ValueError(f"duplicate JSON key: {key}")
        record[key] = value
    return record


def _reject_non_standard_constant(value: str) -> None:
    raise ValueError(f"non-standard JSON constant: {value}")


def _validate_timestamp(value: str, line_number: int) -> None:
    """只接受本教学合同声明的 UTC 秒级时间，避免模糊的本地时间。"""

    if not UTC_SECOND_TIMESTAMP.fullmatch(value):
        raise ValueError(
            f"line {line_number}: created_at must use YYYY-MM-DDTHH:MM:SSZ"
        )
    try:
        datetime.strptime(value, "%Y-%m-%dT%H:%M:%SZ")
    except ValueError as exc:
        raise ValueError(
            f"line {line_number}: created_at is not a calendar timestamp"
        ) from exc


def _read_record(line: str, line_number: int) -> dict[str, str]:
    if not line.strip():
        raise ValueError(f"line {line_number}: blank JSONL record")
    try:
        parsed = json.loads(
            line,
            object_pairs_hook=_object_without_duplicate_keys,
            parse_constant=_reject_non_standard_constant,
        )
    except (json.JSONDecodeError, ValueError) as exc:
        raise ValueError(f"line {line_number}: invalid JSON: {exc}") from exc
    if not isinstance(parsed, dict):
        raise ValueError(f"line {line_number}: record must be a JSON object")

    missing = sorted(REQUIRED_FIELDS - parsed.keys())
    if missing:
        raise ValueError(
            f"line {line_number}: missing fields: {','.join(missing)}"
        )
    unknown = sorted(parsed.keys() - REQUIRED_FIELDS)
    if unknown:
        raise ValueError(
            f"line {line_number}: unknown fields: {','.join(unknown)}"
        )

    record: dict[str, str] = {}
    for field in REQUIRED_FIELDS:
        value = parsed[field]
        if not isinstance(value, str) or not value.strip():
            raise ValueError(
                f"line {line_number}: {field} must be a non-empty string"
            )
        record[field] = value.strip()
    if record["label"] not in LABELS:
        raise ValueError(
            f"line {line_number}: unsupported label {record['label']!r}"
        )
    _validate_timestamp(record["created_at"], line_number)
    return record


def _single_batch_value(values: set[str], field: str) -> str:
    if len(values) != 1:
        raise ValueError(f"batch mixes multiple {field} values")
    return next(iter(values))


def load_batch(path: Path) -> AnnotationBatch:
    """加载严格 JSONL，要求固定双标、统一合同版本和同一输入快照。"""

    if not path.is_file():
        raise ValueError("annotation path must be an existing file")

    grouped: dict[str, dict[str, dict[str, str]]] = defaultdict(dict)
    annotation_ids: set[str] = set()
    annotators: set[str] = set()
    data_versions: set[str] = set()
    guideline_versions: set[str] = set()
    label_set_versions: set[str] = set()
    task_config_versions: set[str] = set()
    with path.open("r", encoding="utf-8-sig") as source:
        for line_number, line in enumerate(source, start=1):
            record = _read_record(line, line_number)
            annotation_id = record["annotation_id"]
            if annotation_id in annotation_ids:
                raise ValueError(
                    f"line {line_number}: duplicate annotation_id {annotation_id!r}"
                )
            annotation_ids.add(annotation_id)

            sample_id = record["sample_id"]
            annotator = record["annotator"]
            if annotator in grouped[sample_id]:
                raise ValueError(
                    f"line {line_number}: duplicate annotation for "
                    f"{sample_id!r} by {annotator!r}"
                )
            grouped[sample_id][annotator] = record
            annotators.add(annotator)
            data_versions.add(record["data_version"])
            guideline_versions.add(record["guideline_version"])
            label_set_versions.add(record["label_set_version"])
            task_config_versions.add(record["task_config_version"])

    if not grouped:
        raise ValueError("annotation file contains no records")
    if len(annotators) != 2:
        raise ValueError(
            f"batch must contain exactly two annotators, found {len(annotators)}"
        )

    ordered_annotators = tuple(sorted(annotators))
    annotator_a, annotator_b = ordered_annotators
    pairs: list[AnnotationPair] = []
    for sample_id, annotations in sorted(grouped.items()):
        if set(annotations) != annotators:
            missing = sorted(annotators - set(annotations))
            raise ValueError(
                f"{sample_id}: missing annotations from {','.join(missing)}"
            )
        source_revisions = {
            record["source_revision"] for record in annotations.values()
        }
        if len(source_revisions) != 1:
            raise ValueError(
                f"{sample_id}: annotators did not use the same source_revision"
            )
        pairs.append(
            AnnotationPair(
                sample_id=sample_id,
                label_a=annotations[annotator_a]["label"],
                label_b=annotations[annotator_b]["label"],
                source_revision=next(iter(source_revisions)),
            )
        )

    return AnnotationBatch(
        data_version=_single_batch_value(data_versions, "data_version"),
        guideline_version=_single_batch_value(
            guideline_versions, "guideline_version"
        ),
        label_set_version=_single_batch_value(
            label_set_versions, "label_set_version"
        ),
        task_config_version=_single_batch_value(
            task_config_versions, "task_config_version"
        ),
        annotators=(annotator_a, annotator_b),
        pairs=tuple(pairs),
    )


def agreement_and_kappa(
    pairs: Sequence[AnnotationPair],
) -> AgreementResult:
    """计算名义标签的一致率、期望一致率、kappa 和混淆计数。"""

    if not pairs:
        raise ValueError("no annotation pairs")
    seen_samples: set[str] = set()
    for pair in pairs:
        if (
            not pair.sample_id
            or not pair.source_revision
            or pair.sample_id in seen_samples
        ):
            raise ValueError(
                "pairs must contain unique non-empty sample_id and source_revision"
            )
        if pair.label_a not in LABELS or pair.label_b not in LABELS:
            raise ValueError("pair contains an unsupported label")
        seen_samples.add(pair.sample_id)
    labels_a = Counter(pair.label_a for pair in pairs)
    labels_b = Counter(pair.label_b for pair in pairs)
    labels = set(labels_a) | set(labels_b)
    count = len(pairs)
    observed = sum(pair.label_a == pair.label_b for pair in pairs) / count
    expected = sum(
        (labels_a[label] / count) * (labels_b[label] / count)
        for label in labels
    )
    denominator = 1.0 - expected
    kappa = (
        None
        if abs(denominator) < 1e-15
        else (observed - expected) / denominator
    )

    confusion_counter = Counter(
        (pair.label_a, pair.label_b)
        for pair in pairs
    )
    confusion = tuple(
        (label_a, label_b, confusion_counter[(label_a, label_b)])
        for label_a in LABELS
        for label_b in LABELS
        if confusion_counter[(label_a, label_b)]
    )
    conflicts = tuple(pair for pair in pairs if pair.label_a != pair.label_b)
    return AgreementResult(
        observed=observed,
        expected=expected,
        kappa=kappa,
        confusion=confusion,
        conflicts=conflicts,
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="审计恰好两名标注者的名义分类 JSONL",
    )
    parser.add_argument("path", type=Path)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    batch = load_batch(args.path)
    result = agreement_and_kappa(batch.pairs)
    annotator_a, annotator_b = batch.annotators
    kappa_text = "undefined" if result.kappa is None else f"{result.kappa:.3f}"
    print(
        f"pairs={len(batch.pairs)} annotators={annotator_a},{annotator_b} "
        f"data_version={batch.data_version} "
        f"guideline_version={batch.guideline_version} "
        f"label_set_version={batch.label_set_version} "
        f"task_config_version={batch.task_config_version}"
    )
    print(
        f"observed={result.observed:.3f} "
        f"expected={result.expected:.3f} kappa={kappa_text}"
    )
    print(f"confusion rows={annotator_a} columns={annotator_b}:")
    for label_a, label_b, count in result.confusion:
        print(f"- {label_a} -> {label_b}: {count}")
    print("conflicts:")
    if not result.conflicts:
        print("- none")
    for pair in result.conflicts:
        print(
            f"- {pair.sample_id} ({pair.source_revision}): "
            f"{pair.label_a} <> {pair.label_b}"
        )
    return 0


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