#!/usr/bin/env python3
"""Validate an RRC PIST predictions JSON artifact.

Validates rrc_pist_predictions_250_v1.json (and similar prediction files):
  - JSON parses
  - required top-level fields are present
  - every prediction has required fields
  - equation_ids are unique
  - matrices are 8×8 with integer entries
  - matrix hashes are consistent
  - no all-zero matrices (degenerate)
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from pathlib import Path
from typing import Any


REQUIRED_TOP = {"schema", "claim_boundary", "predictions"}
REQUIRED_PRED = {
    "equation_id",
    "object_id",
    "name",
    "equation",
    "schema",
    "claim_boundary",
    "matrix_schema",
    "matrix_hash",
    "matrix_8x8",
    "global_vocab_hash",
}


def canonical_matrix_json(matrix: list[list[int]]) -> str:
    return json.dumps(matrix, separators=(",", ":"))


def validate(data: dict[str, Any]) -> list[str]:
    errors: list[str] = []
    missing_top = REQUIRED_TOP - data.keys()
    if missing_top:
        errors.append(f"missing top-level fields: {sorted(missing_top)}")
        return errors

    preds = data.get("predictions", [])
    if not isinstance(preds, list):
        errors.append("'predictions' must be a list")
        return errors

    ids: set[str] = set()
    hashes: set[str] = set()
    for idx, pred in enumerate(preds):
        missing = REQUIRED_PRED - pred.keys()
        if missing:
            errors.append(f"prediction {idx}: missing fields {sorted(missing)}")
            continue

        eq_id = pred["equation_id"]
        if eq_id in ids:
            errors.append(f"duplicate equation_id: {eq_id}")
        ids.add(eq_id)

        # Matrix shape
        matrix = pred["matrix_8x8"]
        if not isinstance(matrix, list) or len(matrix) != 8:
            errors.append(f"prediction {idx} ({eq_id}): matrix_8x8 must have 8 rows")
            continue
        for row_idx, row in enumerate(matrix):
            if not isinstance(row, list) or len(row) != 8:
                errors.append(
                    f"prediction {idx} ({eq_id}): row {row_idx} must have 8 columns"
                )
            elif not all(isinstance(v, int) for v in row):
                errors.append(
                    f"prediction {idx} ({eq_id}): row {row_idx} has non-integer entries"
                )

        # All-zero matrix check
        if all(all(v == 0 for v in row) for row in matrix):
            errors.append(
                f"prediction {idx} ({eq_id}): degenerate all-zero matrix"
            )

        # Matrix hash consistency
        expected_hash = hashlib.sha256(
            canonical_matrix_json(matrix).encode("utf-8")
        ).hexdigest()
        if pred["matrix_hash"] != expected_hash:
            errors.append(
                f"prediction {idx} ({eq_id}): matrix_hash mismatch "
                f"(expected {expected_hash[:16]}..., got {pred['matrix_hash'][:16]}...)"
            )

        hashes.add(pred["matrix_hash"])

    # Report duplicate matrices
    if len(hashes) < len(preds):
        errors.append(
            f"info: {len(preds) - len(hashes)} duplicate matrices "
            f"({len(hashes)} unique out of {len(preds)})"
        )

    return errors


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Validate RRC PIST predictions JSON"
    )
    parser.add_argument("json_file", help="path to predictions JSON file")
    args = parser.parse_args()

    path = Path(args.json_file)
    if not path.exists():
        print(f"ERROR: file not found: {path}", file=sys.stderr)
        return 1

    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as e:
        print(f"ERROR: invalid JSON: {e}", file=sys.stderr)
        return 1

    errors = validate(data)
    # Separate info messages from real errors
    real_errors = [e for e in errors if not e.startswith("info:")]
    infos = [e for e in errors if e.startswith("info:")]

    for info in infos:
        print(f"INFO: {info}", file=sys.stderr)

    if real_errors:
        print("VALIDATION FAILED", file=sys.stderr)
        for err in real_errors:
            print(f"  - {err}", file=sys.stderr)
        return 1

    preds = data.get("predictions", [])
    print(f"OK: {len(preds)} predictions, schema={data.get('schema')}")
    return 0


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