"""Cross-check Rust-produced (seckey, msg, sig) triples against the pure-Python references
`wp34.ecdsa` and `wp34.schnorr` (exact arithmetic). Used by run_bench.py (pass counts go into the
results JSON) and by tests/test_sig_crosscheck.py (pytest gate).

Usage: PYTHONPATH=src python experiments/signature_benchmarks/python_crosscheck.py <triples.json>
"""
from __future__ import annotations
import json, sys, time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
from wp34 import ecdsa as D, schnorr as B  # noqa: E402


def crosscheck(path: str | Path, resign: bool = True) -> dict:
    doc = json.loads(Path(path).read_text())
    items = doc["items"]
    n = len(items)
    c = {
        "n": n, "source": str(path),
        "ecdsa_k256_sig_verified_by_python": 0, "ecdsa_libsecp_sig_verified_by_python": 0,
        "ecdsa_python_resign_identical": 0,
        "schnorr_k256_sig_verified_by_python": 0, "schnorr_libsecp_sig_verified_by_python": 0,
        "schnorr_python_resign_identical": 0,
        "tampered_sig_rejected_by_python": 0,
    }
    t0 = time.time()
    for it in items:
        d = int(it["seckey"], 16)
        msg = bytes.fromhex(it["msg"])
        aux = bytes.fromhex(it["aux_rand"])
        pub = bytes.fromhex(it["pubkey_uncompressed"])
        Q = (int.from_bytes(pub[1:33], "big"), int.from_bytes(pub[33:65], "big"))
        for key, cnt in (("ecdsa_sig_k256", "ecdsa_k256_sig_verified_by_python"),
                         ("ecdsa_sig_libsecp256k1", "ecdsa_libsecp_sig_verified_by_python")):
            sig = bytes.fromhex(it[key])
            r, s = int.from_bytes(sig[:32], "big"), int.from_bytes(sig[32:], "big")
            if D.verify(Q, msg, (r, s)):
                c[cnt] += 1
        if resign:
            r, s = D.sign(d, msg)  # RFC 6979, low-S
            if r.to_bytes(32, "big") + s.to_bytes(32, "big") == bytes.fromhex(it["ecdsa_sig_k256"]):
                c["ecdsa_python_resign_identical"] += 1
        pkx = bytes.fromhex(it["pubkey_xonly"])
        for key, cnt in (("schnorr_sig_k256", "schnorr_k256_sig_verified_by_python"),
                         ("schnorr_sig_libsecp256k1", "schnorr_libsecp_sig_verified_by_python")):
            sig = bytes.fromhex(it[key])
            if B.verify(msg, pkx, sig):
                c[cnt] += 1
        if resign:
            sig = B.sign(msg, d.to_bytes(32, "big"), aux)
            if sig == bytes.fromhex(it["schnorr_sig_k256"]):
                c["schnorr_python_resign_identical"] += 1
        # negative control: flip one bit of s -> Python must reject
        sig = bytearray(bytes.fromhex(it["schnorr_sig_libsecp256k1"]))
        sig[63] ^= 0x01
        if not B.verify(msg, pkx, bytes(sig)):
            c["tampered_sig_rejected_by_python"] += 1
    c["elapsed_s"] = round(time.time() - t0, 2)
    c["resign_checked"] = resign
    keys = ["ecdsa_k256_sig_verified_by_python", "ecdsa_libsecp_sig_verified_by_python",
            "schnorr_k256_sig_verified_by_python", "schnorr_libsecp_sig_verified_by_python",
            "tampered_sig_rejected_by_python"]
    if resign:
        keys += ["ecdsa_python_resign_identical", "schnorr_python_resign_identical"]
    c["all_pass"] = all(c[k] == n for k in keys)
    return c


if __name__ == "__main__":
    print(json.dumps(crosscheck(sys.argv[1]), indent=1))