"""BIP340 Schnorr signatures over secp256k1 — reference implementation following the BIP text, exact arithmetic."""
from __future__ import annotations
import hashlib, random
from . import secp256k1 as k1
from .ec import Curve, sqrt_mod

E: Curve = k1.CURVE
n, p = k1.n, k1.p


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


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


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


def has_even_y(P) -> bool:
    return P is not None and P[1] % 2 == 0


def lift_x(x: int):
    if x >= p:
        return None
    c = (pow(x, 3, p) + 7) % p
    y = pow(c, (p + 1) // 4, p)
    if y * y % p != c:
        return None
    return (x, y if y % 2 == 0 else p - y)


def pubkey_gen(seckey: bytes) -> bytes:
    d0 = int_from_bytes(seckey)
    if not (1 <= d0 <= n - 1):
        raise ValueError("secret key out of range")
    P = E.mul(d0, E.G)
    return bytes_from_int(P[0])


def sign(msg: bytes, seckey: bytes, aux_rand: bytes) -> bytes:
    d0 = int_from_bytes(seckey)
    if not (1 <= d0 <= n - 1):
        raise ValueError("secret key out of range")
    P = E.mul(d0, E.G)
    d = d0 if has_even_y(P) else n - d0
    t = bytes_from_int(d ^ int_from_bytes(tagged_hash("BIP0340/aux", aux_rand)))
    k0 = int_from_bytes(tagged_hash("BIP0340/nonce", t + bytes_from_int(P[0]) + msg)) % n
    if k0 == 0:
        raise RuntimeError("nonce is zero")
    R = E.mul(k0, E.G)
    k = n - k0 if not has_even_y(R) else k0
    e = int_from_bytes(tagged_hash("BIP0340/challenge", bytes_from_int(R[0]) + bytes_from_int(P[0]) + msg)) % n
    sig = bytes_from_int(R[0]) + bytes_from_int((k + e * d) % n)
    if not verify(msg, bytes_from_int(P[0]), sig):
        raise RuntimeError("signature failed self-check")
    return sig


def verify(msg: bytes, pubkey: bytes, sig: bytes) -> bool:
    if len(pubkey) != 32 or len(sig) != 64:
        return False
    P = lift_x(int_from_bytes(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 = E.add(E.mul(s, E.G), E.mul(n - e, P))
    return R is not None and has_even_y(R) and R[0] == r


def batch_verify(items: list[tuple[bytes, bytes, bytes]], rng: random.Random | None = None) -> bool:
    """Batch verification: Σ a_i s_i G = Σ a_i R_i + Σ a_i e_i P_i with random 128-bit weights (a_1 = 1)."""
    rng = rng or random.SystemRandom()
    S = 0
    rhs = None
    for i, (msg, pubkey, sig) in enumerate(items):
        P = lift_x(int_from_bytes(pubkey))
        r, s = int_from_bytes(sig[:32]), int_from_bytes(sig[32:])
        if P is None or r >= p or s >= n:
            return False
        R = lift_x(r)
        if R is None:
            return False
        e = int_from_bytes(tagged_hash("BIP0340/challenge", sig[:32] + pubkey + msg)) % n
        ai = 1 if i == 0 else rng.getrandbits(128)
        S = (S + ai * s) % n
        rhs = E.add(rhs, E.add(E.mul(ai, R), E.mul(ai * e % n, P)))
    return E.mul(S, E.G) == rhs