import string
import t, c


# ==============================================
# MD5 hash algorithm
# ==============================================
MD5_BLOCK_LEN:  t.CDefine = 64    # Group block size
MD5_DIGEST_LEN: t.CDefine = 16    # Digest output length

# MD5 four-round nonlinear functions
def md5_F(x: t.CUInt32T, y: t.CUInt32T, z: t.CUInt32T) -> t.CUInt32T: return (x & y) | (~x & z)
def md5_G(x: t.CUInt32T, y: t.CUInt32T, z: t.CUInt32T) -> t.CUInt32T: return (x & z) | (y & ~z)
def md5_H(x: t.CUInt32T, y: t.CUInt32T, z: t.CUInt32T) -> t.CUInt32T: return x ^ y ^ z
def md5_I(x: t.CUInt32T, y: t.CUInt32T, z: t.CUInt32T) -> t.CUInt32T: return y ^ (x | ~z)
# Circular left shift
def md5_rotl(x: t.CUInt32T, n: t.CInt) -> t.CUInt32T: return (x << n) | (x >> (32 - n))

# MD5 constant table
md5_T: t.CArray[t.CUInt32T, 64] = [
    0xd76aa478, 0xe8c7b756, 0x242070db, 0xc1bdceee, 0xf57c0faf, 0x4787c62a, 0xa8304613, 0xfd469501,
    0x698098d8, 0x8b44f7af, 0xffff5bb1, 0x895cd7be, 0x6b901122, 0xfd987193, 0xa679438e, 0x49b40821,
    0xf61e2562, 0xc040b340, 0x265e5a51, 0xe9b6c7aa, 0xd62f105d, 0x02441453, 0xd8a1e681, 0xe7d3fbc8,
    0x21e1cde6, 0xc33707d6, 0xf4d50d87, 0x455a14ed, 0xa9e3e905, 0xfcefa3f8, 0x676f02d9, 0x8d2a4c8a,
    0xfffa3942, 0x8771f681, 0x6d9d6122, 0xfde5380c, 0xa4beea44, 0x4bdecfa9, 0xf6bb4b60, 0xbebfbc70,
    0x289b7ec6, 0xeaa127fa, 0xd4ef3085, 0x04881d05, 0xd9d4d039, 0xe6db99e5, 0x1fa27cf8, 0xc4ac5665,
    0xf4292244, 0x432aff97, 0xab9423a7, 0xfc93a039, 0x655b59c3, 0x8f0ccc92, 0xffeff47d, 0x85845dd1,
    0x6fa87e4f, 0xfe2ce6e0, 0xa3014314, 0x4e0811a1, 0xf7537e82, 0xbd3af235, 0x2ad7d2bb, 0xeb86d391
]

md5_S: t.CArray[t.CInt, 64] = [
    7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22,
    5,  9, 14, 20, 5,  9, 14, 20, 5,  9, 14, 20, 5,  9, 14, 20,
    4, 11, 16, 23, 4, 11, 16, 23, 4, 11, 16, 23, 4, 11, 16, 23,
    6, 10, 15, 21, 6, 10, 15, 21, 6, 10, 15, 21, 6, 10, 15, 21
]

# MD5 context structure
# ============================================================
@t.Object
class md5:
    state: t.CArray[t.CUInt32T, 4]            # 4 32-bit state registers
    count: t.CUInt64T                         # processed bits
    buf:   t.CArray[t.CUInt8T, MD5_BLOCK_LEN] # Group buffer
    # Initialize MD5 context
    def __init__(self):
        string.memset(self.buf, 0, MD5_BLOCK_LEN)
        self.state[0] = 0x67452301
        self.state[1] = 0xefcdab89
        self.state[2] = 0x98badcfe
        self.state[3] = 0x10325476
        self.count = 0
    # MD5 single-group compression transform
    def transform(self, block: t.CUInt8T | t.CPtr):
        a: t.CUInt32T = self.state[0]
        b: t.CUInt32T = self.state[1]
        c: t.CUInt32T = self.state[2]
        d: t.CUInt32T = self.state[3]
        w: t.CArray[t.CUInt32T, 16]
        # Byte order 32-bit words
        for i in range(16):
            w[i] = ((t.CUInt32T(block[i * 4 + 0]) <<  0) |
                    (t.CUInt32T(block[i * 4 + 1]) <<  8) |
                    (t.CUInt32T(block[i * 4 + 2]) << 16) |
                    (t.CUInt32T(block[i * 4 + 3]) << 24))
        # 4-round 51 iterations
        for i in range(16):
            tmp: t.CUInt32T = md5_F(b, c, d) + a + md5_T[i] + w[i]
            a = d; d = c; c = b
            b += md5_rotl(tmp, md5_S[i])
        # 4-round 51 iterations
        for i in range(16, 32):
            tmp: t.CUInt32T = md5_G(b, c, d) + a + md5_T[i] + w[(5 * i + 1) % 16]
            a = d; d = c; c = b
            b += md5_rotl(tmp, md5_S[i])
        for i in range(32, 48):
            tmp: t.CUInt32T = md5_H(b, c, d) + a + md5_T[i] + w[(3 * i + 5) % 16]
            a = d; d = c; c = b
            b += md5_rotl(tmp, md5_S[i])
        for i in range(48, 64):
            tmp: t.CUInt32T = md5_I(b, c, d) + a + md5_T[i] + w[(7 * i + 0) % 16]
            a = d; d = c; c = b
            b += md5_rotl(tmp, md5_S[i])
        # Accumulate to state registers
        self.state[0] += a
        self.state[1] += b
        self.state[2] += c
        self.state[3] += d
    # Batch input data
    def update(self, s: str):
        data: t.CUInt8T | t.CPtr = t.CUInt8T(s, t.CPtr)
        length: t.CSizeT = len(s)
        idx: t.CSizeT = self.count % MD5_BLOCK_LEN
        self.count += length
        if idx > 0:
            left: t.CSizeT = MD5_BLOCK_LEN - idx
            if length >= left:
                string.memcpy(self.buf + idx, data, left)
                self.transform(self.buf)
                data += left
                length -= left
                idx = 0
        while length >= MD5_BLOCK_LEN:
            self.transform(data)
            data += MD5_BLOCK_LEN
            length -= MD5_BLOCK_LEN
        if length > 0:
            string.memcpy(self.buf + idx, data, length)
    # End calculation, output 16-byte digest
    def final(self, out: t.CArray[t.CUInt8T, MD5_DIGEST_LEN]):
        idx: t.CSizeT = self.count % MD5_BLOCK_LEN
        # Pad 0x80 end
        self.buf[idx] = 0x80
        idx += 1
        # If space is not enough, compress once
        if idx > MD5_BLOCK_LEN - 8:
            string.memset(self.buf + idx, 0, MD5_BLOCK_LEN - idx)
            self.transform(self.buf)
            idx = 0
        # Pad 0 to end 8 bytes
        string.memset(self.buf + idx, 0, MD5_BLOCK_LEN - 8 - idx)
        # Write total bit count
        bits: t.CUInt64T = self.count << 3
        for i in range(8):
            self.buf[MD5_BLOCK_LEN - 8 + i] = (bits >> (i * 8)) & 0xFF
        self.transform(self.buf)
        # State registers to little-endian byte output output
        for i in range(4):
            out[i * 4 + 0] = (self.state[i] >>  0) & 0xFF
            out[i * 4 + 1] = (self.state[i] >>  8) & 0xFF
            out[i * 4 + 2] = (self.state[i] >> 16) & 0xFF
            out[i * 4 + 3] = (self.state[i] >> 24) & 0xFF