import hashlib
import pabtc.secp256k1
import secrets
import typing

# Schnorr Signatures for secp256k1.
# See: https://github.com/bitcoin/bips/blob/master/bip-0340.mediawiki


def prikey_implicit(prikey: pabtc.secp256k1.Fr) -> pabtc.secp256k1.Fr:
    # Make sure the public key is even.
    pubkey = pabtc.secp256k1.G * prikey
    if pubkey == pubkey_implicit(pubkey):
        return +prikey
    else:
        return -prikey


def pubkey_implicit(pubkey: pabtc.secp256k1.Pt) -> pabtc.secp256k1.Pt:
    # Make sure the public key is even.
    if pubkey.y.n & 1:
        return -pubkey
    else:
        return +pubkey


def hash(name: str, data: bytearray) -> bytearray:
    # Cryptographic hash function.
    tag = bytearray(hashlib.sha256(name.encode()).digest())
    out = bytearray(hashlib.sha256(tag + tag + data).digest())
    return out


def sign(prikey: pabtc.secp256k1.Fr, m: pabtc.secp256k1.Fr) -> typing.Tuple[pabtc.secp256k1.Pt, pabtc.secp256k1.Fr]:
    # Sign message m with private key prikey.
    prikey = prikey_implicit(prikey)
    pubkey = pabtc.secp256k1.G * prikey
    k = prikey_implicit(pabtc.secp256k1.Fr(max(1, secrets.randbelow(pabtc.secp256k1.N))))
    r = pabtc.secp256k1.G * k
    e_data = bytearray(r.x.n.to_bytes(32) + pubkey.x.n.to_bytes(32) + m.n.to_bytes(32))
    e_hash = hash('BIP0340/challenge', e_data)
    e = pabtc.secp256k1.Fr(int.from_bytes(e_hash))
    s = k + e * prikey
    return r, s


def verify(pubkey: pabtc.secp256k1.Pt, m: pabtc.secp256k1.Fr, r: pabtc.secp256k1.Pt, s: pabtc.secp256k1.Fr) -> bool:
    # Verify signature (r, s) on message m with public key pubkey.
    pubkey = pubkey_implicit(pubkey)
    e_data = bytearray(r.x.n.to_bytes(32) + pubkey.x.n.to_bytes(32) + m.n.to_bytes(32))
    e_hash = hash('BIP0340/challenge', e_data)
    e = pabtc.secp256k1.Fr(int.from_bytes(e_hash))
    return pabtc.secp256k1.G * s == r + pubkey * e
