"""BLAKE3 in hash mode, written from its specification in pure Python.

Shared by the claim-conformance verifier (`claim_verify.py`), the staged-delivery
verifier (`../staged-conformance/staged_verify.py`) and the HTTP delivery verifier
(`../http-conformance/http_verify.py`), which import from here. It shares no code with
wai-rs, which computes the same hashes with the `blake3` crate.

`subtree_cv` and `parent_cv` expose the tree (the BLAKE3 specification, §2.1): the
non-root chaining value of a subtree that starts at a given byte offset of the whole
input, and a parent node with or without the ROOT flag. The HTTP head proof is built of
these (http-delivery §6).
"""

_IV = [0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A, 0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19]
_PERM = [2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8]
_CHUNK_START, _CHUNK_END, _PARENT, _ROOT = 1, 2, 4, 8
_M32 = 0xFFFFFFFF


def _rotr(x, n):
    return ((x >> n) | (x << (32 - n))) & _M32


def _g(s, a, b, c, d, x, y):
    s[a] = (s[a] + s[b] + x) & _M32
    s[d] = _rotr(s[d] ^ s[a], 16)
    s[c] = (s[c] + s[d]) & _M32
    s[b] = _rotr(s[b] ^ s[c], 12)
    s[a] = (s[a] + s[b] + y) & _M32
    s[d] = _rotr(s[d] ^ s[a], 8)
    s[c] = (s[c] + s[d]) & _M32
    s[b] = _rotr(s[b] ^ s[c], 7)


def _compress(cv, m, counter, blen, flags):
    s = list(cv) + _IV[:4] + [counter & _M32, (counter >> 32) & _M32, blen, flags]
    m = list(m)
    for r in range(7):
        _g(s, 0, 4, 8, 12, m[0], m[1]); _g(s, 1, 5, 9, 13, m[2], m[3])
        _g(s, 2, 6, 10, 14, m[4], m[5]); _g(s, 3, 7, 11, 15, m[6], m[7])
        _g(s, 0, 5, 10, 15, m[8], m[9]); _g(s, 1, 6, 11, 12, m[10], m[11])
        _g(s, 2, 7, 8, 13, m[12], m[13]); _g(s, 3, 4, 9, 14, m[14], m[15])
        if r < 6:
            m = [m[i] for i in _PERM]
    return [s[i] ^ s[i + 8] for i in range(8)] + [s[i + 8] ^ cv[i] for i in range(8)]


def _words(block):
    block = block + b"\0" * (64 - len(block))
    return [int.from_bytes(block[i:i + 4], "little") for i in range(0, 64, 4)]


def _chunk_output(chunk, counter):
    blocks = [chunk[i:i + 64] for i in range(0, len(chunk), 64)] or [b""]
    cv = _IV
    for j, b in enumerate(blocks[:-1]):
        cv = _compress(cv, _words(b), counter, 64, _CHUNK_START if j == 0 else 0)[:8]
    last = blocks[-1]
    return (cv, _words(last), counter, len(last), _CHUNK_END | (_CHUNK_START if len(blocks) == 1 else 0))


def _cv(out):
    cv, w, c, bl, f = out
    return _compress(cv, w, c, bl, f)[:8]


def blake3(data: bytes) -> bytes:
    chunks = [data[i:i + 1024] for i in range(0, len(data), 1024)] or [b""]
    stack = []
    for idx, ch in enumerate(chunks[:-1]):
        cv = _cv(_chunk_output(ch, idx))
        total = idx + 1
        while total & 1 == 0:
            cv = _cv((_IV, stack.pop() + cv, 0, 64, _PARENT))
            total >>= 1
        stack.append(cv)
    out = _chunk_output(chunks[-1], len(chunks) - 1)
    while stack:
        out = (_IV, stack.pop() + _cv(out), 0, 64, _PARENT)
    cv, w, _, bl, f = out
    return b"".join(x.to_bytes(4, "little") for x in _compress(cv, w, 0, bl, f | _ROOT)[:8])


def _bytes(words):
    return b"".join(x.to_bytes(4, "little") for x in words)


def _left_chunks(n):
    """The left subtree of `n` chunks (n >= 2): the largest power of two below n."""
    return 1 << ((n - 1).bit_length() - 1)


def subtree_cv(data: bytes, start: int) -> bytes:
    """The non-root chaining value of `data` as a whole subtree of the input whose first
    chunk is at byte offset `start` (a multiple of 1024): chunk counters run from
    start // 1024."""
    assert start % 1024 == 0
    chunks = [data[i:i + 1024] for i in range(0, len(data), 1024)] or [b""]

    def node(cs, c0):
        if len(cs) == 1:
            return _cv(_chunk_output(cs[0], c0))
        k = _left_chunks(len(cs))
        return _cv((_IV, node(cs[:k], c0) + node(cs[k:], c0 + k), 0, 64, _PARENT))

    return _bytes(node(chunks, start // 1024))


def parent_cv(left: bytes, right: bytes, root: bool = False) -> bytes:
    """A parent node over two chaining values: its chaining value, or with `root` the
    hash it outputs as the root."""
    m = [int.from_bytes((left + right)[i:i + 4], "little") for i in range(0, 64, 4)]
    return _bytes(_compress(_IV, m, 0, 64, _PARENT | (_ROOT if root else 0))[:8])