"""ECVRF-EDWARDS25519-SHA512-TAI, the verifiable random function of RFC 9381 over the Ed25519 curve.

Non-normative; see spec/08-fidelity.md §8.8.
A VRF proof has exactly one valid value for a key and an input, where an Ed25519 signature has many, so a
spot-check seed that is a VRF proof cannot be ground by the checker. The arithmetic is the reference form of
RFC 8032 in plain Python, slow and clear; the keys are ordinary Ed25519 keys, so the roster needs nothing new.
The RFC's Example 16 is pinned in vectors/vrf.jsonl.
"""

import hashlib

SUITE = b"\x03"
COFACTOR = 8
P = 2**255 - 19
Q = 2**252 + 27742317777372353535851937790883648493
D = (-121665 * pow(121666, P - 2, P)) % P
_I = pow(2, (P - 1) // 4, P)


def _inv(x):
    return pow(x, P - 2, P)


def _recover_x(y, sign):
    if y >= P:
        return None
    x2 = (y * y - 1) * _inv(D * y * y + 1) % P
    if x2 == 0:
        return None if sign else 0
    x = pow(x2, (P + 3) // 8, P)
    if (x * x - x2) % P != 0:
        x = x * _I % P
    if (x * x - x2) % P != 0:
        return None
    if (x & 1) != sign:
        x = P - x
    return x


# Points in extended homogeneous coordinates (X, Y, Z, T) with x = X/Z, y = Y/Z, x*y = T/Z (RFC 8032 §5.1.4).
G_Y = 4 * _inv(5) % P
G_X = _recover_x(G_Y, 0)
G = (G_X, G_Y, 1, G_X * G_Y % P)
IDENTITY = (0, 1, 1, 0)


def _add(a, b):
    x1, y1, z1, t1 = a
    x2, y2, z2, t2 = b
    A = (y1 - x1) * (y2 - x2) % P
    B = (y1 + x1) * (y2 + x2) % P
    C = t1 * 2 * D * t2 % P
    Dd = z1 * 2 * z2 % P
    E, F, Gg, H = B - A, Dd - C, Dd + C, B + A
    return (E * F % P, Gg * H % P, F * Gg % P, E * H % P)


def _mul(s, pt):
    out, acc = IDENTITY, pt
    while s > 0:
        if s & 1:
            out = _add(out, acc)
        acc = _add(acc, acc)
        s >>= 1
    return out


def _equal(a, b):
    return (a[0] * b[2] - b[0] * a[2]) % P == 0 and (a[1] * b[2] - b[1] * a[2]) % P == 0


def point_to_string(pt):
    z = _inv(pt[2])
    x, y = pt[0] * z % P, pt[1] * z % P
    return (y | ((x & 1) << 255)).to_bytes(32, "little")


def string_to_point(s):
    """The point a 32-byte string encodes, or None when it is not a valid encoding (RFC 8032 §5.1.3)."""
    if len(s) != 32:
        return None
    y = int.from_bytes(s, "little")
    sign = y >> 255
    y &= (1 << 255) - 1
    x = _recover_x(y, sign)
    if x is None:
        return None
    return (x, y, 1, x * y % P)


def _secret_scalar(sk):
    h = hashlib.sha512(sk).digest()
    a = bytearray(h[:32])
    a[0] &= 248
    a[31] &= 63
    a[31] |= 64
    return int.from_bytes(a, "little"), h[32:]


def public_key(sk):
    x, _ = _secret_scalar(sk)
    return point_to_string(_mul(x, G))


def encode_to_curve(pk_string, alpha):
    """ECVRF_encode_to_curve_try_and_increment (RFC 9381 §5.4.1.1)."""
    for ctr in range(256):
        h = hashlib.sha512(SUITE + b"\x01" + pk_string + alpha + bytes([ctr]) + b"\x00").digest()
        pt = string_to_point(h[:32])
        if pt is not None:
            pt = _mul(COFACTOR, pt)
            if not _equal(pt, IDENTITY):
                return pt
    raise ValueError("no curve point found")


def _challenge(points):
    """ECVRF_challenge_generation (RFC 9381 §5.4.3): the first 16 octets of the hash, as a little-endian integer."""
    s = SUITE + b"\x02" + b"".join(point_to_string(p) for p in points) + b"\x00"
    return int.from_bytes(hashlib.sha512(s).digest()[:16], "little")


def prove(sk, alpha):
    """ECVRF_prove: the 80-byte proof Gamma || c || s for the 32-byte Ed25519 secret `sk` and input `alpha`."""
    x, suffix = _secret_scalar(sk)
    y = _mul(x, G)
    h = encode_to_curve(point_to_string(y), alpha)
    h_string = point_to_string(h)
    gamma = _mul(x, h)
    k = int.from_bytes(hashlib.sha512(suffix + h_string).digest(), "little") % Q
    c = _challenge([y, h, gamma, _mul(k, G), _mul(k, h)])
    s = (k + c * x) % Q
    return point_to_string(gamma) + c.to_bytes(16, "little") + s.to_bytes(32, "little")


def decode_proof(pi):
    if len(pi) != 80:
        return None
    gamma = string_to_point(pi[:32])
    if gamma is None:
        return None
    c = int.from_bytes(pi[32:48], "little")
    s = int.from_bytes(pi[48:80], "little")
    if s >= Q:
        return None
    return gamma, c, s


def proof_to_hash(pi):
    """ECVRF_proof_to_hash: the 64-byte output beta of a proof (RFC 9381 §5.2), or None for a malformed proof."""
    d = decode_proof(pi)
    if d is None:
        return None
    gamma = d[0]
    return hashlib.sha512(SUITE + b"\x03" + point_to_string(_mul(COFACTOR, gamma)) + b"\x00").digest()


def _neg(pt):
    return ((-pt[0]) % P, pt[1], pt[2], (-pt[3]) % P)


def verify(pk_string, alpha, pi):
    """ECVRF_verify: beta when `pi` is the proof for `alpha` under the public key, else None. Never raises."""
    try:
        y = string_to_point(pk_string)
        if y is None or _equal(_mul(COFACTOR, y), IDENTITY):
            return None
        d = decode_proof(pi)
        if d is None:
            return None
        gamma, c, s = d
        h = encode_to_curve(pk_string, alpha)
        u = _add(_mul(s, G), _neg(_mul(c, y)))
        v = _add(_mul(s, h), _neg(_mul(c, gamma)))
        if _challenge([y, h, gamma, u, v]) != c:
            return None
        return proof_to_hash(pi)
    except (ValueError, TypeError):
        return None