# -*- coding: utf-8 -*-
# ============================================================
# gzip - single-file gzip decoder (RFC 1952 + RFC 1951 deflate)
#
# OOP paradigm aligned with zipfile: data is held and freed by the GzipFile object;
# internal bit-stream/Huffman helpers are module-level private functions.
#
# Usage:
#   gz: rawptr[gzip.GzipFile] = gzip.GzipFile("train-images-idx3-ubyte.gz")
#   if gz is None or gz.ok == 0: ...
#   buf: bytes = gz.read_all()     # internally malloc'd, decompression is complete when read_all returns
#   n: sizet = gz.dst_len       # original length (held by the object, freed on close)
#   gz.close()
#
# Covers the three deflate block types (stored/fixed/dynamic); does not decode CRC.
# ============================================================
import stdlib
import string
import win32.fileio as fileio
from stdint import *


# ============================================================
# GzipFile
# ============================================================
class GzipFile:
    ok: int                    # 1 = usable
    err: int                   # error code (0 = none, <0 = error)
    src: bytes                  # in-memory copy of whole .gz file (read in ctor, freed by close)
    src_len: sizet
    dst: bytes                  # decompressed output buffer (created by read_all, freed by close)
    dst_len: sizet             # decompressed length
    _need_close_src: int       # whether src is owned (0 if read failed)

    def __init__(self, path: str):
        self.ok = 0
        self.err = 0
        self.src = None
        self.src_len = 0
        self.dst = None
        self.dst_len = 0
        self._need_close_src = 0
        if path is None:
            self.err = -1
            return
        f: fileio.File = fileio.File(path, fileio.MODE.R)
        if f.closed:
            self.err = -2       # cannot open file
            return
        sz: longlong = f.size()
        if sz < 18:
            f.close()
            self.err = -3       # file too small
            return
        self.src_len = sizet(sz)
        # read_all contract: reads up to max_count-1 bytes and NUL-terminates; request src_len+1.
        buf: bytes = stdlib.malloc(self.src_len + 1)
        if buf is None:
            f.close()
            self.err = -4
            return
        got: LONG = f.read_all(buf, self.src_len + 1)    # 去掉 ULONG 显式窄化：read_all 的 buf_size 形参本就是 sizet，直接以 sizet 传更安全
        f.close()
        if got <= 0 or sizet(got) != self.src_len:
            stdlib.free(buf)
            self.err = -5       # read failed or truncated
            return
        # gzip magic-number check
        ub: uchar* = buf
        if ub[0] != 0x1F or ub[1] != 0x8B or ub[2] != 8:
            stdlib.free(buf)
            self.err = -6
            return
        self.src = buf
        self._need_close_src = 1
        self.ok = 1

    def read_all(self) -> bytes:
        """Decompress the whole member. Returns internal output buffer (or None); length in self.dst_len."""
        if self.ok == 0 or self.src is None:
            return None
        if self.dst is not None:
            return self.dst          # already decompressed
        want: uint32 = gzip_isize(self.src, self.src_len)
        cap: sizet = sizet(want)
        if cap == 0:
            cap = self.src_len * 64  # conservative initial value when no ISIZE
        if cap < 1024:
            cap = 1024
        # If ISIZE mismatches actual size (buffer-too-small), double cap and retry (normally one pass).
        out: BYTEPTR = None
        n: LONG = 0
        while True:
            out = stdlib.malloc(cap)
            if out is None:
                self.err = -7
                return None
            n = _gzip_decode(self.src, self.src_len, out, cap)
            if n == -9 or n == -18 or n == -23:
                stdlib.free(out)
                out = None
                if cap >= 1073741824:
                    self.err = int(n)
                    return None
                cap = cap * 2
                continue
            break
        if n <= 0:
            stdlib.free(out)
            self.err = int(n)   # n is a negative error code; store as-is
            return None
        self.dst = out
        self.dst_len = sizet(n)
        return out

    def close(self):
        if self._need_close_src != 0 and self.src is not None:
            stdlib.free(self.src)
        if self.dst is not None:
            stdlib.free(self.dst)
        self.src = None
        self.dst = None
        self.src_len = 0
        self.dst_len = 0
        self._need_close_src = 0
        self.ok = 0

# ============================================================
# Bit-stream state (module-level, single-thread; reset per _gzip_decode)
# ============================================================
_g_bits: uint64 = 0
_g_bitcnt: int = 0
_g_pos: sizet = 0
_g_src: rawptr = None
_g_src_len: sizet = 0


def gzip_isize(src: rawptr, src_len: sizet) -> uint32:
    """gzip trailing ISIZE (little-endian u32): original length hint."""
    if src is None or src_len < 4:
        return 0
    ub: uchar* = src
    p: sizet = src_len - 4
    # Assemble little-endian u32 (build in 64-bit then narrow)
    v64: uint64 = uint64(ub[p]) \
        + uint64(ub[p + 1]) * 256 \
        + uint64(ub[p + 2]) * 65536 \
        + uint64(ub[p + 3]) * 16777216
    return uint32(v64)


def _reset(src: rawptr, src_len: sizet):
    global _g_bits, _g_bitcnt, _g_pos, _g_src, _g_src_len
    _g_bits = 0
    _g_bitcnt = 0
    _g_pos = 0
    _g_src = src
    _g_src_len = src_len


def _fill(n: int) -> int:
    global _g_bits, _g_bitcnt, _g_pos, _g_src, _g_src_len
    sb: uchar* = _g_src
    while _g_bitcnt < n:
        if _g_pos >= _g_src_len:
            return -1
        _g_bits = _g_bits | (uint64(sb[_g_pos]) << _g_bitcnt)
        _g_pos += 1
        _g_bitcnt += 8
    return 0


def _bits(n: int) -> int:
    global _g_bits, _g_bitcnt
    if _fill(int(n)) != 0:
        return -1
    v: uint64 = _g_bits & ((uint64(1) << n) - 1)
    _g_bits = _g_bits >> n
    _g_bitcnt -= n
    return int(v)


def _byte() -> int:
    return _bits(8)


def _align():
    global _g_bits, _g_bitcnt
    _g_bits = 0
    _g_bitcnt = 0


def _decode(count: int*, first: int*,
            off: int*, sorted: int*) -> int:
    code: int = 0
    ln: int = 1
    while ln <= 15:
        b: int = _bits(1)
        if b < 0:
            return -1
        code = (code << 1) | b
        if count[ln] != 0:
            d: int = code - first[ln]
            if d >= 0 and d < count[ln]:
                return sorted[off[ln] + d]
        ln += 1
    return -1


def _build(first: int*, off: int*,
           count: int*) -> int:
    anylen: int = 0
    l1: int = 1
    while l1 <= 15:
        if count[l1] != 0:
            anylen = 1
        l1 += 1
    if anylen == 0:
        return -1
    left: int = 1
    l2: int = 1
    while l2 <= 15:
        left = left << 1
        if count[l2] > left:
            return -1
        left = left - count[l2]
        l2 += 1
    code: int = 0
    l3: int = 1
    while l3 <= 15:
        code = (code + count[l3 - 1]) << 1
        first[l3] = code
        l3 += 1
    offsum: int = 0
    l4: int = 1
    while l4 <= 15:
        off[l4] = offsum
        offsum += count[l4]
        l4 += 1
    return 0


def _collect(len_arr: int*, cnt: int*,
             sorted_arr: int*, nsym: int,
             pos: int*):
    string.memset(cnt, 0, 17 * 4)
    for s1 in range(int(nsym)):
        if len_arr[s1] != 0:
            cnt[len_arr[s1]] += 1
    run: int = 0
    l5: int = 1
    while l5 <= 15:
        pos[l5] = run
        run += cnt[l5]
        l5 += 1
    for s2 in range(int(nsym)):
        if len_arr[s2] != 0:
            sorted_arr[pos[len_arr[s2]]] = s2
            pos[len_arr[s2]] += 1


def _len_extra(code: int) -> int:
    if code == 285:
        return 0
    if code >= 281:
        return 5
    if code >= 277:
        return 4
    if code >= 273:
        return 3
    if code >= 269:
        return 2
    if code >= 265:
        return 1
    return 0


def _len_base(code: int) -> int:
    if code == 285:
        return 258
    if code >= 281:
        return 131 + ((code - 281) << 5)
    if code >= 277:
        return 67 + ((code - 277) << 4)
    if code >= 273:
        return 35 + ((code - 273) << 3)
    if code >= 269:
        return 19 + ((code - 269) << 2)
    if code >= 265:
        return 11 + ((code - 265) << 1)
    return 3 + (code - 257)


def _dist_extra(code: int) -> int:
    if code < 4:
        return 0
    return (code - 4) // 2 + 1


def _dist_base(code: int) -> int:
    if code < 4:
        return code + 1
    e: int = (code - 4) // 2 + 1
    f: int = (1 << (e + 1)) + 1
    if code % 2 == 0:
        return f
    return f + (1 << e)


def _cl_order_set(arr: int*, idx: int, val: int):
    arr[idx] = val


# ============================================================
# Core decompression (module-private)
# ============================================================
def _gzip_decode(src: rawptr, src_len: sizet,
                 dst: rawptr, dst_cap: sizet) -> LONG:
    if src is None or dst is None or src_len < 18 or dst_cap == 0:
        return -1
    ub: uchar* = src
    if ub[0] != 0x1F or ub[1] != 0x8B or ub[2] != 8:
        return -2
    _reset(src, src_len)
    flg: int = int(ub[3])
    hdr_len: sizet = 10
    if flg & 0x04 != 0:
        xlen: sizet = sizet(ub[hdr_len]) | (sizet(ub[hdr_len + 1]) << 8)
        hdr_len += 2 + xlen
    if flg & 0x08 != 0:
        while hdr_len < src_len and ub[hdr_len] != 0:
            hdr_len += 1
        hdr_len += 1
    if flg & 0x10 != 0:
        while hdr_len < src_len and ub[hdr_len] != 0:
            hdr_len += 1
        hdr_len += 1
    if flg & 0x02 != 0:
        hdr_len += 2
    if hdr_len >= src_len:
        return -3
    global _g_pos
    _g_pos = hdr_len

    dout: uchar* = dst
    wpos: sizet = 0

    lit_len: int* = stdlib.malloc(288 * 4)
    lit_cnt: int* = stdlib.malloc(17 * 4)
    lit_first: int* = stdlib.malloc(17 * 4)
    lit_off: int* = stdlib.malloc(17 * 4)
    lit_sorted: int* = stdlib.malloc(288 * 4)
    dst_len: int* = stdlib.malloc(32 * 4)
    dst_cnt: int* = stdlib.malloc(17 * 4)
    dst_first: int* = stdlib.malloc(17 * 4)
    dst_off: int* = stdlib.malloc(17 * 4)
    dst_sorted: int* = stdlib.malloc(32 * 4)
    cl_len: int* = stdlib.malloc(19 * 4)
    cl_order: int* = stdlib.malloc(19 * 4)
    cl_cnt: int* = stdlib.malloc(17 * 4)
    cl_first: int* = stdlib.malloc(17 * 4)
    cl_off: int* = stdlib.malloc(17 * 4)
    cl_sorted: int* = stdlib.malloc(19 * 4)
    pos_tmp: int* = stdlib.malloc(17 * 4)
    if lit_len is None or lit_cnt is None or lit_first is None \
            or lit_off is None or lit_sorted is None or dst_len is None \
            or dst_cnt is None or dst_first is None or dst_off is None \
            or dst_sorted is None or cl_len is None or cl_order is None \
            or cl_cnt is None or cl_first is None or cl_off is None \
            or cl_sorted is None or pos_tmp is None:
        return -4

    # Zero all table buffers (malloc is uninitialized; avoids garbage tables causing early/bad finish).
    string.memset(lit_len, 0, 288 * 4)
    string.memset(lit_cnt, 0, 17 * 4)
    string.memset(lit_first, 0, 17 * 4)
    string.memset(lit_off, 0, 17 * 4)
    string.memset(lit_sorted, 0, 288 * 4)
    string.memset(dst_len, 0, 32 * 4)
    string.memset(dst_cnt, 0, 17 * 4)
    string.memset(dst_first, 0, 17 * 4)
    string.memset(dst_off, 0, 17 * 4)
    string.memset(dst_sorted, 0, 32 * 4)
    string.memset(cl_len, 0, 19 * 4)
    string.memset(cl_order, 0, 19 * 4)
    string.memset(cl_cnt, 0, 17 * 4)
    string.memset(cl_first, 0, 17 * 4)
    string.memset(cl_off, 0, 17 * 4)
    string.memset(cl_sorted, 0, 19 * 4)
    string.memset(pos_tmp, 0, 17 * 4)

    _cl_order_set(cl_order, 0, 16)
    _cl_order_set(cl_order, 1, 17)
    _cl_order_set(cl_order, 2, 18)
    _cl_order_set(cl_order, 3, 0)
    _cl_order_set(cl_order, 4, 8)
    _cl_order_set(cl_order, 5, 7)
    _cl_order_set(cl_order, 6, 9)
    _cl_order_set(cl_order, 7, 6)
    _cl_order_set(cl_order, 8, 10)
    _cl_order_set(cl_order, 9, 5)
    _cl_order_set(cl_order, 10, 11)
    _cl_order_set(cl_order, 11, 4)
    _cl_order_set(cl_order, 12, 12)
    _cl_order_set(cl_order, 13, 3)
    _cl_order_set(cl_order, 14, 13)
    _cl_order_set(cl_order, 15, 2)
    _cl_order_set(cl_order, 16, 14)
    _cl_order_set(cl_order, 17, 1)
    _cl_order_set(cl_order, 18, 15)

    err: int = 0
    final: int = 0

    while final == 0 and err == 0:
        bfinal: int = _bits(1)
        btype: int = _bits(2)
        if bfinal < 0 or btype < 0:
            err = -5
            break
        final = bfinal
        if btype == 0:
            _align()
            slen: int = _bits(16)
            snlen: int = _bits(16)
            if slen < 0 or snlen < 0:
                err = -6
                break
            if (slen ^ snlen) != 0xFFFF:
                err = -7
                break
            for i in range(int(sizet(slen))):
                ch: int = _byte()
                if ch < 0:
                    err = -8
                    break
                if wpos >= dst_cap:
                    err = -9
                    break
                dout[wpos] = uchar(ch)
                wpos += 1
            if err != 0:
                break
        elif btype == 1:
            string.memset(lit_cnt, 0, 17 * 4)
            j: int = 0
            while j < 144:
                lit_len[j] = 8
                j += 1
            while j < 256:
                lit_len[j] = 9
                j += 1
            while j < 280:
                lit_len[j] = 7
                j += 1
            while j < 288:
                lit_len[j] = 8
                j += 1
            _collect(lit_len, lit_cnt, lit_sorted, 288, pos_tmp)
            if _build(lit_first, lit_off, lit_cnt) != 0:
                err = -10
                break
        else:
            hlit: int = _bits(5) + 257
            hdist: int = _bits(5) + 1
            hclen: int = _bits(4) + 4
            if hlit < 0 or hdist < 0 or hclen < 0:
                err = -11
                break
            string.memset(cl_len, 0, 19 * 4)
            for k in range(int(hclen)):
                vv: int = _bits(3)
                if vv < 0:
                    err = -12
                    break
                cl_len[cl_order[k]] = vv
            if err != 0:
                break
            _collect(cl_len, cl_cnt, cl_sorted, 19, pos_tmp)
            if _build(cl_first, cl_off, cl_cnt) != 0:
                err = -13
                break
            string.memset(lit_len, 0, 288 * 4)
            string.memset(dst_len, 0, 32 * 4)
            sym_i: int = 0
            while sym_i < hlit + hdist:
                sym: int = _decode(cl_cnt, cl_first, cl_off, cl_sorted)
                if sym < 0:
                    err = -14
                    break
                if sym < 16:
                    if sym_i < hlit:
                        lit_len[sym_i] = sym
                    else:
                        dst_len[sym_i - hlit] = sym
                    sym_i += 1
                elif sym == 16:
                    prev: int = 0
                    if sym_i > 0:
                        if sym_i - 1 < hlit:
                            prev = lit_len[sym_i - 1]
                        else:
                            prev = dst_len[sym_i - hlit - 1]
                    rp: int = _bits(2) + 3
                    r2: int = 0
                    while r2 < rp and sym_i < hlit + hdist:
                        if sym_i < hlit:
                            lit_len[sym_i] = prev
                        else:
                            dst_len[sym_i - hlit] = prev
                        sym_i += 1
                        r2 += 1
                elif sym == 17:
                    rp2: int = _bits(3) + 3
                    r3: int = 0
                    while r3 < rp2 and sym_i < hlit + hdist:
                        if sym_i < hlit:
                            lit_len[sym_i] = 0
                        else:
                            dst_len[sym_i - hlit] = 0
                        sym_i += 1
                        r3 += 1
                else:
                    rp3: int = _bits(7) + 11
                    r4: int = 0
                    while r4 < rp3 and sym_i < hlit + hdist:
                        if sym_i < hlit:
                            lit_len[sym_i] = 0
                        else:
                            dst_len[sym_i - hlit] = 0
                        sym_i += 1
                        r4 += 1
            if err != 0:
                break
            _collect(lit_len, lit_cnt, lit_sorted, hlit, pos_tmp)
            if _build(lit_first, lit_off, lit_cnt) != 0:
                err = -15
                break
            _collect(dst_len, dst_cnt, dst_sorted, hdist, pos_tmp)
            if _build(dst_first, dst_off, dst_cnt) != 0:
                err = -16
                break
        if err != 0:
            break

        run: int = 1
        while run == 1 and err == 0:
            sym: int = _decode(lit_cnt, lit_first, lit_off, lit_sorted)
            if sym < 0:
                err = -17
                break
            if sym < 256:
                if wpos >= dst_cap:
                    err = -18
                    break
                dout[wpos] = uchar(sym)
                wpos += 1
            elif sym == 256:
                run = 0
            else:
                le: int = _len_extra(sym)
                ev: int = _bits(le)
                if ev < 0:
                    err = -19
                    break
                length: sizet = sizet(_len_base(sym) + ev)
                ds: int = _decode(dst_cnt, dst_first, dst_off, dst_sorted)
                if ds < 0:
                    err = -20
                    break
                de: int = _dist_extra(ds)
                dv: int = _bits(de)
                if dv < 0:
                    err = -21
                    break
                dist: sizet = sizet(_dist_base(ds) + dv)
                if dist == 0 or dist > wpos:
                    err = -22
                    break
                for ci in range(int(length)):
                    if wpos >= dst_cap:
                        err = -23
                        break
                    dout[wpos] = dout[wpos - dist]
                    wpos += 1
                if err != 0:
                    break

    stdlib.free(lit_len)
    stdlib.free(lit_cnt)
    stdlib.free(lit_first)
    stdlib.free(lit_off)
    stdlib.free(lit_sorted)
    stdlib.free(dst_len)
    stdlib.free(dst_cnt)
    stdlib.free(dst_first)
    stdlib.free(dst_off)
    stdlib.free(dst_sorted)
    stdlib.free(cl_len)
    stdlib.free(cl_order)
    stdlib.free(cl_cnt)
    stdlib.free(cl_first)
    stdlib.free(cl_off)
    stdlib.free(cl_sorted)
    stdlib.free(pos_tmp)

    if err != 0:
        return LONG(err)   # err is already a negative code; do not negate it (else callers treat >0 as success)
    return LONG(wpos)