# Copyright (c) 2021 Pieter Wuille
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
#
# Pure Python RIPEMD160 implementation. Note that this impelentation is not constant time.
# Original source: https://github.com/bitcoin/bitcoin/pull/23716

import typing

# Message schedule indexes for the l path.
ML = [
    0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xa, 0xb, 0xc, 0xd, 0xe, 0xf,
    0x7, 0x4, 0xd, 0x1, 0xa, 0x6, 0xf, 0x3, 0xc, 0x0, 0x9, 0x5, 0x2, 0xe, 0xb, 0x8,
    0x3, 0xa, 0xe, 0x4, 0x9, 0xf, 0x8, 0x1, 0x2, 0x7, 0x0, 0x6, 0xd, 0xb, 0x5, 0xc,
    0x1, 0x9, 0xb, 0xa, 0x0, 0x8, 0xc, 0x4, 0xd, 0x3, 0x7, 0xf, 0xe, 0x5, 0x6, 0x2,
    0x4, 0x0, 0x5, 0x9, 0x7, 0xc, 0x2, 0xa, 0xe, 0x1, 0x3, 0x8, 0xb, 0x6, 0xf, 0xd,
]

# Message schedule indexes for the r path.
MR = [
    0x5, 0xe, 0x7, 0x0, 0x9, 0x2, 0xb, 0x4, 0xd, 0x6, 0xf, 0x8, 0x1, 0xa, 0x3, 0xc,
    0x6, 0xb, 0x3, 0x7, 0x0, 0xd, 0x5, 0xa, 0xe, 0xf, 0x8, 0xc, 0x4, 0x9, 0x1, 0x2,
    0xf, 0x5, 0x1, 0x3, 0x7, 0xe, 0x6, 0x9, 0xb, 0x8, 0xc, 0x2, 0xa, 0x0, 0x4, 0xd,
    0x8, 0x6, 0x4, 0x1, 0x3, 0xb, 0xf, 0x0, 0x5, 0xc, 0x2, 0xd, 0x9, 0x7, 0xa, 0xe,
    0xc, 0xf, 0xa, 0x4, 0x1, 0x5, 0x8, 0x7, 0x6, 0x2, 0xd, 0xe, 0x0, 0x3, 0x9, 0xb,
]

# Rotation counts for the l path.
RL = [
    0xb, 0xe, 0xf, 0xc, 0x5, 0x8, 0x7, 0x9, 0xb, 0xd, 0xe, 0xf, 0x6, 0x7, 0x9, 0x8,
    0x7, 0x6, 0x8, 0xd, 0xb, 0x9, 0x7, 0xf, 0x7, 0xc, 0xf, 0x9, 0xb, 0x7, 0xd, 0xc,
    0xb, 0xd, 0x6, 0x7, 0xe, 0x9, 0xd, 0xf, 0xe, 0x8, 0xd, 0x6, 0x5, 0xc, 0x7, 0x5,
    0xb, 0xc, 0xe, 0xf, 0xe, 0xf, 0x9, 0x8, 0x9, 0xe, 0x5, 0x6, 0x8, 0x6, 0x5, 0xc,
    0x9, 0xf, 0x5, 0xb, 0x6, 0x8, 0xd, 0xc, 0x5, 0xc, 0xd, 0xe, 0xb, 0x8, 0x5, 0x6,
]

# Rotation counts for the r path.
RR = [
    0x8, 0x9, 0x9, 0xb, 0xd, 0xf, 0xf, 0x5, 0x7, 0x7, 0x8, 0xb, 0xe, 0xe, 0xc, 0x6,
    0x9, 0xd, 0xf, 0x7, 0xc, 0x8, 0x9, 0xb, 0x7, 0x7, 0xc, 0x7, 0x6, 0xf, 0xd, 0xb,
    0x9, 0x7, 0xf, 0xb, 0x8, 0x6, 0x6, 0xe, 0xc, 0xd, 0x5, 0xe, 0xd, 0xd, 0x7, 0x5,
    0xf, 0x5, 0x8, 0xb, 0xe, 0xe, 0x6, 0xe, 0x6, 0x9, 0xc, 0x9, 0xc, 0x5, 0xf, 0x8,
    0x8, 0x5, 0xc, 0x9, 0xc, 0x5, 0xe, 0x6, 0x8, 0xd, 0x6, 0x5, 0xf, 0xd, 0xb, 0xb,
]

# K constants for the l path.
KL = [0x00000000, 0x5a827999, 0x6ed9eba1, 0x8f1bbcdc, 0xa953fd4e]

# K constants for the r path.
KR = [0x50a28be6, 0x5c4dd124, 0x6d703ef3, 0x7a6d76e9, 0x00000000]


def rol(x: int, i: int) -> int:
    # Rotate the bottom 32 bits of x left by i bits.
    assert x <= 0xffffffff
    assert i <= 32
    return u32((x << i) | (x >> (32 - i)))


def u32(x: int) -> int:
    return x & 0xffffffff


def function(x: int, y: int, z: int, i: int) -> int:
    # The f1, f2, f3, f4, and f5 functions from the specification.
    assert x <= 0xffffffff
    assert y <= 0xffffffff
    assert z <= 0xffffffff
    assert i <= 4
    if i == 0:
        return x ^ y ^ z
    if i == 1:
        return (x & y) | (~x & z)
    if i == 2:
        return (x | ~y) ^ z
    if i == 3:
        return (x & z) | (y & ~z)
    if i == 4:
        return x ^ (y | ~z)
    raise Exception('unreachable')


def compress(state: typing.List[int], block: typing.List[int]):
    # Compress state with block."""
    # L path variables.
    h0 = state[0]
    h1 = state[1]
    h2 = state[2]
    h3 = state[3]
    h4 = state[4]
    al, bl, cl, dl, el = h0, h1, h2, h3, h4
    # R path variables.
    ar, br, cr, dr, er = h0, h1, h2, h3, h4
    # Message variables.
    x = [int.from_bytes(block[4*i:4*(i+1)], 'little') for i in range(16)]
    # Iterate over the 80 rounds of the compression.
    for j in range(80):
        rn = j >> 4
        # Perform l side of the transformation.
        al = u32(rol(u32(al + function(bl, cl, dl, 0 + rn) + x[ML[j]] + KL[rn]), RL[j]) + el)
        al, bl, cl, dl, el = el, al, bl, rol(cl, 10), dl
        # Perform r side of the transformation.
        ar = u32(rol(u32(ar + function(br, cr, dr, 4 - rn) + x[MR[j]] + KR[rn]), RR[j]) + er)
        ar, br, cr, dr, er = er, ar, br, rol(cr, 10), dr
    # Compose old state, left transform, and right transform into new state.
    state[0] = u32(h1 + cl + dr)
    state[1] = u32(h2 + dl + er)
    state[2] = u32(h3 + el + ar)
    state[3] = u32(h4 + al + br)
    state[4] = u32(h0 + bl + cr)


class Ripemd160:
    # Ripemd160 hasher.

    def __init__(self):
        self.cache = bytearray()
        self.count = 0
        # Initialize state.
        self.state = [0x67452301, 0xefcdab89, 0x98badcfe, 0x10325476, 0xc3d2e1f0]

    def update(self, data: bytearray | bytes):
        self.cache.extend(data)
        self.count += len(data) * 8
        for _ in range(1 << 32):
            if len(self.cache) < 64:
                break
            full_block = self.cache[0x00:0x40]
            self.cache = self.cache[0x40:]
            # Process full 64-byte blocks in the input.
            compress(self.state, list(full_block))
        return self

    def digest(self) -> bytearray:
        # Construct final blocks (with padding and size).
        self.cache.append(0x80)
        size = len(self.cache)
        padn = 0x38 if size <= 56 else 0x78
        self.cache.extend(bytearray(padn - size))
        self.cache.extend(bytearray(self.count.to_bytes(8, 'little')))
        # Process final blocks.
        self.update(bytearray())
        assert len(self.cache) == 0
        # Produce output.
        r = bytearray()
        for h in self.state:
            r.extend(bytearray(h.to_bytes(4, 'little')))
        return r


def ripemd160(data: bytearray = bytearray()) -> Ripemd160:
    # Returns a ripemd160 hash object; optionally initialized with a string.
    assert isinstance(data, bytearray)
    hash = Ripemd160()
    hash.update(data)
    return hash
