"""Pure-Python BLAKE3 (default 32-byte hash mode only), stdlib-only (ADR-0001).

Independently written from the BLAKE3 specification (Aumasson, Neves,

O'Connor, Wilcox-O'Hearn, 2020), section 2 (compression function, chunk

chaining and binary tree). It is deliberately a separate implementation from

the Rust ``blake3`` crate Oak links, so a hash verified here is not a hash the

subject checked with its own code. Throughput is roughly 1-3 MB/s; callers

bound the input size.

"""

from __future__ import annotations

import struct

IV = (0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A,

      0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19)

CHUNK_START, CHUNK_END, PARENT, ROOT = 1, 2, 4, 8

CHUNK_LEN = 1024

BLOCK_LEN = 64

_M = 0xFFFFFFFF

# Message word schedule per round (the spec's permutation applied r times).

_PERM = (2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8)

def _schedules():

    order = list(range(16))

    result = []

    for _ in range(7):

        result.append(tuple(order))

        order = [order[p] for p in _PERM]

    return tuple(result)

_SCHEDULE = _schedules()

_WORDS = struct.Struct('<16I')

def _compress(cv, words, counter, block_len, flags):

    """Return the full 16-word output state of the compression function."""

    s0, s1, s2, s3, s4, s5, s6, s7 = cv

    s8, s9, s10, s11 = IV[0], IV[1], IV[2], IV[3]

    s12, s13, s14, s15 = counter & _M, (counter >> 32) & _M, block_len, flags

    for schedule in _SCHEDULE:

        m = [words[i] for i in schedule]

        # Column step.

        s0 = (s0 + s4 + m[0]) & _M; s12 ^= s0; s12 = ((s12 >> 16) | (s12 << 16)) & _M

        s8 = (s8 + s12) & _M; s4 ^= s8; s4 = ((s4 >> 12) | (s4 << 20)) & _M

        s0 = (s0 + s4 + m[1]) & _M; s12 ^= s0; s12 = ((s12 >> 8) | (s12 << 24)) & _M

        s8 = (s8 + s12) & _M; s4 ^= s8; s4 = ((s4 >> 7) | (s4 << 25)) & _M

        s1 = (s1 + s5 + m[2]) & _M; s13 ^= s1; s13 = ((s13 >> 16) | (s13 << 16)) & _M

        s9 = (s9 + s13) & _M; s5 ^= s9; s5 = ((s5 >> 12) | (s5 << 20)) & _M

        s1 = (s1 + s5 + m[3]) & _M; s13 ^= s1; s13 = ((s13 >> 8) | (s13 << 24)) & _M

        s9 = (s9 + s13) & _M; s5 ^= s9; s5 = ((s5 >> 7) | (s5 << 25)) & _M

        s2 = (s2 + s6 + m[4]) & _M; s14 ^= s2; s14 = ((s14 >> 16) | (s14 << 16)) & _M

        s10 = (s10 + s14) & _M; s6 ^= s10; s6 = ((s6 >> 12) | (s6 << 20)) & _M

        s2 = (s2 + s6 + m[5]) & _M; s14 ^= s2; s14 = ((s14 >> 8) | (s14 << 24)) & _M

        s10 = (s10 + s14) & _M; s6 ^= s10; s6 = ((s6 >> 7) | (s6 << 25)) & _M

        s3 = (s3 + s7 + m[6]) & _M; s15 ^= s3; s15 = ((s15 >> 16) | (s15 << 16)) & _M

        s11 = (s11 + s15) & _M; s7 ^= s11; s7 = ((s7 >> 12) | (s7 << 20)) & _M

        s3 = (s3 + s7 + m[7]) & _M; s15 ^= s3; s15 = ((s15 >> 8) | (s15 << 24)) & _M

        s11 = (s11 + s15) & _M; s7 ^= s11; s7 = ((s7 >> 7) | (s7 << 25)) & _M

        # Diagonal step.

        s0 = (s0 + s5 + m[8]) & _M; s15 ^= s0; s15 = ((s15 >> 16) | (s15 << 16)) & _M

        s10 = (s10 + s15) & _M; s5 ^= s10; s5 = ((s5 >> 12) | (s5 << 20)) & _M

        s0 = (s0 + s5 + m[9]) & _M; s15 ^= s0; s15 = ((s15 >> 8) | (s15 << 24)) & _M

        s10 = (s10 + s15) & _M; s5 ^= s10; s5 = ((s5 >> 7) | (s5 << 25)) & _M

        s1 = (s1 + s6 + m[10]) & _M; s12 ^= s1; s12 = ((s12 >> 16) | (s12 << 16)) & _M

        s11 = (s11 + s12) & _M; s6 ^= s11; s6 = ((s6 >> 12) | (s6 << 20)) & _M

        s1 = (s1 + s6 + m[11]) & _M; s12 ^= s1; s12 = ((s12 >> 8) | (s12 << 24)) & _M

        s11 = (s11 + s12) & _M; s6 ^= s11; s6 = ((s6 >> 7) | (s6 << 25)) & _M

        s2 = (s2 + s7 + m[12]) & _M; s13 ^= s2; s13 = ((s13 >> 16) | (s13 << 16)) & _M

        s8 = (s8 + s13) & _M; s7 ^= s8; s7 = ((s7 >> 12) | (s7 << 20)) & _M

        s2 = (s2 + s7 + m[13]) & _M; s13 ^= s2; s13 = ((s13 >> 8) | (s13 << 24)) & _M

        s8 = (s8 + s13) & _M; s7 ^= s8; s7 = ((s7 >> 7) | (s7 << 25)) & _M

        s3 = (s3 + s4 + m[14]) & _M; s14 ^= s3; s14 = ((s14 >> 16) | (s14 << 16)) & _M

        s9 = (s9 + s14) & _M; s4 ^= s9; s4 = ((s4 >> 12) | (s4 << 20)) & _M

        s3 = (s3 + s4 + m[15]) & _M; s14 ^= s3; s14 = ((s14 >> 8) | (s14 << 24)) & _M

        s9 = (s9 + s14) & _M; s4 ^= s9; s4 = ((s4 >> 7) | (s4 << 25)) & _M

    return (s0 ^ s8, s1 ^ s9, s2 ^ s10, s3 ^ s11, s4 ^ s12, s5 ^ s13, s6 ^ s14, s7 ^ s15,

            s8 ^ cv[0], s9 ^ cv[1], s10 ^ cv[2], s11 ^ cv[3],

            s12 ^ cv[4], s13 ^ cv[5], s14 ^ cv[6], s15 ^ cv[7])

def _block_words(block: bytes):

    if len(block) < BLOCK_LEN:

        block = block + bytes(BLOCK_LEN - len(block))

    return _WORDS.unpack(block)

def _chunk_output(chunk: bytes, counter: int):

    """(cv, words, counter, block_len, flags) of a chunk's final block, not yet finalized."""

    cv = IV

    blocks = max(1, (len(chunk) + BLOCK_LEN - 1) // BLOCK_LEN)

    for index in range(blocks - 1):

        flags = CHUNK_START if index == 0 else 0

        cv = _compress(cv, _block_words(chunk[index * BLOCK_LEN:(index + 1) * BLOCK_LEN]),

                       counter, BLOCK_LEN, flags)[:8]

    last = chunk[(blocks - 1) * BLOCK_LEN:]

    flags = CHUNK_END | (CHUNK_START if blocks == 1 else 0)

    return cv, _block_words(last), counter, len(last), flags

def _parent_output(left, right):

    return IV, tuple(left) + tuple(right), 0, BLOCK_LEN, PARENT

def digest(data: bytes) -> bytes:

    """BLAKE3 hash (32 bytes) of ``data``."""

    data = bytes(data)

    chunks = max(1, (len(data) + CHUNK_LEN - 1) // CHUNK_LEN)

    stack = []

    for index in range(chunks - 1):

        output = _chunk_output(data[index * CHUNK_LEN:(index + 1) * CHUNK_LEN], index)

        cv = _compress(*output)[:8]

        total = index + 1

        while total & 1 == 0:

            cv = _compress(*_parent_output(stack.pop(), cv))[:8]

            total >>= 1

        stack.append(cv)

    output = _chunk_output(data[(chunks - 1) * CHUNK_LEN:], chunks - 1)

    while stack:

        output = _parent_output(stack.pop(), _compress(*output)[:8])

    cv, words, counter, block_len, flags = output

    return struct.pack('<8I', *_compress(cv, words, counter, block_len, flags | ROOT)[:8])

def hexdigest(data: bytes) -> str:

    return digest(data).hex()