#!/usr/bin/env python3
"""Minimal Nostr/BIP340 crypto helpers.

Pure Python secp256k1 point arithmetic + BIP340 Schnorr sign/verify.
"""

from __future__ import annotations

import hashlib
from typing import Optional, Tuple

P = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F
N = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141
GX = 0x79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798
GY = 0x483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8

Point = Optional[Tuple[int, int]]
G: Point = (GX, GY)


def _sha256(b: bytes) -> bytes:
    return hashlib.sha256(b).digest()


def _tagged_hash(tag: str, msg: bytes) -> bytes:
    t = _sha256(tag.encode("utf-8"))
    return _sha256(t + t + msg)


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


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


def _mod_inv(a: int, n: int) -> int:
    return pow(a, -1, n)


def _is_infinite(p: Point) -> bool:
    return p is None


def _has_even_y(p: Point) -> bool:
    assert p is not None
    return (p[1] % 2) == 0


def _point_add(p1: Point, p2: Point) -> Point:
    if _is_infinite(p1):
        return p2
    if _is_infinite(p2):
        return p1
    assert p1 is not None and p2 is not None
    x1, y1 = p1
    x2, y2 = p2

    if x1 == x2 and (y1 + y2) % P == 0:
        return None

    if x1 == x2 and y1 == y2:
        lam = (3 * x1 * x1) * _mod_inv(2 * y1 % P, P) % P
    else:
        lam = (y2 - y1) * _mod_inv((x2 - x1) % P, P) % P

    x3 = (lam * lam - x1 - x2) % P
    y3 = (lam * (x1 - x3) - y1) % P
    return x3, y3


def _point_mul(k: int, p: Point) -> Point:
    if k % N == 0 or p is None:
        return None
    result: Point = None
    addend = p
    while k:
        if k & 1:
            result = _point_add(result, addend)
        addend = _point_add(addend, addend)
        k >>= 1
    return result


def _lift_x(x: int) -> Point:
    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
    if y % 2 != 0:
        y = P - y
    return (x, y)


def normalize_secret(secret32: bytes) -> int:
    if len(secret32) != 32:
        raise ValueError("secret key must be 32 bytes")
    d = _int_from_bytes(secret32)
    if d == 0 or d >= N:
        raise ValueError("secret key must be in range [1, n-1]")
    return d


def pubkey_gen_xonly(secret32: bytes) -> bytes:
    d = normalize_secret(secret32)
    p = _point_mul(d, G)
    if p is None:
        raise ValueError("invalid secret key")
    return _bytes_from_int(p[0])


def schnorr_sign(msg32: bytes, secret32: bytes, aux_rand32: bytes = b"\x00" * 32) -> bytes:
    if len(msg32) != 32:
        raise ValueError("message hash must be 32 bytes")
    if len(aux_rand32) != 32:
        raise ValueError("aux_rand must be 32 bytes")

    d0 = normalize_secret(secret32)
    p = _point_mul(d0, G)
    if p is None:
        raise ValueError("invalid secret key")

    d = d0 if _has_even_y(p) else N - d0
    px = _bytes_from_int(p[0])

    t = bytes(a ^ b for a, b in zip(_bytes_from_int(d), _tagged_hash("BIP0340/aux", aux_rand32)))
    k0 = _int_from_bytes(_tagged_hash("BIP0340/nonce", t + px + msg32)) % N
    if k0 == 0:
        raise ValueError("nonce generation failed")

    r_point = _point_mul(k0, G)
    assert r_point is not None
    k = k0 if _has_even_y(r_point) else N - k0
    r = _bytes_from_int(r_point[0])

    e = _int_from_bytes(_tagged_hash("BIP0340/challenge", r + px + msg32)) % N
    s = (k + e * d) % N
    sig = r + _bytes_from_int(s)

    if not schnorr_verify(msg32, px, sig):
        raise ValueError("self-verification failed")
    return sig


def schnorr_verify(msg32: bytes, pubkey_xonly32: bytes, sig64: bytes) -> bool:
    if len(msg32) != 32 or len(pubkey_xonly32) != 32 or len(sig64) != 64:
        return False

    px = _int_from_bytes(pubkey_xonly32)
    p = _lift_x(px)
    if p is None:
        return False

    r = _int_from_bytes(sig64[:32])
    s = _int_from_bytes(sig64[32:])
    if r >= P or s >= N:
        return False

    e = _int_from_bytes(_tagged_hash("BIP0340/challenge", sig64[:32] + pubkey_xonly32 + msg32)) % N
    r_point = _point_add(_point_mul(s, G), _point_mul((N - e) % N, p))
    if r_point is None:
        return False
    if not _has_even_y(r_point):
        return False
    return r_point[0] == r


def derive_secret_from_string(seed: str) -> bytes:
    d = _int_from_bytes(_sha256(seed.encode("utf-8"))) % N
    if d == 0:
        d = 1
    return _bytes_from_int(d)
