from stdint import *
import stddef
import string
import stdlib
import memhub
import t, c


# ============================================================
# DEFLATE / zlib constants
# ============================================================
ZDEFLATE_WINDOW_BITS: t.CDefine = 15
ZDEFLATE_WINDOW_SIZE: t.CDefine =    (1 << ZDEFLATE_WINDOW_BITS)
ZDEFLATE_WINDOW_MASK: t.CDefine =    (ZDEFLATE_WINDOW_SIZE - 1)

ZDEFLATE_MIN_MATCH: t.CDefine =      3
ZDEFLATE_MAX_MATCH: t.CDefine =      258

ZDEFLATE_HASH_BITS: t.CDefine =      15
ZDEFLATE_HASH_SIZE: t.CDefine =      (1 << ZDEFLATE_HASH_BITS)
ZDEFLATE_HASH_MASK: t.CDefine =      (ZDEFLATE_HASH_SIZE - 1)

ZDEFLATE_MAX_CODES:         t.CDefine =      288
ZDEFLATE_MAX_DIST_CODES:    t.CDefine = 32
ZDEFLATE_MAX_BITS:          t.CDefine =       15
ZDEFLATE_MAX_CODELEN_CODES: t.CDefine = 19

ZDEFLATE_LIT_COUNT:  t.CDefine =      286
ZDEFLATE_DIST_COUNT: t.CDefine =     30

ZDEFLATE_LEN_SYMBOLS_BASE: t.CDefine = 257
ZDEFLATE_END_OF_BLOCK:     t.CDefine =     256

# ============================================================
# Length / Distance extra bits tables (RFC 1951)
# ============================================================
zdeflate_len_extra_bits: t.CArray[t.CInt, 29] = [
    0, 0, 0, 0, 0, 0, 0, 0,
    1, 1, 1, 1,
    2, 2, 2, 2,
    3, 3, 3, 3,
    4, 4, 4, 4,
    5, 5, 5, 5,
    0
]

zdeflate_len_base: t.CArray[t.CInt, 29] = [
    3, 4, 5, 6, 7, 8, 9, 10,
    11, 13, 15, 17,
    19, 23, 27, 31,
    35, 43, 51, 59,
    67, 83, 99, 115,
    131, 163, 195, 227,
    258
]

zdeflate_dist_extra_bits: t.CArray[t.CInt, 30] = [
    0, 0, 0, 0,
    1, 1,
    2, 2,
    3, 3,
    4, 4,
    5, 5,
    6, 6,
    7, 7,
    8, 8,
    9, 9,
    10, 10,
    11, 11,
    12, 12,
    13, 13
]

zdeflate_dist_base: t.CArray[t.CInt, 30] = [
    1, 2, 3, 4,
    5, 7,
    9, 13,
    17, 25,
    33, 49,
    65, 97,
    129, 193,
    257, 385,
    513, 769,
    1025, 1537,
    2049, 3073,
    4097, 6145,
    8193, 12289,
    16385, 24577
]

zdeflate_codelen_order: t.CArray[int, 19] = [
    16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15
]

# ============================================================
# Bit writer - writes bits LSB first into a byte buffer
# ============================================================
@t.Object
class zbit_writer:
    pool: memhub.MemManager | t.CPtr
    buf: BYTE | t.CPtr
    cap: t.CSizeT
    byte_pos: t.CSizeT
    bit_pos: t.CInt

    def __init__(self, pool: memhub.MemManager | t.CPtr):
        self.pool = pool
        self.cap = 4096
        if pool == None:
            self.buf = BYTEPTR(malloc(self.cap))
        else:
            self.buf = BYTEPTR(pool.alloc(self.cap))
        memset(self.buf, 0, self.cap)
        self.byte_pos = 0
        self.bit_pos = 0

    def ensure(self, need: t.CSizeT):
        while self.byte_pos + need + 4 >= self.cap:
            new_cap: t.CSizeT = self.cap * 2
            if self.pool == None:
                new_buf: BYTE | t.CPtr = BYTEPTR(realloc(self.buf, new_cap))
            else:
                new_buf: BYTE | t.CPtr = BYTEPTR(self.pool.alloc(new_cap))
                if new_buf:
                    memcpy(new_buf, self.buf, self.byte_pos + (self.bit_pos + 7) / 8)
                    self.pool.free(self.buf)
            if not new_buf: return
            memset(new_buf + self.cap, 0, new_cap - self.cap)
            self.buf = new_buf
            self.cap = new_cap

    def write_bits(self, value: UINT, nbits: t.CInt):
        self.ensure((nbits + 7) / 8 + 1)
        for i in range(nbits):
            if value & (UINT(1) << i):
                self.buf[self.byte_pos] |= (UINT(1) << self.bit_pos)
            self.bit_pos += 1
            if self.bit_pos == 8:
                self.bit_pos = 0
                self.byte_pos += 1

    def write_bits_rev(self, code: UINT, nbits: t.CInt):
        self.ensure((nbits + 7) / 8 + 1)
        for i in range(nbits - 1, -1, -1):
            if code & (UINT(1) << i):
                self.buf[self.byte_pos] |= (UINT(1) << self.bit_pos)
            self.bit_pos += 1
            if self.bit_pos == 8:
                self.bit_pos = 0
                self.byte_pos += 1

    def stored_block(self, data: UINT8PTR, length: t.CSizeT, final: t.CInt):
        self.write_bits(1 if final else 0, 1)
        self.write_bits(0, 2)
        self.align()

        pos: t.CSizeT = 0
        while pos < length:
            block_length: t.CSizeT = length - pos
            if block_length > 65535:
                block_length = 65535
            n: t.CUInt16T = t.CUInt16T(block_length)
            ncomp: t.CUInt16T = ~n
            self.write_bits(n, 16)
            self.write_bits(ncomp, 16)

            self.ensure(block_length)
            i: t.CSizeT = 0
            for i in range(block_length):
                self.buf[self.byte_pos] = data[pos + i]
                self.byte_pos += 1
            pos += block_length
            if pos < length:
                self.write_bits(0, 1)
                self.write_bits(0, 2)
                self.align()

    def write_zlib_header(self, wbits: t.CInt, level: t.CInt):
        cinfo: t.CInt = wbits - 8
        if cinfo < 1: cinfo = 1
        if cinfo > 7: cinfo = 7
        cmf: t.CUnsignedChar = t.CUnsignedChar((cinfo << 4) | 8)

        flevel: t.CInt = 0
        if   level >= 5: flevel = 3
        elif level >= 3: flevel = 1

        flg: t.CUnsignedChar = t.CUnsignedChar(flevel << 6)
        flg |= t.CUnsignedChar(31 - (cmf * 256 + flg) % 31)

        self.buf[self.byte_pos + 0] = cmf
        self.buf[self.byte_pos + 1] = flg
        self.byte_pos += 2


    def write_gzip_header(self):
        self.buf[self.byte_pos + 0] = 0x1F
        self.buf[self.byte_pos + 1] = 0x8B
        self.buf[self.byte_pos + 2] = 0x08
        self.buf[self.byte_pos + 3] = 0x00
        self.buf[self.byte_pos + 4] = 0x00
        self.buf[self.byte_pos + 5] = 0x00
        self.buf[self.byte_pos + 6] = 0x00
        self.buf[self.byte_pos + 7] = 0x00
        self.buf[self.byte_pos + 8] = 0x00
        self.buf[self.byte_pos + 9] = 0xFF
        self.byte_pos += 10

    def align(self):
        if self.bit_pos > 0:
            self.bit_pos = 0
            self.byte_pos += 1

    def total(self) -> t.CSizeT:
        return self.byte_pos + (1 if (self.bit_pos > 0) else 0)

    def free(self):
        if self.pool == None:
            free(self.buf)
        else:
            self.pool.free(self.buf)
        self.buf = None
        self.cap = 0
        self.byte_pos = 0
        self.bit_pos = 0

    def __del__(self):
        self.free()

# ============================================================
# Bit reader - reads bits LSB first from a byte buffer
# ============================================================
@t.Object
class zbit_reader:
    buf: BYTE | t.CPtr
    length: t.CSizeT
    byte_pos: t.CSizeT
    bit_pos: t.CInt

    def init(self, buf: BYTE | t.CPtr, length: t.CSizeT):
        self.buf = buf
        self.length = length
        self.byte_pos = 0
        self.bit_pos = 0

    def read_bits(self, nbits: t.CInt, out: UINTPTR) -> t.CInt:
        val: UINT = 0
        for i in range(nbits):
            if self.byte_pos >= self.length: return -1
            if self.buf[self.byte_pos] & (UINT(1) << self.bit_pos):
                val |= (UINT(1) << i)
            self.bit_pos += 1
            if self.bit_pos == 8:
                self.bit_pos = 0
                self.byte_pos += 1
        c.Set(c.Deref(out), val)
        return 0

    def read_bits_rev(self, nbits: t.CInt, out: UINTPTR) -> t.CInt:
        val: UINT = 0
        for i in range(nbits - 1, 0, -1):
            if self.byte_pos >= self.length: return -1
            if self.buf[self.byte_pos] & (UINT(1) << self.bit_pos):
                val |= (UINT(1) << i)
            self.bit_pos += 1
            if self.bit_pos == 8:
                self.bit_pos = 0
                self.byte_pos += 1
        c.Set(c.Deref(out), val)
        return 0

    def align(self):
        if self.bit_pos > 0:
            self.bit_pos = 0
            self.byte_pos += 1

# ============================================================
# Memory allocation helper for bare metal
# Can be replaced with custom allocator
# ============================================================
def zdef_alloc(pool: memhub.MemManager | t.CPtr, n: t.CSizeT) -> VOIDPTR:
    if pool == None: return malloc(n)
    return pool.alloc(n)
def zdef_free(pool: memhub.MemManager | t.CPtr, p: VOIDPTR):
    if pool == None:
        free(p)
        return
    pool.free(p)
