from stdint import *
import zlib.pyzlib as pyzlib
import zlib.zdeflate as zdeflate
import zlib.zinflate as zinflate
import zlib.zdef as zdef
import zlib.zhuff as zhuff
import string
import stdlib
import testcheck
import t, c


def fill_data(data: BYTEPTR, size: sizet, seed: sizet, repeating: int):
    i: sizet
    for i in range(size):
        if repeating:
            data[i] = BYTE(65 + (i % 3))
        else:
            data[i] = BYTE((i * 31 + seed * 17 + (i * i) % 251) % 256)


def huffman_probe(seed: sizet, size: sizet, repeating: int):
    """Recompute the literal code table with zhuff for one dataset and print its
    distribution (active/no_code/sum15 and code-count histogram per bit width).
    """
    data: BYTEPTR = BYTEPTR(stdlib.malloc(size))
    fill_data(data, size, seed, repeating)
    freqs: t.CArray[int, 288]
    string.memset(freqs, 0, freqs.__sizeof__())
    i: sizet
    for i in range(size):
        freqs[data[i]] += 1
    freqs[256] = 1

    tree = zhuff.zhuff_tree()
    tree.build_codes(freqs, 257, 15)

    hist: t.CArray[int, 16]
    string.memset(hist, 0, hist.__sizeof__())
    active: int = 0
    no_code: int = 0
    sum15: int = 0
    k: int
    bits: int
    for k in range(257):
        bits = tree.codes[k].bits
        if freqs[k] > 0 and bits == 0:
            no_code += 1
        if bits > 0 and bits < 16:
            active += 1
            hist[bits] += 1
            sum15 += 1 << (15 - bits)

    testcheck.info(f"[probe seed={seed} size={size} rep={repeating}] "
                   f"active={active} no_code={no_code} sum15={sum15} (期望 32768)")
    for k in range(1, 16):
        if hist[k] > 0:
            testcheck.info(f"    bits={k} codes={hist[k]}")

    stdlib.free(data)


def count_freqs_probe(seed: sizet, size: sizet, repeating: int):
    """Call the real zdeflate_count_freqs and compare it with a naive histogram
    (prints sum_lit/active/hlit/hdist; sum_lit should equal size+1).
    """
    data: BYTEPTR = BYTEPTR(stdlib.malloc(size))
    fill_data(data, size, seed, repeating)

    lit: t.CArray[int, 288]
    dist: t.CArray[int, 32]
    string.memset(lit, 0, lit.__sizeof__())
    string.memset(dist, 0, dist.__sizeof__())

    s = zdeflate.zdeflate_create(None, pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.MAX_WBITS, 8, 0)
    if s:
        zdeflate.zdeflate_count_freqs(lit, dist, s, data, size)
        s.destroy()

    # Signedness cross-check: if UINT8PTR were treated as signed i8, bytes >=128
    # would read negative, undercounting freqs and dropping symbols -> truncated stream.
    up: UINT8PTR = UINT8PTR(data)
    ge8_up: int = 0
    ge8_bp: int = 0
    i: sizet
    for i in range(size):
        if up[i] >= 128: ge8_up += 1
        if data[i] >= 128: ge8_bp += 1
    testcheck.info(f"[countprobe] >=128 计数: UINT8PTR={ge8_up} BYTEPTR={ge8_bp} "
                   f"(两者应相等且≈一半)")

    sum_lit: int = 0
    active: int = 0
    k: int
    for k in range(288):
        sum_lit += lit[k]
        if k < 256 and lit[k] > 0:
            active += 1
    sum_dist: int = 0
    for k in range(32):
        sum_dist += dist[k]
    hlit: int = 286
    while hlit > 257 and lit[hlit - 1] == 0:
        hlit -= 1
    hdist: int = 30
    while hdist > 1 and dist[hdist - 1] == 0:
        hdist -= 1

    testcheck.info(f"[countprobe seed={seed} size={size} rep={repeating}] "
                   f"sum_lit={sum_lit} (期望 {size + 1}) active={active} "
                   f"sum_dist={sum_dist} hlit={hlit} hdist={hdist}")

    stdlib.free(data)


def roundtrip(seed: sizet, size: sizet, repeating: int, wbits: int) -> int:
    data: BYTEPTR = BYTEPTR(stdlib.malloc(size))
    fill_data(data, size, seed, repeating)

    out_len: sizet = 0
    comp: BYTEPTR = pyzlib.compress(None, data, size,
                                    pyzlib.Z_DEFAULT_COMPRESSION, wbits, &out_len)
    if not comp:
        testcheck.fail(f"seed={seed} size={size} rep={repeating} wbits={wbits}: "
                       f"compress 失败 err={pyzlib.get_error_code()} {pyzlib.get_error()}")
        stdlib.free(data)
        return -1

    dec_len: sizet = 0
    dec: BYTEPTR = pyzlib.decompress(None, comp, out_len, wbits,
                                     pyzlib.DEF_BUF_SIZE, &dec_len)
    ok: int = 1
    if not dec:
        ok = 0
        testcheck.fail(f"seed={seed} size={size} rep={repeating} wbits={wbits}: "
                       f"decompress 失败 comp={out_len} err={pyzlib.get_error_code()} {pyzlib.get_error()}")
    elif dec_len != size or string.memcmp(dec, data, size) != 0:
        ok = 0
        testcheck.fail(f"seed={seed} size={size} rep={repeating} wbits={wbits}: "
                       f"内容不匹配 comp={out_len} dec={dec_len}")
    else:
        testcheck.ok(f"seed={seed} size={size} rep={repeating} wbits={wbits}: "
                     f"comp={out_len} dec={dec_len}")

    if dec: stdlib.free(dec)
    stdlib.free(comp)
    stdlib.free(data)
    return 0 if ok else -1


def main() -> int:
    testcheck.begin("ZlibDbg: zlib 多 seed/size/wbits 往返")

    # Run the code-table probe first to locate the cause of comp below entropy floor
    testcheck.section("字面量码表探针")
    huffman_probe(0, 250, 0)
    huffman_probe(9, 20000, 0)
    huffman_probe(0, 256, 1)

    # Real frequency stats vs naive histogram comparison
    testcheck.section("频率统计探针 (zdeflate_count_freqs)")
    count_freqs_probe(0, 250, 0)
    count_freqs_probe(9, 20000, 0)
    count_freqs_probe(0, 256, 1)

    testcheck.section("随机数据 (wbits=15)")
    sizes: t.CArray[sizet, None] = [250, 255, 256, 257, 260, 300, 500, 1000, 4096, 20000]
    n: int
    for n in range(10):
        roundtrip(sizet(n), sizes[n], 0, pyzlib.MAX_WBITS)

    testcheck.section("重复数据 (wbits=15)")
    sizes2: t.CArray[sizet, None] = [256, 257, 300, 1000, 20000]
    for n in range(5):
        roundtrip(sizet(n), sizes2[n], 1, pyzlib.MAX_WBITS)

    testcheck.section("gzip 容器 (wbits=31)")
    sizes3: t.CArray[sizet, None] = [300, 1000]
    for n in range(2):
        roundtrip(sizet(n), sizes3[n], 0, pyzlib.MAX_WBITS + 16)

    return testcheck.end()