"""Keyed hashing, key derivation, Merkle roots and Ed25519.

Non-normative; see spec/05-hashes.md and spec/08-fidelity.md.
Ed25519 follows the reference code of RFC 8032, section 6; it is slow and is meant only for test vectors.
"""

import hashlib
import hmac

from . import jcs


def hkdf(ikm, salt, info, length=32):
    """HKDF-SHA-256, RFC 5869."""
    prk = hmac.new(salt, ikm, hashlib.sha256).digest()
    out, t, i = b"", b"", 1
    while len(out) < length:
        t = hmac.new(prk, t + info + bytes([i]), hashlib.sha256).digest()
        out += t
        i += 1
    return out[:length]


def team_key(secret):
    return hkdf(secret, b"scienting-team-key/1", b"scienting-hash/1")


def key_id(key):
    return hashlib.sha256(b"scienting-key-id/1" + key).hexdigest()[:16]


def page_hash(page_fp_hex):
    return hashlib.sha256(b"scienting-page-hash/1" + bytes.fromhex(page_fp_hex)).hexdigest()


def page_key(page_fps_hex, context):
    ikm = b"".join(bytes.fromhex(h) for h in page_fps_hex)
    return hkdf(ikm, b"scienting-page-key/1", context.encode("utf-8"))


# --- Merkle trees (RFC 6962 hashing) ------------------------------------------


def _leaf(data):
    return hashlib.sha256(b"\x00" + data).digest()


def _node(left, right):
    return hashlib.sha256(b"\x01" + left + right).digest()


def _mth(leaves):
    n = len(leaves)
    if n == 1:
        return leaves[0]
    k = 1
    while k * 2 < n:
        k *= 2
    return _node(_mth(leaves[:k]), _mth(leaves[k:]))


def merkle_root(items):
    """Root over receipt items: each item is a list such as ["node", id, hash]."""
    if not items:
        raise ValueError("a Merkle root needs at least one item")
    return _mth([_leaf(jcs.dumpb(list(it))) for it in items]).hex()


def receipt_items(receipt):
    items = [["node", nid, h] for nid, h in receipt.get("nodes", [])]
    items += [list(r) for r in receipt.get("records", [])]
    return items


# --- Ed25519 (RFC 8032) -------------------------------------------------------

_p = 2**255 - 19
_d = -121665 * pow(121666, _p - 2, _p) % _p
_q = 2**252 + 27742317777372353535851937790883648493


def _sha512(b):
    return hashlib.sha512(b).digest()


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


def _add(P, Q):
    A = (P[1] - P[0]) * (Q[1] - Q[0]) % _p
    B = (P[1] + P[0]) * (Q[1] + Q[0]) % _p
    C = 2 * P[3] * Q[3] * _d % _p
    D = 2 * P[2] * Q[2] % _p
    E, F, G, H = B - A, D - C, D + C, B + A
    return (E * F % _p, G * H % _p, F * G % _p, E * H % _p)


def _mul(s, P):
    Q = (0, 1, 1, 0)
    while s > 0:
        if s & 1:
            Q = _add(Q, P)
        P = _add(P, P)
        s >>= 1
    return Q


def _equal(P, Q):
    if (P[0] * Q[2] - Q[0] * P[2]) % _p != 0:
        return False
    return (P[1] * Q[2] - Q[1] * P[2]) % _p == 0


_modp_sqrt_m1 = pow(2, (_p - 1) // 4, _p)


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


_gy = 4 * _inv(5) % _p
_gx = _recover_x(_gy, 0)
_G = (_gx, _gy, 1, _gx * _gy % _p)


def _compress(P):
    zinv = _inv(P[2])
    x = P[0] * zinv % _p
    y = P[1] * zinv % _p
    return int.to_bytes(y | ((x & 1) << 255), 32, "little")


def _decompress(s):
    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_expand(secret):
    h = _sha512(secret)
    a = int.from_bytes(h[:32], "little")
    a &= (1 << 254) - 8
    a |= 1 << 254
    return a, h[32:]


def ed25519_public(secret):
    a, _ = _secret_expand(secret)
    return _compress(_mul(a, _G))


def ed25519_sign(secret, msg):
    a, prefix = _secret_expand(secret)
    A = _compress(_mul(a, _G))
    r = int.from_bytes(_sha512(prefix + msg), "little") % _q
    R = _compress(_mul(r, _G))
    h = int.from_bytes(_sha512(R + A + msg), "little") % _q
    s = (r + h * a) % _q
    return R + int.to_bytes(s, 32, "little")


def ed25519_verify(public, msg, signature):
    if len(public) != 32 or len(signature) != 64:
        return False
    A = _decompress(public)
    if not A:
        return False
    R = _decompress(signature[:32])
    if not R:
        return False
    s = int.from_bytes(signature[32:], "little")
    if s >= _q:
        return False
    h = int.from_bytes(_sha512(signature[:32] + public + msg), "little") % _q
    return _equal(_mul(s, _G), _add(R, _mul(h, A)))