# ============================================================+
# Identity IoT — Ed25519 for MicroPython
# (c) Copyright : ae Aeonian Engineering Limited - Hong Kong
# (c) Copyright : WIDE di D. Papa - Naples - Italy
# ============================================================+
# Based on Bernstein et al. reference implementation (public domain)
# Iterative scalarmult — no recursion, safe for MicroPython stack
# ============================================================+

from sha512 import sha512

_q = 2**255 - 19
_l = 2**252 + 27742317777372353535851937790883648493
_d = -121665 * pow(121666, _q - 2, _q) % _q
_I = pow(2, (_q - 1) // 4, _q)


def _xrecover(y):
    x2 = (y * y - 1) * pow(_d * y * y + 1, _q - 2, _q)
    x  = pow(x2, (_q + 3) // 8, _q)
    if (x * x - x2) % _q != 0:
        x = x * _I % _q
    if x % 2 != 0:
        x = _q - x
    return x


_By = 4 * pow(5, _q - 2, _q) % _q
_Bx = _xrecover(_By)
_B  = [_Bx % _q, _By % _q, 1, _Bx * _By % _q]


def _edwards_add(P, Q):
    x1, y1, z1, t1 = P
    x2, y2, z2, t2 = Q
    A = (y1 - x1) * (y2 - x2) % _q
    B = (y1 + x1) * (y2 + x2) % _q
    C = t1 * 2 * _d * t2 % _q
    D = z1 * 2 * z2 % _q
    E = B - A; F = D - C; G = D + C; H = B + A
    return [E * F % _q, G * H % _q, F * G % _q, E * H % _q]


def _scalarmult(P, e):
    # Iterative double-and-add — no recursion
    Q = [0, 1, 1, 0]  # neutral element
    while e > 0:
        if e & 1:
            Q = _edwards_add(Q, P)
        P = _edwards_add(P, P)
        e >>= 1
    return Q


def _encodeint(y):
    bits = [(y >> i) & 1 for i in range(256)]
    return bytes([sum([bits[i * 8 + j] << j for j in range(8)]) for i in range(32)])


def _encodepoint(P):
    x, y, z, _ = P
    zi = pow(z, _q - 2, _q)
    x  = x * zi % _q
    y  = y * zi % _q
    bits = [(y >> i) & 1 for i in range(255)] + [x & 1]
    return bytes([sum([bits[i * 8 + j] << j for j in range(8)]) for i in range(32)])


def _bit(h, i):
    return (h[i // 8] >> (i % 8)) & 1


def _hint(m):
    return int.from_bytes(sha512(m), 'little')


def _decodepoint(s):
    y = sum(_bit(s, i) * 2 ** i for i in range(255))
    x = _xrecover(y)
    if x & 1 != _bit(s, 255):
        x = _q - x
    return [x, y, 1, x * y % _q]


# ============================================================
# Public API
# ============================================================

def ed25519_keypair_from_seed(seed: bytes):
    if len(seed) != 32:
        raise ValueError('Seed must be 32 bytes')
    h  = sha512(seed)
    a  = 2 ** 254 + sum(2 ** i * _bit(h, i) for i in range(3, 254))
    A  = _scalarmult(_B, a)
    pk = _encodepoint(A)
    return pk, seed + pk


def ed25519_sign(seed: bytes, message: bytes) -> bytes:
    if len(seed) != 32:
        raise ValueError('Seed must be 32 bytes')
    h  = sha512(seed)
    a  = 2 ** 254 + sum(2 ** i * _bit(h, i) for i in range(3, 254))
    pk = ed25519_keypair_from_seed(seed)[0]
    r  = _hint(h[32:64] + message)
    R  = _scalarmult(_B, r)
    S  = (r + _hint(_encodepoint(R) + pk + message) * a) % _l
    return _encodepoint(R) + _encodeint(S)


def ed25519_verify(pk: bytes, message: bytes, signature: bytes) -> bool:
    if len(signature) != 64 or len(pk) != 32:
        return False
    try:
        R_enc = signature[:32]
        S     = int.from_bytes(signature[32:], 'little')
        if S >= _l:
            return False
        A   = _decodepoint(pk)
        R   = _decodepoint(R_enc)
        h   = _hint(R_enc + pk + message)
        lhs = _scalarmult(_B, 8 * S)
        rhs = _edwards_add(_scalarmult(R, 8), _scalarmult(A, 8 * h))
        return _encodepoint(lhs) == _encodepoint(rhs)
    except Exception:
        return False