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

Non-normative; see spec/05-hashes.md and spec/08-fidelity.md.
HKDF and Ed25519 come from the `cryptography` package.
"""

import hashlib

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

from . import jcs


def hkdf(ikm, salt, info, length=32):
    """HKDF-SHA-256, RFC 5869."""
    return HKDF(algorithm=hashes.SHA256(), length=length, salt=salt, info=info).derive(ikm)


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


EMPTY_ROOT = hashlib.sha256(b"").hexdigest()


def history_root(items):
    """Root over an append-only history, with RFC 9162's hash of the empty list for no items."""
    if not items:
        return EMPTY_ROOT
    return _mth([_leaf(jcs.dumpb(list(it))) for it in items]).hex()


def leaf_hash(item):
    return _leaf(jcs.dumpb(list(item))).hex()


# --- RFC 9162 proofs (§2.1.3 and §2.1.4) ---------------------------------------


def _path(m, leaves):
    """PATH(m, D[n]): the inclusion path for the leaf at index m, as RFC 9162 §2.1.3.1 defines it."""
    n = len(leaves)
    if n == 1:
        return []
    k = 1
    while k * 2 < n:
        k *= 2
    if m < k:
        return _path(m, leaves[:k]) + [_mth(leaves[k:])]
    return _path(m - k, leaves[k:]) + [_mth(leaves[:k])]


def inclusion_path(items, index):
    """The audit path of `items[index]` in the tree over `items`, as hex hashes, leaf-most first."""
    if not 0 <= index < len(items):
        raise ValueError("index outside the tree")
    return [h.hex() for h in _path(index, [_leaf(jcs.dumpb(list(it))) for it in items])]


def verify_inclusion(root, leaf, index, size, path):
    """Whether `path` proves that the leaf hash `leaf` is item `index` of a tree of `size` with `root`
    (RFC 9162 §2.1.3.2). Never raises."""
    try:
        if not 0 <= index < size:
            return False
        fn, sn = index, size - 1
        r = bytes.fromhex(leaf)
        for p in path:
            p = bytes.fromhex(p)
            if sn == 0:
                return False
            if fn & 1 or fn == sn:
                r = _node(p, r)
                while fn & 1 == 0 and fn != 0:
                    fn >>= 1
                    sn >>= 1
            else:
                r = _node(r, p)
            fn >>= 1
            sn >>= 1
        return sn == 0 and r.hex() == root
    except (ValueError, TypeError):
        return False


def _subproof(m, leaves, b):
    n = len(leaves)
    if m == n:
        return [] if b else [_mth(leaves)]
    k = 1
    while k * 2 < n:
        k *= 2
    if m <= k:
        return _subproof(m, leaves[:k], b) + [_mth(leaves[k:])]
    return _subproof(m - k, leaves[k:], False) + [_mth(leaves[:k])]


def consistency_proof(items, m):
    """PROOF(m, D[n]): what shows the tree over the first m items is a prefix of the tree over all n
    (RFC 9162 §2.1.4.1), as hex hashes; empty when m is 0 or n."""
    n = len(items)
    if not 0 <= m <= n:
        raise ValueError("m outside the tree")
    if m == 0 or m == n:
        return []
    return [h.hex() for h in _subproof(m, [_leaf(jcs.dumpb(list(it))) for it in items], True)]


def verify_consistency(root_m, root_n, m, n, proof):
    """Whether `proof` shows the tree of size m with `root_m` is a prefix of the tree of size n with `root_n`
    (RFC 9162 §2.1.4.2). Never raises."""
    try:
        if m == 0:
            return not proof and root_m == EMPTY_ROOT
        if m == n:
            return not proof and root_m == root_n
        if not 0 < m < n or not proof:
            return False
        proof = [bytes.fromhex(p) for p in proof]
        fn, sn = m - 1, n - 1
        while fn & 1:
            fn >>= 1
            sn >>= 1
        if fn == 0:
            fr = sr = bytes.fromhex(root_m)
        else:
            fr = sr = proof[0]
            proof = proof[1:]
        for c in proof:
            if sn == 0:
                return False
            if fn & 1 or fn == sn:
                fr = _node(c, fr)
                sr = _node(c, sr)
                while fn & 1 == 0 and fn != 0:
                    fn >>= 1
                    sn >>= 1
            else:
                sr = _node(sr, c)
            fn >>= 1
            sn >>= 1
        return sn == 0 and fr.hex() == root_m and sr.hex() == root_n
    except (ValueError, TypeError):
        return False


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


def ed25519_public(secret):
    """The 32-byte public key of a 32-byte Ed25519 secret."""
    return Ed25519PrivateKey.from_private_bytes(secret).public_key().public_bytes_raw()


def ed25519_sign(secret, msg):
    return Ed25519PrivateKey.from_private_bytes(secret).sign(msg)


def ed25519_verify(public, msg, signature):
    """True when `signature` is a valid Ed25519 signature of `msg` under `public`; never raises."""
    if len(public) != 32 or len(signature) != 64:
        return False
    try:
        Ed25519PublicKey.from_public_bytes(public).verify(signature, msg)
    except (InvalidSignature, ValueError):
        return False
    return True