#!/usr/bin/env python3
"""
Standalone QBEACON verifier. Independent of the qbeacon package.

    python verify.py pulses.jsonl prover_pubkey.pem [genesis_hex]

Dependencies: the Python standard library plus `cryptography`. NOT numpy.

This file deliberately duplicates protocol logic rather than importing it:
two independent implementations that must agree.

Every pulse records its `device_class` and `soundness_assumptions`; the
verifier prints them for the ledger you check.

(c) Liju James / Jameson Machining Inc. -- 2026-07-26
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import json
import math
import struct
import sys
from collections import defaultdict

from cryptography.hazmat.primitives import serialization

PROTOCOL_VERSION = "qbeacon/v3"

DOM_SETTINGS = b"qbeacon/v3/settings"
DOM_CHALLENGE = b"qbeacon/v3/challenge"
DOM_LEAF = b"qbeacon/v3/leaf"
DOM_PULSE = b"qbeacon/v3/pulse"
DOM_R_RAW = b"qbeacon/v3/R_raw"
DOM_GENESIS = b"qbeacon/v3/genesis"
DOM_RELAY_COMMIT = b"qbeacon/v3/relay-commit"
DOM_EXTRACTOR_SEED = b"qbeacon/v3/extractor-seed"

# Tier A: a relay pulse carries no Bell pairs, no outcomes, no CHSH value --
# the generic settings/commitment/challenge pipeline below does not apply to
# it at all. Checked via _verify_relay_pulse instead. Kept as a tuple, same
# as qbeacon.prover.RELAY_DEVICE_CLASSES, in case a second relay source is
# ever added.
RELAY_DEVICE_CLASSES = ("anu_qrng_relay",)

TSIRELSON = 2.0 * math.sqrt(2.0)
PBR_KAPPA = 0.55
P_WIN_QUANTUM = math.cos(math.pi / 8) ** 2
EPS_BELL = 2.0**-64
EPS_STAT = 1e-6


# ─── Hashing ─────────────────────────────────────────────────────────────
def tagged_hash(tag, *parts):
    h = hashlib.sha256()
    h.update(len(tag).to_bytes(4, "big"))
    h.update(tag)
    for p in parts:
        h.update(len(p).to_bytes(4, "big"))
        h.update(p)
    return h.digest()


def canonical_json(obj):
    return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode(
        "utf-8"
    )


# ─── Merkle (RFC 6962) ───────────────────────────────────────────────────
def leaf_hash(data):
    return hashlib.sha256(b"\x00" + data).digest()


def node_hash(left, right):
    return hashlib.sha256(b"\x01" + left + right).digest()


def _split(n):
    k = 1
    while k << 1 < n:
        k <<= 1
    return k


def merkle_root(leaves):
    if not leaves:
        return hashlib.sha256(b"").digest()
    layer = [leaf_hash(d) for d in leaves]

    def rec(off, size):
        if size == 1:
            return layer[off]
        k = _split(size)
        return node_hash(rec(off, k), rec(off + k, size - k))

    return rec(0, len(layer))


# ─── Settings ────────────────────────────────────────────────────────────
def _expand_bits(seed, n_bits):
    out = bytearray()
    counter = 0
    while len(out) * 8 < n_bits:
        out += hashlib.sha256(seed + counter.to_bytes(4, "big")).digest()
        counter += 1
    bits = []
    for byte in out:
        for k in range(7, -1, -1):
            bits.append((byte >> k) & 1)
            if len(bits) == n_bits:
                return bits
    return bits


def genesis_hash(label=b"default"):
    return tagged_hash(DOM_GENESIS, label).hex()


def settings_seed(prev_hash, pulse_id, n_pairs):
    return tagged_hash(
        DOM_SETTINGS,
        bytes.fromhex(prev_hash),
        pulse_id.to_bytes(4, "big"),
        n_pairs.to_bytes(4, "big"),
    ).hex()


def derive_bases(seed_hex, n_pairs):
    bits = _expand_bits(bytes.fromhex(seed_hex), n_pairs * 2)
    return bits[0::2][:n_pairs], bits[1::2][:n_pairs]


def derive_challenge(commitment_hex, prev_hash, pulse_id):
    return tagged_hash(
        DOM_CHALLENGE,
        bytes.fromhex(commitment_hex),
        bytes.fromhex(prev_hash),
        pulse_id.to_bytes(4, "big"),
    ).hex()


def derive_test_bits(challenge_hex, n_pairs):
    return _expand_bits(bytes.fromhex(challenge_hex), n_pairs)[:n_pairs]


# ─── Transcript ──────────────────────────────────────────────────────────
def unpack_outcomes(b64, n_pairs):
    raw = base64.b64decode(b64)
    if len(raw) < (2 * n_pairs + 7) // 8:
        raise ValueError("transcript too short")
    bits = []
    for byte in raw:
        for k in range(7, -1, -1):
            bits.append((byte >> k) & 1)
    a = [1 - 2 * bits[2 * i] for i in range(n_pairs)]
    b = [1 - 2 * bits[2 * i + 1] for i in range(n_pairs)]
    return a, b


def leaf_data(i, ba, bb, a, b):
    return DOM_LEAF + struct.pack(">IBBBB", i, ba, bb, (1 - a) // 2, (1 - b) // 2)


# ─── CHSH ────────────────────────────────────────────────────────────────
def chsh_S(revealed):
    sums, counts = defaultdict(int), defaultdict(int)
    for r in revealed:
        key = (r["ba"], r["bb"])
        sums[key] += r["a"] * r["b"]
        counts[key] += 1
    e = {k: sums[k] / counts[k] for k in sums}
    return e.get((0, 0), 0.0) + e.get((0, 1), 0.0) + e.get((1, 0), 0.0) - e.get((1, 1), 0.0)


def is_win(a, b, ba, bb):
    return (((1 - a) // 2) ^ ((1 - b) // 2)) == (ba & bb)


def pbr_pvalue(revealed, kappa=PBR_KAPPA):
    if not revealed:
        return 1.0
    win = math.log2(1 + kappa / 4)
    loss = math.log2(1 - 3 * kappa / 4)
    log2_T = 0.0
    for r in revealed:
        log2_T += win if is_win(r["a"], r["b"], r["ba"], r["bb"]) else loss
    if log2_T <= 0:
        return 1.0
    return 2.0 ** (-log2_T) if log2_T < 1000 else 0.0


def tsirelson_allowance(n_min_cell, eps=EPS_STAT):
    if n_min_cell <= 0:
        return float("inf")
    return math.sqrt(8.0 * math.log(2.0 / eps) / n_min_cell)


def no_signaling_violation(revealed):
    """Largest dependence of one party's marginal on the OTHER party's setting.

    Zero in expectation for any genuine two-station device. Catches the
    partially-signaling prover that sits below the Tsirelson ceiling.
    """
    acc = {}
    for r in revealed:
        key = (r["ba"], r["bb"])
        if key not in acc:
            acc[key] = [0, 0, 0]
        acc[key][0] += r["a"] == 1
        acc[key][1] += r["b"] == 1
        acc[key][2] += 1
    if len(acc) < 4:
        return 0.0
    worst = 0.0
    for bb in (0, 1):
        worst = max(
            worst, abs(acc[(0, bb)][1] / acc[(0, bb)][2] - acc[(1, bb)][1] / acc[(1, bb)][2])
        )
    for ba in (0, 1):
        worst = max(
            worst, abs(acc[(ba, 0)][0] / acc[(ba, 0)][2] - acc[(ba, 1)][0] / acc[(ba, 1)][2])
        )
    return worst


def no_signaling_allowance(n_min_cell, eps=EPS_STAT):
    if n_min_cell <= 0:
        return float("inf")
    return math.sqrt(2.0 * math.log(16.0 / eps) / n_min_cell)


def min_revealed_for_grinding(lam=64, log2_q_max=60, kappa=PBR_KAPPA):
    gain = P_WIN_QUANTUM * math.log2(1 + kappa / 4) + (1 - P_WIN_QUANTUM) * math.log2(
        1 - 3 * kappa / 4
    )
    return math.ceil((lam + log2_q_max) / gain)


# ─── R ───────────────────────────────────────────────────────────────────
def r_raw(seed_hex, hidden_leaves):
    return tagged_hash(DOM_R_RAW, bytes.fromhex(seed_hex), b"".join(hidden_leaves)).hex()


def pulse_hash(d):
    return tagged_hash(DOM_PULSE, canonical_json(d)).hex()


# ─── Relay (Tier A) extraction ──────────────────────────────────────────
# Ported from qbeacon.entropy / qbeacon.prover / qbeacon.verifier's
# _verify_relay_pulse -- the two must agree, which is the entire point of
# this file. See qbeacon.entropy.extractor_seed's docstring for why the
# extractor seed is fiat-shamir(prev_hash, pulse_id), NOT derived from the
# outcomes/source bytes themselves (a v3 bug class this fixes).
def _bytes_to_bits(data):
    bits = []
    for byte in data:
        for k in range(7, -1, -1):
            bits.append((byte >> k) & 1)
    return bits


def extractor_seed(prev_hash, pulse_id, challenge, challenge_source):
    if challenge_source != "fiat-shamir":
        return bytes.fromhex(challenge)
    return hashlib.sha256(
        DOM_EXTRACTOR_SEED + bytes.fromhex(prev_hash) + pulse_id.to_bytes(4, "big")
    ).digest()


def _expand_seed_shake(seed, n_bits):
    need = (n_bits + 7) // 8
    raw = hashlib.shake_256(DOM_EXTRACTOR_SEED + seed).digest(need)
    bits = []
    for byte in raw:
        for k in range(7, -1, -1):
            bits.append((byte >> k) & 1)
            if len(bits) == n_bits:
                return bits
    return bits


def toeplitz_extract(source_bits, seed, out_len_bits):
    n = len(source_bits)
    m = out_len_bits
    if m <= 0:
        return b""
    if m > n:
        raise ValueError(f"cannot extract {m} bits from a {n}-bit source")
    t = _expand_seed_shake(seed, n + m - 1)
    out_bits = []
    for j in range(m):
        acc = 0
        base = j + n - 1
        for i in range(n):
            if source_bits[i]:
                acc ^= t[base - i]
        out_bits.append(acc)
    out = bytearray()
    for i in range(0, len(out_bits), 8):
        byte = 0
        for k in range(8):
            byte = (byte << 1) | (out_bits[i + k] if i + k < len(out_bits) else 0)
        out.append(byte)
    return bytes(out)


def _verify_relay_pulse(d, prev_hash):
    """anu_qrng_relay (and any future RELAY_DEVICE_CLASSES member): checks
    signature/chain/version have already passed by the time this runs.
    Verifies the commitment binds the published raw source bytes to this
    chain position, and that R is the honest Toeplitz extract of those
    bytes -- NOT a Bell test, no settings/commitment-tree/challenge to check."""
    info = {"S_hat": 0.0, "p_value": None, "device_class": d.get("device_class", "?")}

    if not (d.get("n_pairs") == 0 and d.get("outcomes_b64", "") == ""):
        return False, "TRANSCRIPT_MALFORMED", info

    try:
        source_bytes = bytes.fromhex(d.get("R_raw") or "")
    except ValueError:
        return False, "TRANSCRIPT_MALFORMED", info
    if not source_bytes:
        return False, "TRANSCRIPT_MALFORMED", info

    prev_bytes = bytes.fromhex(prev_hash) if prev_hash else b""
    expected_commitment = tagged_hash(
        DOM_RELAY_COMMIT, prev_bytes, d["pulse_id"].to_bytes(8, "big"), source_bytes
    ).hex()
    if expected_commitment != d["commitment"]:
        return False, "ROOT_MISMATCH", info

    try:
        source_bits = _bytes_to_bits(source_bytes)
        ext_seed = extractor_seed(
            d["prev_hash"], d["pulse_id"], d.get("challenge", ""), d.get("challenge_source", "")
        )
        r_recomputed = toeplitz_extract(
            source_bits, ext_seed, d.get("output_len_bits", 0)
        ).hex()
    except Exception:
        return False, "R_MISMATCH", info
    if r_recomputed != d.get("R", ""):
        return False, "R_MISMATCH", info

    return True, "OK", info


# ─── Verification ────────────────────────────────────────────────────────
def verify_pulse(d, pk, prev_hash, eps_bell=EPS_BELL, enforce_grinding=True):
    """Returns (ok, reason_code, detail_dict)."""
    info = {"S_hat": None, "p_value": None, "device_class": d.get("device_class", "?")}

    payload = canonical_json({k: v for k, v in d.items() if k != "signature"})
    try:
        pk.verify(bytes.fromhex(d.get("signature", "")), payload)
    except Exception:
        return False, "SIG_INVALID", info

    if d["prev_hash"] != prev_hash:
        return False, "CHAIN_BROKEN", info
    if d.get("protocol_version") != PROTOCOL_VERSION:
        return False, "VERSION_MISMATCH", info

    if d.get("device_class") in RELAY_DEVICE_CLASSES:
        return _verify_relay_pulse(d, prev_hash)

    n = d["n_pairs"]
    try:
        oa, ob = unpack_outcomes(d["outcomes_b64"], n)
    except Exception:
        return False, "TRANSCRIPT_MALFORMED", info

    seed = settings_seed(d["prev_hash"], d["pulse_id"], n)
    if seed != d["settings_seed"]:
        return False, "SETTINGS_MISMATCH", info
    ba, bb = derive_bases(seed, n)

    leaves = [leaf_data(i, ba[i], bb[i], oa[i], ob[i]) for i in range(n)]
    if merkle_root(leaves).hex() != d["commitment"]:
        return False, "ROOT_MISMATCH", info

    if d.get("challenge_source") == "fiat-shamir":
        if derive_challenge(d["commitment"], d["prev_hash"], d["pulse_id"]) != d["challenge"]:
            return False, "CHALLENGE_MISMATCH", info

    t = derive_test_bits(d["challenge"], n)
    rev_idx = [i for i in range(n) if t[i] == 1]
    hid_idx = [i for i in range(n) if t[i] == 0]
    revealed = [{"i": i, "ba": ba[i], "bb": bb[i], "a": oa[i], "b": ob[i]} for i in rev_idx]

    if len(revealed) < 100:
        return False, "INSUFFICIENT_REVEALED", info
    counts = defaultdict(int)
    for r in revealed:
        counts[(r["ba"], r["bb"])] += 1
    if len(counts) < 4:
        return False, "INSUFFICIENT_REVEALED", info
    n_min_cell = min(counts.values())

    budget = d.get("grinding_budget_bits", 60)
    if enforce_grinding and d.get("challenge_source") == "fiat-shamir":
        if len(revealed) < min_revealed_for_grinding(log2_q_max=budget):
            return False, "GRINDING_BUDGET_UNMET", info

    S = chsh_S(revealed)
    info["S_hat"] = S
    allowance = tsirelson_allowance(n_min_cell)
    if abs(S) > TSIRELSON + allowance:
        return False, "PROVER_LYING_SUPERQUANTUM", info

    ns = no_signaling_violation(revealed)
    info["no_signaling_violation"] = ns
    if ns > no_signaling_allowance(n_min_cell):
        return False, "SIGNALING_DETECTED", info

    p = pbr_pvalue(revealed)
    info["p_value"] = p
    threshold = eps_bell / (2.0**budget) if d.get("challenge_source") == "fiat-shamir" else eps_bell
    if p > threshold:
        return False, "BELL_NOT_VIOLATED", info

    if r_raw(seed, [leaves[i] for i in hid_idx]) != d["R_raw"]:
        return False, "R_MISMATCH", info

    # Checks R_raw, the committed beacon material that binds the output to the
    # transcript; R is a deterministic post-processing of it.
    return True, "OK", info


def main(argv=None):
    ap = argparse.ArgumentParser(description="Standalone QBEACON v3 verifier")
    ap.add_argument("ledger")
    ap.add_argument("pubkey")
    ap.add_argument(
        "genesis",
        nargs="?",
        default=None,
        help="expected genesis prev_hash (default: the protocol genesis)",
    )
    args = ap.parse_args(argv)

    pk = serialization.load_pem_public_key(open(args.pubkey, "rb").read())
    prev = args.genesis or genesis_hash()

    n_ok = n_total = 0
    assumptions = set()
    device_classes = set()

    with open(args.ledger, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            d = json.loads(line)
            n_total += 1
            ok, code, info = verify_pulse(d, pk, prev)
            n_ok += ok
            assumptions.update(d.get("soundness_assumptions", []))
            device_classes.add(d.get("device_class", "?"))
            s = "n/a" if info["S_hat"] is None else f"{info['S_hat']:+.4f}"
            p = "n/a" if info["p_value"] is None else f"{info['p_value']:.2e}"
            print(
                f"pulse #{d['pulse_id']:4d}  S={s:>8}  p={p:>9}  {'OK  ' if ok else 'FAIL'}  {code}"
            )
            prev = pulse_hash(d)

    print(f"\n{n_ok}/{n_total} pulses verified")
    print(f"device classes: {', '.join(sorted(device_classes))}")
    print("\nVerified: every pulse is well-formed, signed and chained, and its")
    print("statistics are re-derived from the published outcomes.")
    print("Stated assumptions carried by this ledger:")
    for a in sorted(assumptions):
        print(f"  - {a}")
    return 0 if n_ok == n_total and n_total > 0 else 1


if __name__ == "__main__":
    sys.exit(main())