"""Validate and compare DMHY annotation pipeline interchange JSONL.

This tool is intentionally independent from the prefix-tree and DAG runners.
It defines the shared annotation-unit and annotation-patch contract that both
pipelines can emit while they evolve separately.
"""

from __future__ import annotations

import argparse
import json
import sys
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable


SCHEMA_VERSION = "dmhy-annotation-v1"
SOURCE_KINDS = {"prefix_tree", "prefix_dag"}
PATCH_STATUSES = {"ok", "fallback", "failed"}
ANNOTATION_ARRAY_FIELDS = (
    "episode_title_suffixes",
    "media_suffixes",
    "title_candidates",
)


@dataclass(frozen=True)
class JsonlRow:
    path: Path
    line_number: int
    value: Any


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Validate and compare DMHY annotation-unit and patch JSONL files"
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    units = subparsers.add_parser("validate-units", help="Validate annotation-unit JSONL")
    units.add_argument("units_jsonl", type=Path)
    units.add_argument("--manifest-output", type=Path)

    patches = subparsers.add_parser("validate-patches", help="Validate annotation-patch JSONL")
    patches.add_argument("patches_jsonl", type=Path)
    patches.add_argument("--units", type=Path, help="Optional units JSONL for terminal_id mismatch checks")
    patches.add_argument("--manifest-output", type=Path)

    compare = subparsers.add_parser("compare-patches", help="Compare two annotation-patch JSONL files")
    compare.add_argument("left_patches_jsonl", type=Path)
    compare.add_argument("right_patches_jsonl", type=Path)
    compare.add_argument("--left-label", default="left")
    compare.add_argument("--right-label", default="right")
    compare.add_argument("--manifest-output", type=Path)
    compare.add_argument("--summary-limit", type=int, default=8)

    return parser.parse_args()


def iter_jsonl(path: Path) -> Iterable[JsonlRow]:
    if not path.exists():
        raise SystemExit(f"JSONL not found: {path}")
    with path.open("r", encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            try:
                value = json.loads(line)
            except json.JSONDecodeError as exc:
                yield JsonlRow(path, line_number, {"__json_error__": str(exc)})
                continue
            yield JsonlRow(path, line_number, value)


def error_key(message: str) -> str:
    return message.split(":", 1)[0]


def add_error(errors: list[str], field_errors: Counter[str], message: str) -> None:
    errors.append(message)
    field_errors[error_key(message)] += 1


def is_string_list(value: Any) -> bool:
    return isinstance(value, list) and all(isinstance(item, str) for item in value)


def is_int_not_bool(value: Any) -> bool:
    return isinstance(value, int) and not isinstance(value, bool)


def validate_unit_row(row: Any) -> tuple[list[str], str | None, list[str]]:
    field_errors: Counter[str] = Counter()
    errors: list[str] = []
    if not isinstance(row, dict):
        add_error(errors, field_errors, "row_type: expected object")
        return errors, None, []

    unit_id = row.get("unit_id")
    terminal_ids = row.get("terminal_ids")
    if not isinstance(unit_id, str) or not unit_id:
        add_error(errors, field_errors, "unit_id: expected non-empty string")
    if row.get("source_kind") not in SOURCE_KINDS:
        add_error(errors, field_errors, "source_kind: expected prefix_tree or prefix_dag")
    source_id = row.get("source_id")
    if not isinstance(source_id, (str, int)) or isinstance(source_id, bool):
        add_error(errors, field_errors, "source_id: expected string or integer")
    if not is_string_list(terminal_ids):
        add_error(errors, field_errors, "terminal_ids: expected string array")
        clean_terminal_ids: list[str] = []
    else:
        clean_terminal_ids = terminal_ids
    if not is_int_not_bool(row.get("weight")):
        add_error(errors, field_errors, "weight: expected integer")

    context = row.get("context")
    if not isinstance(context, dict):
        add_error(errors, field_errors, "context: expected object")
    else:
        for key in ("prefixes", "digit_skeletons", "edge_labels"):
            if not is_string_list(context.get(key)):
                add_error(errors, field_errors, f"context.{key}: expected string array")
        if context.get("notes") is not None and not isinstance(context.get("notes"), str):
            add_error(errors, field_errors, "context.notes: expected string or null")

    examples = row.get("examples")
    if not isinstance(examples, dict):
        add_error(errors, field_errors, "examples: expected object")
    else:
        for key in ("values", "suffixes"):
            if not is_string_list(examples.get(key)):
                add_error(errors, field_errors, f"examples.{key}: expected string array")

    expected_output = row.get("expected_output")
    if not isinstance(expected_output, dict):
        add_error(errors, field_errors, "expected_output: expected object")
    elif expected_output.get("schema_version") != SCHEMA_VERSION:
        add_error(errors, field_errors, f"expected_output.schema_version: expected {SCHEMA_VERSION}")

    return errors, unit_id if isinstance(unit_id, str) else None, clean_terminal_ids


def normalize_patch_row(row: Any) -> dict[str, Any]:
    """Return the canonical patch view, accepting legacy flat patch rows too."""
    if not isinstance(row, dict):
        return {}
    annotation = row.get("annotation")
    if isinstance(annotation, dict):
        annotation_obj = annotation
    else:
        annotation_obj = {
            key: row.get(key)
            for key in (*ANNOTATION_ARRAY_FIELDS, "llm_label", "notes")
            if key in row
        }
    status = row.get("status")
    if status is None:
        source = str(row.get("source") or "")
        if row.get("fallback") is True:
            status = "fallback"
        elif source.startswith("responses"):
            status = "ok"
        else:
            status = "ok"
    errors = row.get("errors")
    if errors is None:
        errors = []
    return {
        "unit_id": row.get("unit_id", row.get("terminal_id")),
        "terminal_ids": row.get("terminal_ids", [row["terminal_id"]] if "terminal_id" in row else None),
        "annotation": annotation_obj,
        "status": status,
        "errors": errors,
    }


def validate_patch_row(row: Any) -> tuple[list[str], str | None, list[str], dict[str, Any]]:
    field_errors: Counter[str] = Counter()
    errors: list[str] = []
    canonical = normalize_patch_row(row)
    if not isinstance(row, dict):
        add_error(errors, field_errors, "row_type: expected object")
        return errors, None, [], canonical

    unit_id = canonical.get("unit_id")
    terminal_ids = canonical.get("terminal_ids")
    annotation = canonical.get("annotation")
    if not isinstance(unit_id, str) or not unit_id:
        add_error(errors, field_errors, "unit_id: expected non-empty string")
    if not is_string_list(terminal_ids):
        add_error(errors, field_errors, "terminal_ids: expected string array")
        clean_terminal_ids: list[str] = []
    else:
        clean_terminal_ids = terminal_ids
    if not isinstance(annotation, dict):
        add_error(errors, field_errors, "annotation: expected object")
    else:
        for key in ANNOTATION_ARRAY_FIELDS:
            if not is_string_list(annotation.get(key)):
                add_error(errors, field_errors, f"annotation.{key}: expected string array")
        if annotation.get("llm_label") is not None and not isinstance(annotation.get("llm_label"), str):
            add_error(errors, field_errors, "annotation.llm_label: expected string or null")
        if annotation.get("notes") is not None and not isinstance(annotation.get("notes"), str):
            add_error(errors, field_errors, "annotation.notes: expected string or null")
    if canonical.get("status") not in PATCH_STATUSES:
        add_error(errors, field_errors, "status: expected ok, fallback, or failed")
    if not is_string_list(canonical.get("errors")):
        add_error(errors, field_errors, "errors: expected string array")
    return errors, unit_id if isinstance(unit_id, str) else None, clean_terminal_ids, canonical


def terminal_set(ids: Iterable[str]) -> set[str]:
    return {str(item) for item in ids if str(item)}


def unit_index(path: Path) -> tuple[dict[str, set[str]], dict[str, Any]]:
    units: dict[str, set[str]] = {}
    report = validate_units(path)
    for item in report["valid_units"]:
        units[item["unit_id"]] = set(item["terminal_ids"])
    return units, report


def summarize_report(report: dict[str, Any], manifest_output: Path | None) -> None:
    printable = {key: value for key, value in report.items() if not key.startswith("valid_")}
    text = json.dumps(printable, ensure_ascii=False, indent=2, sort_keys=True)
    print(text)
    if manifest_output is not None:
        manifest_output.parent.mkdir(parents=True, exist_ok=True)
        manifest_output.write_text(text + "\n", encoding="utf-8")


def validate_units(path: Path) -> dict[str, Any]:
    field_errors: Counter[str] = Counter()
    duplicate_unit_ids = 0
    seen_unit_ids: set[str] = set()
    valid_units: list[dict[str, Any]] = []
    total = 0
    error_rows = 0
    for jsonl_row in iter_jsonl(path):
        total += 1
        row = jsonl_row.value
        if isinstance(row, dict) and "__json_error__" in row:
            field_errors["json"] += 1
            error_rows += 1
            continue
        errors, unit_id, terminal_ids = validate_unit_row(row)
        for message in errors:
            field_errors[error_key(message)] += 1
        if unit_id:
            if unit_id in seen_unit_ids:
                duplicate_unit_ids += 1
                field_errors["unit_id.duplicate"] += 1
                errors.append("unit_id.duplicate")
            seen_unit_ids.add(unit_id)
        if errors:
            error_rows += 1
        else:
            valid_units.append({"unit_id": unit_id, "terminal_ids": terminal_ids})
    return {
        "schema": "dmhy_annotation_units",
        "schema_version": SCHEMA_VERSION,
        "path": str(path),
        "total_rows": total,
        "error_rows": error_rows,
        "error_count": sum(field_errors.values()),
        "common_field_errors": dict(field_errors.most_common()),
        "duplicate_unit_ids": duplicate_unit_ids,
        "terminal_coverage_count": len({tid for unit in valid_units for tid in unit["terminal_ids"]}),
        "valid_units": valid_units,
    }


def validate_patches(path: Path, units_path: Path | None = None) -> dict[str, Any]:
    field_errors: Counter[str] = Counter()
    status_counts: Counter[str] = Counter()
    duplicate_unit_ids = 0
    seen_unit_ids: set[str] = set()
    valid_patches: list[dict[str, Any]] = []
    total = 0
    error_rows = 0
    terminal_mismatch_count = 0
    unknown_unit_count = 0
    units_by_id: dict[str, set[str]] = {}
    units_report: dict[str, Any] | None = None
    if units_path is not None:
        units_by_id, units_report = unit_index(units_path)

    for jsonl_row in iter_jsonl(path):
        total += 1
        row = jsonl_row.value
        if isinstance(row, dict) and "__json_error__" in row:
            field_errors["json"] += 1
            error_rows += 1
            continue
        errors, unit_id, terminal_ids, canonical = validate_patch_row(row)
        for message in errors:
            field_errors[error_key(message)] += 1
        if unit_id:
            if unit_id in seen_unit_ids:
                duplicate_unit_ids += 1
                field_errors["unit_id.duplicate"] += 1
                errors.append("unit_id.duplicate")
            seen_unit_ids.add(unit_id)
            if units_by_id:
                expected = units_by_id.get(unit_id)
                if expected is None:
                    unknown_unit_count += 1
                    field_errors["unit_id.unknown"] += 1
                    errors.append("unit_id.unknown")
                elif set(terminal_ids) != expected:
                    terminal_mismatch_count += 1
                    field_errors["terminal_ids.mismatch"] += 1
                    errors.append("terminal_ids.mismatch")
        if isinstance(canonical.get("status"), str):
            status_counts[str(canonical["status"])] += 1
        if errors:
            error_rows += 1
        else:
            valid_patches.append(canonical)

    return {
        "schema": "dmhy_annotation_patches",
        "schema_version": SCHEMA_VERSION,
        "path": str(path),
        "units_path": str(units_path) if units_path else None,
        "total_rows": total,
        "error_rows": error_rows,
        "error_count": sum(field_errors.values()),
        "common_field_errors": dict(field_errors.most_common()),
        "duplicate_unit_ids": duplicate_unit_ids,
        "terminal_id_mismatch_count": terminal_mismatch_count,
        "unknown_unit_count": unknown_unit_count,
        "status_counts": {status: status_counts.get(status, 0) for status in sorted(PATCH_STATUSES)},
        "terminal_coverage_count": len({tid for patch in valid_patches for tid in patch["terminal_ids"]}),
        "units_validation": None
        if units_report is None
        else {key: value for key, value in units_report.items() if not key.startswith("valid_")},
        "valid_patches": valid_patches,
    }


def annotation_empty_rates(patches: list[dict[str, Any]]) -> dict[str, float]:
    if not patches:
        return {key: 0.0 for key in (*ANNOTATION_ARRAY_FIELDS, "all_annotation_arrays")}
    counts: Counter[str] = Counter()
    for patch in patches:
        annotation = patch.get("annotation") or {}
        all_empty = True
        for key in ANNOTATION_ARRAY_FIELDS:
            empty = not annotation.get(key)
            counts[key] += int(empty)
            all_empty = all_empty and empty
        counts["all_annotation_arrays"] += int(all_empty)
    return {key: round(counts[key] / len(patches), 6) for key in (*ANNOTATION_ARRAY_FIELDS, "all_annotation_arrays")}


def text_summary(patches: list[dict[str, Any]], limit: int) -> dict[str, list[dict[str, Any]]]:
    notes: Counter[str] = Counter()
    errors: Counter[str] = Counter()
    for patch in patches:
        annotation = patch.get("annotation") or {}
        note = annotation.get("notes")
        if isinstance(note, str) and note.strip():
            notes[note.strip()[:240]] += 1
        for error in patch.get("errors") or []:
            if isinstance(error, str) and error.strip():
                errors[error.strip()[:240]] += 1
    return {
        "notes": [{"text": text, "count": count} for text, count in notes.most_common(limit)],
        "errors": [{"text": text, "count": count} for text, count in errors.most_common(limit)],
    }


def compact_patch_stats(path: Path, label: str, summary_limit: int) -> dict[str, Any]:
    report = validate_patches(path)
    patches = report["valid_patches"]
    terminals = {tid for patch in patches for tid in patch["terminal_ids"]}
    statuses = Counter(str(patch.get("status")) for patch in patches)
    return {
        "label": label,
        "path": str(path),
        "validation": {key: value for key, value in report.items() if not key.startswith("valid_")},
        "status_counts": {status: statuses.get(status, 0) for status in sorted(PATCH_STATUSES)},
        "terminal_coverage_count": len(terminals),
        "terminals": terminals,
        "annotation_empty_rates": annotation_empty_rates(patches),
        "text_summary": text_summary(patches, summary_limit),
    }


def compare_patches(
    left_path: Path,
    right_path: Path,
    left_label: str,
    right_label: str,
    summary_limit: int,
) -> dict[str, Any]:
    left = compact_patch_stats(left_path, left_label, summary_limit)
    right = compact_patch_stats(right_path, right_label, summary_limit)
    left_terminals = left.pop("terminals")
    right_terminals = right.pop("terminals")
    return {
        "schema": "dmhy_annotation_patch_comparison",
        "schema_version": SCHEMA_VERSION,
        "left": left,
        "right": right,
        "terminal_intersection_count": len(left_terminals & right_terminals),
        "left_only_terminal_count": len(left_terminals - right_terminals),
        "right_only_terminal_count": len(right_terminals - left_terminals),
        "terminal_union_count": len(left_terminals | right_terminals),
    }


def main() -> None:
    args = parse_args()
    if args.command == "validate-units":
        report = validate_units(args.units_jsonl)
        summarize_report(report, args.manifest_output)
        raise SystemExit(1 if report["error_count"] else 0)
    if args.command == "validate-patches":
        report = validate_patches(args.patches_jsonl, args.units)
        summarize_report(report, args.manifest_output)
        raise SystemExit(1 if report["error_count"] else 0)
    if args.command == "compare-patches":
        report = compare_patches(
            args.left_patches_jsonl,
            args.right_patches_jsonl,
            args.left_label,
            args.right_label,
            max(1, args.summary_limit),
        )
        summarize_report(report, args.manifest_output)
        left_errors = report["left"]["validation"]["error_count"]
        right_errors = report["right"]["validation"]["error_count"]
        raise SystemExit(1 if left_errors or right_errors else 0)
    raise SystemExit(f"unknown command: {args.command}")


if __name__ == "__main__":
    try:
        main()
    except BrokenPipeError:
        sys.exit(1)