"""BIP-340 Schnorr signatures over secp256k1 (pure Python).

This is the reference implementation from BIP-340 (public domain), used for the
``nostr-secp256k1-schnorr-sha256-jcs`` suite. We sign with an all-zero auxiliary
randomness value, matching the Rust ``sign_schnorr_no_aux_rand`` and the Node
``schnorr.sign(digest, sk, ZERO_AUX)`` calls — so every implementation produces
byte-identical, deterministic signatures.

We use a pure-Python implementation because the native ``coincurve``/``secp256k1``
bindings fail to build on bleeding-edge CPython. Signing/verifying short
fixed-size digests is well within pure-Python performance for a CLI tool.
"""

from __future__ import annotations

import hashlib

p = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F
n = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141
G = (
    0x79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798,
    0x483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8,
)

Point = tuple[int, int] | None


def tagged_hash(tag: str, msg: bytes) -> bytes:
    tag_hash = hashlib.sha256(tag.encode()).digest()
    return hashlib.sha256(tag_hash + tag_hash + msg).digest()


def is_infinite(P: Point) -> bool:
    return P is None


def x(P: Point) -> int:
    assert P is not None
    return P[0]


def y(P: Point) -> int:
    assert P is not None
    return P[1]


def point_add(P1: Point, P2: Point) -> Point:
    if P1 is None:
        return P2
    if P2 is None:
        return P1
    if x(P1) == x(P2) and y(P1) != y(P2):
        return None
    if P1 == P2:
        lam = (3 * x(P1) * x(P1) * pow(2 * y(P1), p - 2, p)) % p
    else:
        lam = ((y(P2) - y(P1)) * pow(x(P2) - x(P1), p - 2, p)) % p
    x3 = (lam * lam - x(P1) - x(P2)) % p
    return (x3, (lam * (x(P1) - x3) - y(P1)) % p)


def point_mul(P: Point, scalar: int) -> Point:
    R: Point = None
    for i in range(256):
        if (scalar >> i) & 1:
            R = point_add(R, P)
        P = point_add(P, P)
    return R


def bytes_from_int(x_: int) -> bytes:
    return x_.to_bytes(32, byteorder="big")


def bytes_from_point(P: Point) -> bytes:
    return bytes_from_int(x(P))


def lift_x(b: bytes) -> Point:
    x_ = int.from_bytes(b, byteorder="big")
    if x_ >= p:
        return None
    y_sq = (pow(x_, 3, p) + 7) % p
    y_ = pow(y_sq, (p + 1) // 4, p)
    if pow(y_, 2, p) != y_sq:
        return None
    return (x_, y_ if y_ & 1 == 0 else p - y_)


def int_from_bytes(b: bytes) -> int:
    return int.from_bytes(b, byteorder="big")


def has_even_y(P: Point) -> bool:
    assert P is not None
    return y(P) % 2 == 0


def pubkey_gen(seckey: bytes) -> bytes:
    """Return the 32-byte x-only public key for a 32-byte secret key."""
    d0 = int_from_bytes(seckey)
    if not (1 <= d0 <= n - 1):
        raise ValueError("schnorr: secret key out of range")
    P = point_mul(G, d0)
    assert P is not None
    return bytes_from_point(P)


def sign(msg: bytes, seckey: bytes, aux_rand: bytes = b"\x00" * 32) -> bytes:
    """Produce a 64-byte BIP-340 Schnorr signature over ``msg``."""
    d0 = int_from_bytes(seckey)
    if not (1 <= d0 <= n - 1):
        raise ValueError("schnorr: secret key out of range")
    P = point_mul(G, d0)
    assert P is not None
    d = d0 if has_even_y(P) else n - d0
    t = (d ^ int_from_bytes(tagged_hash("BIP0340/aux", aux_rand))).to_bytes(32, "big")
    k0 = int_from_bytes(tagged_hash("BIP0340/nonce", t + bytes_from_point(P) + msg)) % n
    if k0 == 0:
        raise ValueError("schnorr: nonce generation failed")
    R = point_mul(G, k0)
    assert R is not None
    k = k0 if has_even_y(R) else n - k0
    e = (
        int_from_bytes(
            tagged_hash("BIP0340/challenge", bytes_from_point(R) + bytes_from_point(P) + msg)
        )
        % n
    )
    sig = bytes_from_point(R) + ((k + e * d) % n).to_bytes(32, "big")
    if not verify(msg, bytes_from_point(P), sig):
        raise ValueError("schnorr: produced an invalid signature")
    return sig


def verify(msg: bytes, pubkey: bytes, sig: bytes) -> bool:
    """Verify a 64-byte BIP-340 Schnorr signature ``sig`` over ``msg``."""
    if len(pubkey) != 32 or len(sig) != 64:
        return False
    P = lift_x(pubkey)
    r = int_from_bytes(sig[0:32])
    s = int_from_bytes(sig[32:64])
    if P is None or r >= p or s >= n:
        return False
    e = (
        int_from_bytes(tagged_hash("BIP0340/challenge", sig[0:32] + pubkey + msg)) % n
    )
    R = point_add(point_mul(G, s), point_mul(P, n - e))
    if R is None or not has_even_y(R) or x(R) != r:
        return False
    return True