#!/usr/bin/env python3
"""Independent verifier for the dimensional KRK tablebase certificate.

This module imports nothing from dimensional_chess_solver.py. It independently
reconstructs legality, D4 orbits, transitions, terminal conditions, Bellman
values, distance-to-mate equations, policies, and the concrete principal line.
"""

from __future__ import annotations

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


WHITE = 0
BLACK = 1
CAPTURE_DRAW = -1


class VerificationError(RuntimeError):
    pass


@dataclass(frozen=True, slots=True)
class Position:
    white_king: int
    white_rook: int
    black_king: int
    side: int


def pack(position: Position) -> int:
    return (((position.white_king * 64 + position.white_rook) * 64 + position.black_king) * 2 + position.side)


def unpack(number: int) -> Position:
    side = number & 1
    number >>= 1
    black_king = number & 63
    number >>= 6
    white_rook = number & 63
    white_king = number >> 6
    return Position(white_king, white_rook, black_king, side)


def image_of_square(square: int, operation: int) -> int:
    x, y = square & 7, square >> 3
    if operation == 0:
        a, b = x, y
    elif operation == 1:
        a, b = 7 - x, y
    elif operation == 2:
        a, b = x, 7 - y
    elif operation == 3:
        a, b = 7 - x, 7 - y
    elif operation == 4:
        a, b = y, x
    elif operation == 5:
        a, b = 7 - y, x
    elif operation == 6:
        a, b = y, 7 - x
    elif operation == 7:
        a, b = 7 - y, 7 - x
    else:
        raise VerificationError("invalid symmetry")
    return (b << 3) | a


def image(position: Position, operation: int) -> Position:
    return Position(
        image_of_square(position.white_king, operation),
        image_of_square(position.white_rook, operation),
        image_of_square(position.black_king, operation),
        position.side,
    )


def orbit_key(position: Position) -> int:
    return min(pack(image(position, operation)) for operation in range(8))


def touching(first: int, second: int) -> bool:
    return max(abs((first & 7) - (second & 7)), abs((first >> 3) - (second >> 3))) <= 1


def rook_line(rook: int, target: int, obstruction: int) -> bool:
    rx, ry = rook & 7, rook >> 3
    tx, ty = target & 7, target >> 3
    if rx == tx:
        increment = 8 if ty > ry else -8
    elif ry == ty:
        increment = 1 if tx > rx else -1
    else:
        return False
    probe = rook + increment
    while probe != target:
        if probe == obstruction:
            return False
        probe += increment
    return True


def admissible(position: Position) -> bool:
    pieces = (position.white_king, position.white_rook, position.black_king)
    if position.side not in (WHITE, BLACK) or any(not 0 <= value < 64 for value in pieces):
        return False
    if len(set(pieces)) != 3 or touching(position.white_king, position.black_king):
        return False
    if position.side == WHITE and rook_line(position.white_rook, position.black_king, position.white_king):
        return False
    return True


def king_steps(origin: int) -> Iterable[int]:
    x, y = origin & 7, origin >> 3
    for dx in (-1, 0, 1):
        for dy in (-1, 0, 1):
            if dx == 0 and dy == 0:
                continue
            a, b = x + dx, y + dy
            if 0 <= a < 8 and 0 <= b < 8:
                yield (b << 3) | a


def rook_steps(position: Position) -> Iterable[int]:
    for increment in (1, -1, 8, -8):
        current = position.white_rook
        while True:
            destination = current + increment
            if not 0 <= destination < 64:
                break
            if increment in (1, -1) and (destination >> 3) != (current >> 3):
                break
            if destination in (position.white_king, position.black_king):
                break
            yield destination
            current = destination


def concrete_children(position: Position) -> tuple[Position | None, ...]:
    children: list[Position | None] = []
    if position.side == WHITE:
        for destination in king_steps(position.white_king):
            if destination != position.white_rook and not touching(destination, position.black_king):
                children.append(Position(destination, position.white_rook, position.black_king, BLACK))
        for destination in rook_steps(position):
            children.append(Position(position.white_king, destination, position.black_king, BLACK))
    else:
        for destination in king_steps(position.black_king):
            if destination == position.white_king or touching(position.white_king, destination):
                continue
            if destination == position.white_rook:
                children.append(None)
            elif not rook_line(position.white_rook, destination, position.white_king):
                children.append(Position(position.white_king, position.white_rook, destination, WHITE))
    return tuple(children)


def child_keys(position: Position) -> tuple[int, ...]:
    return tuple(sorted({CAPTURE_DRAW if child is None else orbit_key(child) for child in concrete_children(position)}))


def checked(position: Position) -> bool:
    return rook_line(position.white_rook, position.black_king, position.white_king)


def digest(path: Path) -> str:
    value = hashlib.sha256()
    with path.open("rb") as handle:
        while block := handle.read(1 << 20):
            value.update(block)
    return value.hexdigest()


def read_json(path: Path) -> Any:
    with path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def read_certificate(path: Path) -> list[dict[str, Any]]:
    rows = []
    with path.open("r", encoding="utf-8") as handle:
        for number, line in enumerate(handle, 1):
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                raise VerificationError(f"invalid certificate JSON on line {number}") from exc
            rows.append(row)
    return rows


def enumerate_universe() -> tuple[set[int], int]:
    keys: set[int] = set()
    raw = 0
    for white_king in range(64):
        for white_rook in range(64):
            if white_rook == white_king:
                continue
            for black_king in range(64):
                for side in (WHITE, BLACK):
                    position = Position(white_king, white_rook, black_king, side)
                    if admissible(position):
                        raw += 1
                        keys.add(orbit_key(position))
    return keys, raw


def verify(run_dir: Path) -> dict[str, Any]:
    certificate_path = run_dir / "krk_tablebase_certificate.jsonl"
    manifest_path = run_dir / "krk_manifest.json"
    pv_path = run_dir / "krk_hardest_principal_variation.json"
    for path in (certificate_path, manifest_path, pv_path):
        if not path.exists():
            raise VerificationError(f"missing {path.name}")
    manifest = read_json(manifest_path)
    if manifest.get("schema_version") != "dimensional-chess-krk/1.0":
        raise VerificationError("unexpected manifest schema")
    if digest(certificate_path) != manifest["certificate_sha256"]:
        raise VerificationError("certificate digest mismatch")
    if digest(pv_path) != manifest["principal_variation_sha256"]:
        raise VerificationError("principal-variation digest mismatch")

    rows = read_certificate(certificate_path)
    codes = [int(row["state"]) for row in rows]
    if codes != sorted(set(codes)):
        raise VerificationError("certificate states are duplicated or unsorted")
    supplied = set(codes)
    expected, raw_count = enumerate_universe()
    if supplied != expected:
        missing = len(expected - supplied)
        extra = len(supplied - expected)
        raise VerificationError(f"certificate universe mismatch: {missing} missing, {extra} extra")
    if raw_count != manifest["raw_legal_states"] or len(expected) != manifest["d4_canonical_states"]:
        raise VerificationError("manifest state counts disagree with independent enumeration")

    table = {int(row["state"]): row for row in rows}
    symmetry_comparisons = 0
    edge_count = 0
    for code in codes:
        representative = unpack(code)
        if orbit_key(representative) != code or not admissible(representative):
            raise VerificationError(f"state {code} is not a legal canonical representative")
        base_children = child_keys(representative)
        base_check = checked(representative)
        edge_count += len(base_children)
        for operation in range(8):
            transformed = image(representative, operation)
            if orbit_key(transformed) != code:
                raise VerificationError(f"orbit key changed under symmetry at {code}")
            if child_keys(transformed) != base_children or checked(transformed) != base_check:
                raise VerificationError(f"unsafe projection collision at {code}")
            symmetry_comparisons += 1

        row = table[code]
        label = row.get("wdl")
        dtm = row.get("dtm_plies")
        policy = row.get("policy")
        nonterminal_children = [child for child in base_children if child != CAPTURE_DRAW]
        for child in nonterminal_children:
            if child not in table:
                raise VerificationError(f"state {code} points outside the declared universe")
        losing_children = [child for child in nonterminal_children if table[child]["wdl"] == "L"]
        winning_children = [child for child in nonterminal_children if table[child]["wdl"] == "W"]
        drawing_children = [child for child in nonterminal_children if table[child]["wdl"] == "D"]

        checkmate = not base_children and representative.side == BLACK and checked(representative)
        if label == "W":
            if not losing_children:
                raise VerificationError(f"winning state {code} has no losing child")
            minimum = min(int(table[child]["dtm_plies"]) for child in losing_children)
            if dtm != minimum + 1 or policy not in losing_children or table[int(policy)]["dtm_plies"] != minimum:
                raise VerificationError(f"winning Bellman or policy equation failed at {code}")
        elif label == "L":
            if checkmate:
                if dtm != 0 or policy is not None:
                    raise VerificationError(f"checkmate certificate failed at {code}")
            else:
                if CAPTURE_DRAW in base_children or not base_children or len(winning_children) != len(base_children):
                    raise VerificationError(f"losing state {code} has a non-winning escape")
                maximum = max(int(table[child]["dtm_plies"]) for child in winning_children)
                if dtm != maximum + 1 or policy not in winning_children or table[int(policy)]["dtm_plies"] != maximum:
                    raise VerificationError(f"losing Bellman or delay policy failed at {code}")
        elif label == "D":
            if losing_children:
                raise VerificationError(f"draw state {code} can move to a loss")
            has_draw_escape = CAPTURE_DRAW in base_children or bool(drawing_children) or not base_children
            if not has_draw_escape or dtm is not None:
                raise VerificationError(f"draw fixed-point equation failed at {code}")
            if CAPTURE_DRAW in base_children:
                if policy != "ROOK_CAPTURE_DRAW":
                    raise VerificationError(f"draw capture policy failed at {code}")
            elif drawing_children:
                if policy not in drawing_children:
                    raise VerificationError(f"draw-preserving policy failed at {code}")
            elif policy is not None:
                raise VerificationError(f"terminal stalemate has a policy at {code}")
        else:
            raise VerificationError(f"invalid WDL label at {code}")

    if edge_count != manifest["canonical_edges_including_capture_draw"]:
        raise VerificationError("edge count disagrees with manifest")
    if symmetry_comparisons != manifest["symmetry_transition_comparisons"]:
        raise VerificationError("symmetry comparison count disagrees with manifest")

    counts = Counter(row["wdl"] for row in rows)
    if dict(counts) != manifest["wdl_counts"]:
        raise VerificationError("WDL counts disagree with manifest")
    finite_distances = [int(row["dtm_plies"]) for row in rows if row["wdl"] in ("W", "L")]
    if max(finite_distances) != manifest["maximum_dtm_plies"]:
        raise VerificationError("maximum DTM disagrees with manifest")
    if any(row["wdl"] == "L" and unpack(int(row["state"])).side == WHITE for row in rows):
        raise VerificationError("unexpected White-to-move loss in KRK")
    if any(row["wdl"] == "W" and unpack(int(row["state"])).side == BLACK for row in rows):
        raise VerificationError("unexpected Black-to-move win in KRK")

    line = read_json(pv_path)
    if not isinstance(line, list) or len(line) != manifest["principal_variation_length"]:
        raise VerificationError("principal variation length mismatch")
    for index, row in enumerate(line):
        concrete = unpack(int(row["concrete_state"]))
        canonical = orbit_key(concrete)
        if not admissible(concrete) or canonical != row["canonical_state"]:
            raise VerificationError(f"bad concrete position on PV row {index}")
        certificate = table[canonical]
        if row["wdl"] != certificate["wdl"] or row["dtm_plies"] != certificate["dtm_plies"]:
            raise VerificationError(f"PV value mismatch on row {index}")
        if index + 1 < len(line):
            following = unpack(int(line[index + 1]["concrete_state"]))
            if following not in concrete_children(concrete):
                raise VerificationError(f"PV contains an illegal move at row {index}")
            if int(line[index + 1]["dtm_plies"]) != int(row["dtm_plies"]) - 1:
                raise VerificationError(f"PV distance does not decrease at row {index}")
            if not row.get("optimal_move"):
                raise VerificationError(f"PV omits move notation at row {index}")
    final_position = unpack(int(line[-1]["concrete_state"]))
    if line[-1]["dtm_plies"] != 0 or concrete_children(final_position) or not checked(final_position):
        raise VerificationError("principal variation does not terminate in checkmate")

    return {
        "verification": "PASS",
        "raw_legal_states": raw_count,
        "canonical_states": len(codes),
        "canonical_edges": edge_count,
        "symmetry_transition_comparisons": symmetry_comparisons,
        "wdl_counts": dict(counts),
        "maximum_dtm_plies": max(finite_distances),
        "principal_variation_legal_plies": len(line) - 1,
        "certificate_sha256": digest(certificate_path),
        "claim": "Exact independent verification of the declared KRK abstraction; not full chess.",
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run-dir", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    try:
        result = verify(parse_args().run_dir)
    except (VerificationError, KeyError, ValueError, TypeError, json.JSONDecodeError) as exc:
        print(f"INDEPENDENT CHESS VERIFICATION FAIL: {exc}", file=sys.stderr)
        raise SystemExit(2) from exc
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()