import string
import t, c



# ==============================================
# SHA1 Hash Algorithm
# ==============================================
SHA1_BLOCK_LEN:  t.CDefine = 64
SHA1_DIGEST_LEN: t.CDefine = 20

def sha1_rotl(x: uint32, n: int) -> uint32: return (x << n) | (x >> (32 - n))
# 3-round logical functions
def sha1_f1(b: uint32, c: uint32, d: uint32) -> uint32: return (b & c) | (~b & d)
def sha1_f2(b: uint32, c: uint32, d: uint32) -> uint32: return b ^ c ^ d
def sha1_f3(b: uint32, c: uint32, d: uint32) -> uint32: return (b & c) | (b & d) | (c & d)

# SHA1 Context
@t.Object
class sha1:
    state: t.CArray[uint32, 5]
    count: uint64
    buf:   t.CArray[uint8, SHA1_BLOCK_LEN]
    def __init__(self):
        string.memset(self.buf, 0, SHA1_BLOCK_LEN)
        self.state[0] = 0x67452301
        self.state[1] = 0xefcdab89
        self.state[2] = 0x98badcfe
        self.state[3] = 0x10325476
        self.state[4] = 0xc3d2e1f0
        self.count = 0
    # Single-group compression transform
    def transform(self, block: rowptr[uint8]):
        w: t.CArray[uint32, 80]
        data: rowptr[uint8] = block
        # First 16 words
        for i in range(16):
            w[i] = (((uint32(data[i * 4 + 0])) << 24) |
                    ((uint32(data[i * 4 + 1])) << 16) |
                    ((uint32(data[i * 4 + 2])) <<  8) |
                    ((uint32(data[i * 4 + 3])) <<  0))
        # Extend to 80 words
        for i in range(16, 80):
            w[i] = sha1_rotl(w[i-3] ^ w[i-8] ^ w[i-14] ^ w[i-16],1)
        a: uint32 = self.state[0]
        b: uint32 = self.state[1]
        c: uint32 = self.state[2]
        d: uint32 = self.state[3]
        e: uint32 = self.state[4]
        k: uint32
        tmp: uint32
        # 4-round iterations
        for i in range(0, 20):
            k = 0x5a827999
            tmp = sha1_rotl(a,5) + sha1_f1(b, c, d) + e + k + w[i]
            e = d; d = c
            c = sha1_rotl(b, 30)
            b = a; a = tmp
        for i in range(20, 40):
            k = 0x6ed9eba1
            tmp = sha1_rotl(a,5) + sha1_f2(b, c, d) + e + k + w[i]
            e = d; d = c
            c = sha1_rotl(b, 30)
            b = a; a = tmp
        for i in range(40, 60):
            k = 0x8f1bbcdc
            tmp = sha1_rotl(a,5) + sha1_f3(b, c, d) + e + k + w[i]
            e = d; d = c
            c = sha1_rotl(b, 30)
            b = a; a = tmp
        for i in range(60, 80):
            k = 0xca62c1d6
            tmp = sha1_rotl(a,5) + sha1_f2(b, c, d) + e + k + w[i]
            e = d; d = c
            c = sha1_rotl(b, 30)
            b = a; a = tmp
        self.state[0] += a
        self.state[1] += b
        self.state[2] += c
        self.state[3] += d
        self.state[4] += e
    def update(self, s: str):
        data: rowptr[uint8] = uint8(s, t.CPtr)
        length: t.CSizeT = len(s)
        idx: t.CSizeT = self.count % SHA1_BLOCK_LEN
        self.count += length
        if idx > 0:
            left: t.CSizeT = SHA1_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 >= SHA1_BLOCK_LEN:
            self.transform(data)
            data += SHA1_BLOCK_LEN
            length -= SHA1_BLOCK_LEN
        if length > 0:
            string.memcpy(self.buf + idx, data, length)
    def final(self, out: t.CArray[uint8, SHA1_DIGEST_LEN]):
        idx: t.CSizeT = self.count % SHA1_BLOCK_LEN
        self.buf[idx] = 0x80
        idx += 1
        if idx > SHA1_BLOCK_LEN - 8:
            string.memset(self.buf + idx, 0, SHA1_BLOCK_LEN - idx)
            self.transform(self.buf)
            idx = 0
        string.memset(self.buf + idx,0,SHA1_BLOCK_LEN-8-idx)
        bits: uint64 = self.count << 3
        # SHA1 big-endian write length
        for i in range(8):
            self.buf[SHA1_BLOCK_LEN - 8 + i] = (bits >> (56 - i * 8)) & 0xFF
        self.transform(self.buf)
        # State registers to big-endian byte output output output
        for i in range(5):
            out[i * 4 + 0] = (self.state[i] >> 24) & 0xFF
            out[i * 4 + 1] = (self.state[i] >> 16) & 0xFF
            out[i * 4 + 2] = (self.state[i] >>  8) & 0xFF
            out[i * 4 + 3] = (self.state[i] >>  0) & 0xFF