from stdint import *
import zlib.pyzlib as pyzlib
import zlib.zhuff as zhuff
import zlib.zdef as zdef
from stdio import printf
from string import strlen, memcmp, memcpy, memset
from stdlib import malloc, free, realloc
import stdlib
import string
import memhub
import t, c

pool: t.CPtr = None

test_passed: t.CInt = 0
test_failed: t.CInt = 0
def check(name: str, condition: t.CInt, detail: str):
    global test_passed, test_failed
    if condition:
        test_passed += 1
        printf("  [PASS] %s\n", name)
    else:
        test_failed += 1
        printf("  [FAIL] %s  -- %s\n", name, detail if detail else "")

def print_hex_test(data: BYTEPTR, length: t.CSizeT, max_show: t.CSizeT):
    if not data or length == 0:
        printf("(empty)\n")
        return
    show: t.CSizeT = max_show if (length > max_show) else length
    i: t.CSizeT
    for i in range(show):
        printf("%02X ", data[i])
    if length > max_show:
        printf("... (%zu bytes total)", length)
    printf("\n")

def print_separator():
    printf("  -------------------------------------------\n")

def section_header(title: str):
    printf("\n+-- %s --+\n", title)

def section_footer():
    printf("+--------------------------------------------+\n")

# ============================================================

def test_version():
    section_header("Version Info")

    ver: str = pyzlib.ZLIB_VERSION
    printf("  ZLIB_VERSION (compile-time) : %s\n", ver)
    check("ZLIB_VERSION not None", ver != None, "version string is None")
    check("ZLIB_VERSION not empty", strlen(ver) > 0, "version string is empty")

    runtime_ver: str = pyzlib.runtime_version()
    printf("  ZLIB_RUNTIME_VERSION       : %s\n", runtime_ver)
    check("runtime version not None", runtime_ver != None, "runtime version is None")
    check("runtime version not empty", strlen(runtime_ver) > 0, "runtime version is empty")

    section_footer()


def test_compress_decompress():
    section_header("compress / decompress")

    inp: str = "Hello, World! This is a test of zlib compression and decompression."
    input_length: t.CSizeT = strlen(inp)
    printf("  Input (%zu bytes): \"%s\"\n", input_length, inp)

    out_length: t.CSizeT = 0
    compressed: BYTEPTR = pyzlib.compress(pool, BYTEPTR(inp), input_length,
                                              pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.MAX_WBITS, c.Addr(out_length))
    check("compress returns non-None", compressed != None, "compress returned None")
    check("compress output size > 0", out_length > 0, "compressed size is 0")
    check("compress output produced", out_length > 0, "compressed size is 0")

    printf("  Compressed (%zu bytes): ", out_length)
    print_hex_test(compressed, out_length, 32)
    printf("  Compression ratio: %.1f%%\n", t.CDouble(out_length) / input_length * 100)

    dec_length: t.CSizeT = 0
    decompressed: BYTEPTR = pyzlib.decompress(pool, compressed, out_length,
                                                   pyzlib.MAX_WBITS, pyzlib.DEF_BUF_SIZE, c.Addr(dec_length))
    check("decompress returns non-None", decompressed != None, "decompress returned None")
    check("decompress output size matches input", dec_length == input_length, "size mismatch")
    check("decompress output matches input", decompressed != None and memcmp(decompressed, inp, input_length) == 0, "content mismatch")

    printf("  Decompressed (%zu bytes): \"%.*s\"\n", dec_length, t.CInt(dec_length), str(decompressed))

    free(compressed)
    free(decompressed)
    section_footer()


def test_compress_levels():
    section_header("Compression Levels")

    inp: str = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
    input_length: t.CSizeT = strlen(inp)
    printf("  Input (%zu bytes): \"%s\"\n", input_length, inp)

    out_length: t.CSizeT = 0

    c0: BYTEPTR = pyzlib.compress(pool, BYTEPTR(inp), input_length,
                                       pyzlib.Z_NO_COMPRESSION, pyzlib.MAX_WBITS, c.Addr(out_length))
    check("level 0 compress OK", c0 != None, "level 0 returned None")
    printf("  Level 0 (Z_NO_COMPRESSION): %zu bytes\n", out_length)
    length0: t.CSizeT = out_length
    free(c0)

    c1: BYTEPTR = pyzlib.compress(pool, BYTEPTR(inp), input_length,
                                       pyzlib.Z_BEST_SPEED, pyzlib.MAX_WBITS, c.Addr(out_length))
    check("level 1 compress OK", c1 != None, "level 1 returned None")
    printf("  Level 1 (Z_BEST_SPEED)    : %zu bytes\n", out_length)
    free(c1)

    c6: BYTEPTR = pyzlib.compress(pool, BYTEPTR(inp), input_length,
                                       pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.MAX_WBITS, c.Addr(out_length))
    check("level -1 compress OK", c6 != None, "default level returned None")
    printf("  Level -1 (DEFAULT)        : %zu bytes\n", out_length)
    free(c6)

    c9: BYTEPTR = pyzlib.compress(pool, BYTEPTR(inp), input_length,
                                       pyzlib.Z_BEST_COMPRESSION, pyzlib.MAX_WBITS, c.Addr(out_length))
    check("level 9 compress OK", c9 != None, "level 9 returned None")
    printf("  Level 9 (Z_BEST_COMPRESS) : %zu bytes\n", out_length)
    length9: t.CSizeT = out_length
    free(c9)

    check("level 9 smaller than level 0", length9 < length0, "level 9 not smaller than level 0")

    section_footer()


def test_compressobj():
    section_header("compressobj (incremental)")

    part1: str = "Hello, "
    part2: str = "World!"
    printf("  Part 1: \"%s\"\n", part1)
    printf("  Part 2: \"%s\"\n", part2)

    out_length: t.CSizeT = 0

    s: pyzlib.Compress | t.CPtr = pyzlib.compressobj(pool, pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.DEFLATED,
                                         pyzlib.MAX_WBITS, pyzlib.DEF_MEM_LEVEL,
                                         pyzlib.Z_DEFAULT_STRATEGY, None, 0)
    check("compressobj created", s != None, "compressobj is None")
    print_separator()

    c1: BYTEPTR = s.compress(BYTEPTR(part1), strlen(part1), c.Addr(out_length))
    check("compress part1 OK", c1 != None, "compress part1 failed")
    printf("  Compressed part1: %zu bytes  -> ", out_length)
    print_hex_test(c1, out_length, 16)
    total: t.CSizeT = out_length
    all: BYTEPTR = None
    if out_length > 0 and c1:
        all = BYTEPTR(malloc(total))
        if all: memcpy(all, c1, out_length)
    free(c1)

    c2: BYTEPTR = s.compress(BYTEPTR(part2),
                                                strlen(part2), c.Addr(out_length))
    check("compress part2 OK", c2 != None, "compress part2 failed")
    printf("  Compressed part2: %zu bytes  -> ", out_length)
    print_hex_test(c2, out_length, 16)
    if out_length > 0 and c2:
        new_all: BYTEPTR = BYTEPTR(realloc(all, total + out_length))
        if new_all:
            all = new_all
            memcpy(all + total, c2, out_length)
            total += out_length
    free(c2)

    c3: BYTEPTR = s.flush(pyzlib.Z_FINISH, c.Addr(out_length))
    check("flush (Z_FINISH) OK", c3 != None, "flush failed")
    printf("  Flush result    : %zu bytes  -> ", out_length)
    print_hex_test(c3, out_length, 16)
    if out_length > 0 and c3:
        new_all: BYTEPTR = realloc(all, total + out_length) # 隐式转换
        if new_all:
            all = new_all
            memcpy(all + total, c3, out_length)
            total += out_length
    free(c3)

    print_separator()

    if all:
        dec_length: t.CSizeT = 0
        dec: BYTEPTR = pyzlib.decompress(pool, all, total, pyzlib.MAX_WBITS, pyzlib.DEF_BUF_SIZE, c.Addr(dec_length))
        check("decompress incremental OK", dec != None, "decompress returned None")
        check("decompress incremental size", dec != None and dec_length == strlen(part1) + strlen(part2), "size mismatch")
        check("decompress incremental content",
              dec != None and memcmp(dec, "Hello, World!", dec_length) == 0, "content mismatch")
        if dec:
            printf("  Decompressed: \"%.*s\" (%zu bytes)\n", t.CInt(dec_length), str(dec), dec_length)
        free(dec)
        free(all)

    s.delete() # del s
    section_footer()

def test_decompressobj():
    section_header("decompressobj")

    inp: str = "Test data for decompress object."
    input_length: t.CSizeT = strlen(inp)
    printf("  Input (%zu bytes): \"%s\"\n", input_length, inp)

    out_length: t.CSizeT = 0
    compressed: BYTEPTR = pyzlib.compress(pool, BYTEPTR(inp), input_length,
                                              pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.MAX_WBITS, c.Addr(out_length))
    printf("  Compressed: %zu bytes\n", out_length)

    d: pyzlib.Decompress | t.CPtr = pyzlib.decompressobj(pool, pyzlib.MAX_WBITS, None, 0)
    check("decompressobj created", d != None, "decompressobj is None")
    check("eof is false initially", d.eof == 0, "eof should be 0")
    print_separator()

    dec_length: t.CSizeT = 0
    dec: BYTEPTR = d.decompress(compressed, out_length,
                                                0, c.Addr(dec_length))
    check("decompress OK", dec != None, "decompress returned None")
    check("decompress size", dec != None and dec_length == input_length, "size mismatch")
    check("decompress content", dec != None and memcmp(dec, inp, input_length) == 0, "content mismatch")
    check("eof is true after stream end", d.eof == 1, "eof should be 1")

    if dec:
        printf("  Decompressed (%zu bytes): \"%.*s\"\n", dec_length, t.CInt(dec_length), str(dec))
    printf("  eof = %d\n", d.eof)

    free(dec)
    free(compressed)
    d.delete()
    section_footer()

def test_decompressobj_max_lengthgth():
    section_header("decompressobj max_lengthgth")

    inp: str = "This is a longer string for testing max_lengthgth parameter in decompress."
    input_length: t.CSizeT = strlen(inp)
    printf("  Input (%zu bytes): \"%s\"\n", input_length, inp)

    out_length: t.CSizeT = 0
    compressed: BYTEPTR = pyzlib.compress(pool, BYTEPTR(inp), input_length,
                                              pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.MAX_WBITS, c.Addr(out_length))
    printf("  Compressed: %zu bytes\n", out_length)

    d: pyzlib.Decompress | t.CPtr = pyzlib.decompressobj(pool, pyzlib.MAX_WBITS, None, 0)

    dec_length: t.CSizeT = 0
    dec: BYTEPTR = d.decompress(compressed, out_length,
                                                10, c.Addr(dec_length))
    check("max_lengthgth decompress OK", dec != None, "decompress returned None")
    check("max_lengthgth limits output", dec_length <= 10, "output exceeded max_lengthgth")
    if dec:
        printf("  First chunk (max_lengthgth=10): \"%.*s\" (%zu bytes)\n", t.CInt(dec_length), str(dec), dec_length)
    free(dec)

    tail_length: t.CSizeT = 0
    tail: BYTEPTR = d.unconsumed_tail(c.Addr(tail_length))
    check("unconsumed_tail exists", tail != None or tail_length == 0, "no unconsumed_tail")
    printf("  unconsumed_tail: %zu bytes remaining\n", tail_length)
    printf("  eof = %d\n", d.eof)

    if tail_length > 0 or not d.eof:
        flushed: BYTEPTR = d.flush(pyzlib.DEF_BUF_SIZE, c.Addr(dec_length))
        check("flush after max_lengthgth OK", flushed != None, "flush returned None")
        if flushed: printf("  Flushed (%zu bytes): \"%.*s\"\n", dec_length, t.CInt(dec_length), str(flushed))
        free(flushed)

    printf("  eof after flush = %d\n", d.eof)

    d.delete()
    free(compressed)
    section_footer()


def test_decompressobj_unused_data():
    section_header("decompressobj unused_data")

    inp: str = "Hello"
    input_length: t.CSizeT = strlen(inp)
    printf("  Input: \"%s\"\n", inp)

    out_length: t.CSizeT = 0
    compressed: BYTEPTR = pyzlib.compress(pool, BYTEPTR(inp), input_length,
                                              pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.MAX_WBITS, c.Addr(out_length))
    printf("  Compressed: %zu bytes\n", out_length)

    total_length: t.CSizeT = out_length + 5
    combined: BYTEPTR = BYTEPTR(malloc(total_length))
    memcpy(combined, compressed, out_length)
    memcpy(combined + out_length, "EXTRA", 5)
    free(compressed)
    printf("  Combined (compressed + \"EXTRA\"): %zu bytes\n", total_length)

    d: pyzlib.Decompress | t.CPtr = pyzlib.decompressobj(pool, pyzlib.MAX_WBITS, None, 0)

    dec_length: t.CSizeT = 0
    dec: BYTEPTR = d.decompress(combined, total_length,
                                                0, c.Addr(dec_length))
    check("decompress with extra OK", dec != None, "decompress returned None")
    check("eof after stream end", d.eof == 1, "eof should be 1")
    if dec:
        printf("  Decompressed: \"%.*s\" (%zu bytes)\n", t.CInt(dec_length), str(dec), dec_length)
    free(dec)

    print_separator()

    unused_length: t.CSizeT = 0
    unused: BYTEPTR = d.unused_data(c.Addr(unused_length))
    check("unused_data exists", unused != None and unused_length > 0, "no unused_data")
    check("unused_data is EXTRA", unused_length == 5 and unused != None and memcmp(unused, "EXTRA", 5) == 0, "unused_data mismatch")
    printf("  unused_data (%zu bytes): \"%.*s\"\n", unused_length, t.CInt(unused_length), str(unused))

    d.delete()
    free(combined)
    section_footer()


def test_compress_copy():
    section_header("Compress.copy()")

    inp: str = "Copy test data for compress object"
    input_length: t.CSizeT = strlen(inp)
    printf("  Input: \"%s\"\n", inp)

    out_length: t.CSizeT = 0
    c1: pyzlib.Compress | t.CPtr = pyzlib.compressobj(pool, pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.DEFLATED,
                                          pyzlib.MAX_WBITS, pyzlib.DEF_MEM_LEVEL,
                                          pyzlib.Z_DEFAULT_STRATEGY, None, 0)
    c1.compress(BYTEPTR(inp), input_length, c.Addr(out_length))
    printf("  Original: compressed %zu bytes so far\n", out_length)

    c2: pyzlib.Compress | t.CPtr = c1.copy()
    check("compress copy OK", c2 != None, "copy returned None")
    printf("  Copied compressobj\n")

    print_separator()

    f1: BYTEPTR = c1.flush(pyzlib.Z_FINISH, c.Addr(out_length))
    check("flush original OK", f1 != None, "flush original failed")
    length1: t.CSizeT = out_length
    printf("  Original flush: %zu bytes\n", length1)
    free(f1)

    f2: BYTEPTR = c2.flush(pyzlib.Z_FINISH, c.Addr(out_length))
    check("flush copy OK", f2 != None, "flush copy failed")
    length2: t.CSizeT = out_length
    printf("  Copy flush   : %zu bytes\n", length2)
    free(f2)

    check("copy produces same output", length1 == length2, "output lengthgths differ")

    c1.delete()
    c2.delete()
    section_footer()


def test_decompress_copy():
    section_header("Decompress.copy()")

    inp: str = "Decompress copy test"
    input_length: t.CSizeT = strlen(inp)
    printf("  Input: \"%s\"\n", inp)

    out_length: t.CSizeT = 0
    compressed: BYTEPTR = pyzlib.compress(pool, BYTEPTR(inp), input_length,
                                              pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.MAX_WBITS, c.Addr(out_length))
    printf("  Compressed: %zu bytes\n", out_length)

    d1: pyzlib.Decompress | t.CPtr = pyzlib.decompressobj(pool, pyzlib.MAX_WBITS, None, 0)
    d2: pyzlib.Decompress | t.CPtr = d1.copy()
    check("decompress copy OK", d2 != None, "copy returned None")
    printf("  Copied decompressobj\n")

    print_separator()

    dec_length1: t.CSizeT = 0
    dec_length2: t.CSizeT = 0
    dec1: BYTEPTR = d1.decompress(compressed, out_length,
                                                  0, c.Addr(dec_length1))
    dec2: BYTEPTR = d2.decompress(compressed, out_length,
                                                  0, c.Addr(dec_length2))
    check("decompress copy output size", dec_length1 == dec_length2, "sizes differ")
    check("decompress copy output content",
          dec1 and dec2 and memcmp(dec1, dec2, dec_length1) == 0, "content differs")

    if dec1: printf("  Original: \"%.*s\" (%zu bytes)\n", t.CInt(dec_length1), str(dec1), dec_length1)
    if dec2: printf("  Copy    : \"%.*s\" (%zu bytes)\n", t.CInt(dec_length2), str(dec2), dec_length2)

    free(dec1)
    free(dec2)
    free(compressed)
    d1.delete()
    d2.delete()
    section_footer()


def test_adler32():
    section_header("adler32 checksum")

    checksum: ULONG = pyzlib.zlib_adler32(BYTEPTR("Hello"), 5, 1)
    printf("  adler32(\"Hello\") = 0x%08lX (%lu)\n", checksum, checksum)
    check("adler32 returns non-zero", checksum != 0, "checksum is 0")

    incremental: ULONG = pyzlib.zlib_adler32(BYTEPTR("Hel"), 3, 1)
    printf("  adler32(\"Hel\") = 0x%08lX\n", incremental)
    incremental = pyzlib.zlib_adler32(BYTEPTR("lo"), 2, incremental)
    printf("  adler32(\"lo\", prev) = 0x%08lX\n", incremental)
    check("adler32 incremental matches one-shot", incremental == checksum, "incremental != one-shot")
    printf("  Incremental == One-shot: %s\n", incremental == "YES" if checksum else "NO")

    section_footer()


def test_crc32():
    section_header("crc32 checksum")

    checksum: ULONG = pyzlib.zlib_crc32(BYTEPTR("Hello"), 5, 0)
    printf("  crc32(\"Hello\") = 0x%08lX (%lu)\n", checksum, checksum)
    check("crc32 returns non-zero", checksum != 0, "checksum is 0")

    incremental: ULONG = pyzlib.zlib_crc32(BYTEPTR("Hel"), 3, 0)
    printf("  crc32(\"Hel\") = 0x%08lX\n", incremental)
    incremental = pyzlib.zlib_crc32(BYTEPTR("lo"), 2, incremental)
    printf("  crc32(\"lo\", prev) = 0x%08lX\n", incremental)
    check("crc32 incremental matches one-shot", incremental == checksum, "incremental != one-shot")
    printf("  Incremental == One-shot: %s\n", incremental == "YES" if checksum else "NO")

    section_footer()


def test_empty_compress():
    section_header("Empty data compress/decompress")

    out_length: t.CSizeT = 0
    compressed: BYTEPTR = pyzlib.compress(pool, BYTEPTR(""), 0,
                                              pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.MAX_WBITS, c.Addr(out_length))
    check("compress empty OK", compressed != None, "compress empty returned None")
    check("compress empty output > 0", out_length > 0, "compressed empty is 0 bytes")
    printf("  Empty input compressed: %zu bytes (header only)\n", out_length)
    printf("  Hex: ")
    print_hex_test(compressed, out_length, 32)

    dec_length: t.CSizeT = 0
    decompressed: BYTEPTR = pyzlib.decompress(pool, compressed, out_length,
                                                   pyzlib.MAX_WBITS, pyzlib.DEF_BUF_SIZE, c.Addr(dec_length))
    check("decompress empty OK", decompressed != None, "decompress empty returned None")
    check("decompress empty size == 0", dec_length == 0, "decompressed size != 0")
    printf("  Decompressed back: %zu bytes\n", dec_length)

    free(compressed)
    free(decompressed)
    section_footer()


def test_gzip_format():
    section_header("Gzip format (wbits=31)")

    inp: str = "Gzip format test"
    input_length: t.CSizeT = strlen(inp)
    printf("  Input: \"%s\" (%zu bytes)\n", inp, input_length)

    gzip_wbits: t.CInt = pyzlib.MAX_WBITS + 16
    printf("  wbits = MAX_WBITS + 16 = %d\n", gzip_wbits)

    out_length: t.CSizeT = 0
    compressed: BYTEPTR = pyzlib.compress(pool, BYTEPTR(inp), input_length,
                                              pyzlib.Z_DEFAULT_COMPRESSION, gzip_wbits, c.Addr(out_length))
    check("gzip compress OK", compressed != None, "gzip compress failed")
    printf("  Gzip compressed: %zu bytes\n", out_length)
    printf("  Header bytes: ")
    print_hex_test(compressed, 4 if (out_length > 4) else out_length, 4)

    dec_length: t.CSizeT = 0
    decompressed: BYTEPTR = pyzlib.decompress(pool, compressed, out_length,
                                                   gzip_wbits, pyzlib.DEF_BUF_SIZE, c.Addr(dec_length))
    check("gzip decompress OK", decompressed != None, "gzip decompress failed")
    check("gzip decompress size", dec_length == input_length, "size mismatch")
    check("gzip decompress content", decompressed != None and memcmp(decompressed, inp, input_length) == 0, "content mismatch")
    if decompressed: printf("  Decompressed: \"%.*s\" (%zu bytes)\n", t.CInt(dec_length), str(decompressed), dec_length)

    free(compressed)
    free(decompressed)
    section_footer()


def test_raw_deflate():
    section_header("Raw deflate (wbits=-15)")

    inp: str = "Raw deflate test"
    input_length: t.CSizeT = strlen(inp)
    printf("  Input: \"%s\" (%zu bytes)\n", inp, input_length)

    raw_wbits: t.CInt = -pyzlib.MAX_WBITS
    printf("  wbits = -MAX_WBITS = %d\n", raw_wbits)

    out_length: t.CSizeT = 0
    compressed: BYTEPTR = pyzlib.compress(pool, BYTEPTR(inp), input_length,
                                              pyzlib.Z_DEFAULT_COMPRESSION, raw_wbits, c.Addr(out_length))
    check("raw deflate compress OK", compressed != None, "raw deflate compress failed")
    printf("  Raw compressed: %zu bytes\n", out_length)

    dec_length: t.CSizeT = 0
    decompressed: BYTEPTR = pyzlib.decompress(pool, compressed, out_length,
                                                   raw_wbits, pyzlib.DEF_BUF_SIZE, c.Addr(dec_length))
    check("raw deflate decompress OK", decompressed != None, "raw deflate decompress failed")
    check("raw deflate decompress size", dec_length == input_length, "size mismatch")
    check("raw deflate decompress content", decompressed != None and memcmp(decompressed, inp, input_length) == 0, "content mismatch")
    if decompressed: printf("  Decompressed: \"%.*s\" (%zu bytes)\n", t.CInt(dec_length), str(decompressed), dec_length)

    free(compressed)
    free(decompressed)
    section_footer()


def test_zdict():
    section_header("Preset dictionary (zdict)")

    dict: str = "common words that appear frequently in the data"
    inp: str = "common words that appear frequently"
    input_length: t.CSizeT = strlen(inp)
    printf("  Dictionary: \"%s\" (%zu bytes)\n", dict, strlen(dict))
    printf("  Input     : \"%s\" (%zu bytes)\n", inp, input_length)

    out_length: t.CSizeT = 0
    s: pyzlib.Compress | t.CPtr = pyzlib.compressobj(pool, pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.DEFLATED,
                                         pyzlib.MAX_WBITS, pyzlib.DEF_MEM_LEVEL,
                                         pyzlib.Z_DEFAULT_STRATEGY,
                                         BYTEPTR(dict), strlen(dict))
    check("compressobj with zdict OK", s != None, "compressobj with zdict failed")

    comp_data: BYTEPTR = s.compress(BYTEPTR(inp),
                                                       input_length, c.Addr(out_length))
    check("compress with zdict OK", comp_data != None, "compress with zdict failed")
    printf("  Compressed with zdict: %zu bytes\n", out_length)
    free(comp_data)

    flush_data: BYTEPTR = s.flush(pyzlib.Z_FINISH, c.Addr(out_length))
    check("flush with zdict OK", flush_data != None, "flush with zdict failed")
    printf("  Flush: %zu bytes\n", out_length)
    free(flush_data)

    print_separator()

    total_comp: t.CSizeT = out_length
    s.delete()

    s = pyzlib.compressobj(pool, pyzlib.Z_DEFAULT_COMPRESSION, pyzlib.DEFLATED,
                         pyzlib.MAX_WBITS, pyzlib.DEF_MEM_LEVEL,
                         pyzlib.Z_DEFAULT_STRATEGY, None, 0)
    comp_no_dict: BYTEPTR = s.compress(BYTEPTR(inp),
                                                          input_length, c.Addr(out_length))
    flush_no_dict: BYTEPTR = s.flush(pyzlib.Z_FINISH, c.Addr(out_length))
    no_dict_total: t.CSizeT = 0
    if comp_no_dict: no_dict_total += out_length
    if flush_no_dict: no_dict_total += out_length
    printf("  Without zdict: ~%zu bytes compressed\n", no_dict_total)
    printf("  With zdict   : %zu bytes compressed\n", total_comp)
    free(comp_no_dict)
    free(flush_no_dict)
    s.delete()

    section_footer()


def test_huffman_tree():
    section_header("Huffman tree construction")

    printf("  --- Fixed Huffman tree (literal/lengthgth) ---\n")
    lit_tree = zhuff.zhuff_tree()
    lit_tree.build_fixed_lit_tree()

    check("fixed lit tree count == 288", lit_tree.count == 288, "count mismatch")
    check("fixed lit tree max_bits == 9", lit_tree.max_bits == 9, "max_bits mismatch")

    has_8bit: t.CInt = 0
    has_9bit: t.CInt = 0
    has_7bit: t.CInt = 0
    for i in range(143 + 1):
        if lit_tree.codes[i].bits == 8:
            has_8bit += 1
    for i in range(143 + 1, 255 + 1):
        if lit_tree.codes[i].bits == 9:
            has_9bit += 1
    for i in range(255 + 1, 279 + 1):
        if lit_tree.codes[i].bits == 7:
            has_7bit += 1
    check("0-143 are 8-bit", has_8bit == 144, "wrong count")
    check("144-255 are 9-bit", has_9bit == 112, "wrong count")
    check("256-279 are 7-bit", has_7bit == 24, "wrong count")

    printf("  0-143: %d codes with 8 bits\n", has_8bit)
    printf("  144-255: %d codes with 9 bits\n", has_9bit)
    printf("  256-279: %d codes with 7 bits\n", has_7bit)
    printf("  280-287: 8-bit (end-of-block 256 = 7-bit)\n")
    printf("  EOB (256): code=0x%X, bits=%d\n", lit_tree.codes[256].code, lit_tree.codes[256].bits)

    printf("\n  --- Fixed Huffman tree (distance) ---\n")
    dist_tree = zhuff.zhuff_tree()
    dist_tree.build_fixed_dist_tree()

    check("fixed dist tree count == 32", dist_tree.count == 32, "count mismatch")
    check("fixed dist tree max_bits == 5", dist_tree.max_bits == 5, "max_bits mismatch")

    all_5bit: t.CInt = 1
    for i in range(32):
        if dist_tree.codes[i].bits != 5:
            all_5bit = 0
    check("all 32 distance codes are 5-bit", all_5bit, "not all 5-bit")
    printf("  All 32 distance codes: 5 bits each\n")

    printf("\n  --- Dynamic Huffman tree from frequencies ---\n")
    freqs: t.CArray[t.CInt, 257]
    memset(freqs, 0, freqs.__sizeof__())
    freqs['A'] = 50
    freqs['B'] = 25
    freqs['C'] = 12
    freqs['D'] = 6
    freqs['E'] = 3
    freqs['F'] = 1
    freqs[256] = 1

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

    check("dynamic tree count == 257", dyn_tree.count == 257, "count mismatch")
    check("dynamic tree max_bits <= 15", dyn_tree.max_bits <= 15, "max_bits overflow")

    total_codes: t.CInt = 0
    max_length_found: t.CInt = 0
    for i in range(257):
        if dyn_tree.codes[i].bits > 0:
            total_codes += 1
            if dyn_tree.codes[i].bits > max_length_found:
                max_length_found = dyn_tree.codes[i].bits
    check("dynamic tree has 7 active codes", total_codes == 7, "wrong active count")
    check("dynamic tree max code lengthgth found <= 15", max_length_found <= 15, "code lengthgth overflow")

    printf("  Frequencies: A=50, B=25, C=12, D=6, E=3, F=1, EOB=1\n")
    printf("  Active codes: %d, max code lengthgth: %d\n", total_codes, max_length_found)

    printf("  Code assignments:\n")
    names: t.CArray[str, None] = ["A","B","C","D","E","F"] # 一个都是char* 的数组
    syms: t.CArray[t.CInt, None] = ['A','B','C','D','E','F']
    for i in range(6):
        printf("    '%s' (freq=%3d): code=0x%04X, bits=%d\n",
                names[i], freqs[syms[i]],
                dyn_tree.codes[syms[i]].code,
                dyn_tree.codes[syms[i]].bits)



    printf("    EOB (freq=1): code=0x%04X, bits=%d\n",
            dyn_tree.codes[256].code, dyn_tree.codes[256].bits)
    a_bits: t.CInt = dyn_tree.codes['A'].bits
    f_bits: t.CInt = dyn_tree.codes['F'].bits
    check("high-freq symbol has shorter code", a_bits <= f_bits, "A should be <= F")
    printf("  High-freq 'A' (%d bits) <= low-freq 'F' (%d bits): %s\n",
            a_bits, f_bits, "YES" if (a_bits <= f_bits) else "NO")

    printf("\n  --- Huffman encode/decode roundtrip ---\n")
    freqs: t.CArray[t.CInt, 288]
    memset(freqs, 0, freqs.__sizeof__())
    freqs['H'] = 10
    freqs['e'] = 8
    freqs['l'] = 20
    freqs['o'] = 8
    freqs[' '] = 5
    freqs['W'] = 3
    freqs['r'] = 5
    freqs['d'] = 3
    freqs['!'] = 2
    freqs[256] = 1

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

    dec_tree = zhuff.zhuff_decode_tree()
    dec_tree.build_decode_tree(c.Addr(enc_tree))

    writer = zdef.zbit_writer()

    symbols: t.CArray[t.CInt, 16]
    symbols[0] = 'H'; symbols[1] = 'e'; symbols[2] = 'l'; symbols[3] = 'l'
    symbols[4] = 'o'; symbols[5] = ' '; symbols[6] = 'W'; symbols[7] = 'o'
    symbols[8] = 'r'; symbols[9] = 'l'; symbols[10] = 'd'; symbols[11] = '!'
    symbols[12] = 256
    num_symbols: t.CInt = 13

    for i in range(num_symbols):
        enc_tree.encode_symbol(symbols[i], c.Addr(writer))

    reader = zdef.zbit_reader()
    reader.buf = writer.buf
    reader.length = writer.byte_pos + (1 if (writer.bit_pos > 0) else 0)
    reader.byte_pos = 0
    reader.bit_pos = 0

    roundtrip_ok: t.CInt = 1
    for i in range(num_symbols):
        decoded: t.CInt = dec_tree.decode_symbol(c.Addr(reader))
        if decoded != symbols[i]:
            roundtrip_ok = 0
            printf("  MISMATCH at symbol %d: expected %d, got %d\n", i, symbols[i], decoded)
            break
    check("encode/decode roundtrip matches", roundtrip_ok, "roundtrip failed")
    printf("  Encoded %d symbols into %zu bytes, decoded all correctly\n",
            num_symbols, writer.byte_pos + (1 if (writer.bit_pos > 0) else 0))

    free(writer.buf)

    printf("\n  --- Overflow handling (many symbols, limited max_bits) ---\n")
    freqs: t.CArray[t.CInt, 257]
    for i in range(256):
        freqs[i] = 1
    freqs[256] = 1

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

    overflow_ok: t.CInt = 1
    for i in range(257):
        if tree.codes[i].bits > 15 or tree.codes[i].bits < 0:
            overflow_ok = 0
            break
    check("all code lengthgths <= 15 with 257 symbols", overflow_ok, "code lengthgth overflow")

    bl_count: t.CArray[t.CInt, 16] = [0]
    for i in range(257):
        if tree.codes[i].bits > 0:
            bl_count[tree.codes[i].bits] += 1
    printf("  Code lengthgth distribution (257 equal-freq symbols, max_bits=15):\n")
    for i in range(1, 15 + 1):
        if bl_count[i] > 0:
            printf("    %2d bits: %3d codes\n", i, bl_count[i])

    dt = zhuff.zhuff_decode_tree()
    dt.build_decode_tree(c.Addr(tree))

    w = zdef.zbit_writer()
    tree.encode_symbol(0,   c.Addr(w))
    tree.encode_symbol(128, c.Addr(w))
    tree.encode_symbol(255, c.Addr(w))
    tree.encode_symbol(256, c.Addr(w))

    r = zdef.zbit_reader()
    r.buf = w.buf
    r.length = w.byte_pos + (1 if (w.bit_pos > 0) else 0)
    r.byte_pos = 0
    r.bit_pos  = 0

    s0: t.CInt = dt.decode_symbol(c.Addr(r))
    s1: t.CInt = dt.decode_symbol(c.Addr(r))
    s2: t.CInt = dt.decode_symbol(c.Addr(r))
    s3: t.CInt = dt.decode_symbol(c.Addr(r))

    check("overflow roundtrip: symbol 0", s0 == 0, "decode mismatch")
    check("overflow roundtrip: symbol 128", s1 == 128, "decode mismatch")
    check("overflow roundtrip: symbol 255", s2 == 255, "decode mismatch")
    check("overflow roundtrip: symbol 256 (EOB)", s3 == 256, "decode mismatch")

    free(w.buf)


    printf("\n  --- Kraft inequality check ---\n")
    freqs: t.CArray[t.CInt, 257]
    memset(freqs, 0, freqs.__sizeof__())
    freqs['X'] = 100
    freqs['Y'] = 50
    freqs['Z'] = 25
    freqs[256] = 1

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

    kraft_sum: t.CDouble = 0.0
    for i in range(257):
        if tree.codes[i].bits > 0:
            kraft_sum += 1.0 / (t.CDouble(ULONGLONG(1) << tree.codes[i].bits))
    check("Kraft inequality: sum <= 1.0", kraft_sum <= 1.0 + 1e-9, "Kraft violated")
    printf("  Kraft sum = %.10f (must be <= 1.0)\n", kraft_sum)

    section_footer()


def test_error_handling():
    section_header("Error handling")

    out_length: t.CSizeT = 0
    printf("  Trying to decompress garbage data...\n")

    result: BYTEPTR = pyzlib.decompress(pool, BYTEPTR("garbage"), 7,
                                             pyzlib.MAX_WBITS, pyzlib.DEF_BUF_SIZE, c.Addr(out_length))
    check("decompress garbage returns None", result == None, "should have returned None")
    printf("  Error code   : %d\n", pyzlib.get_error_code())
    printf("  Error message: \"%s\"\n", pyzlib.get_error())
    check("error message set", strlen(pyzlib.get_error()) > 0, "error message is empty")
    check("error code is set", pyzlib.get_error_code() != pyzlib.Z_OK, "error code is Z_OK")

    print_separator()

    pyzlib.clear_error()
    check("clear error clears code", pyzlib.get_error_code() == 0, "error code not cleared")
    printf("  After clear: code=%d, msg=\"%s\"\n", pyzlib.get_error_code(), pyzlib.get_error())

    section_footer()


def test_constants():
    section_header("Constants")
    printf("  Compression Levels:\n")
    printf("    Z_NO_COMPRESSION       = %d\n", pyzlib.Z_NO_COMPRESSION)
    printf("    Z_BEST_SPEED           = %d\n", pyzlib.Z_BEST_SPEED)
    printf("    Z_BEST_COMPRESSION     = %d\n", pyzlib.Z_BEST_COMPRESSION)
    printf("    Z_DEFAULT_COMPRESSION  = %d\n", pyzlib.Z_DEFAULT_COMPRESSION)
    printf("  Methods:\n")
    printf("    DEFLATED               = %d\n", pyzlib.DEFLATED)
    printf("  Flush Modes:\n")
    printf("    Z_NO_FLUSH             = %d\n", pyzlib.Z_NO_FLUSH)
    printf("    Z_PARTIAL_FLUSH        = %d\n", pyzlib.Z_PARTIAL_FLUSH)
    printf("    Z_SYNC_FLUSH           = %d\n", pyzlib.Z_SYNC_FLUSH)
    printf("    Z_FULL_FLUSH           = %d\n", pyzlib.Z_FULL_FLUSH)
    printf("    Z_FINISH               = %d\n", pyzlib.Z_FINISH)
    printf("    Z_BLOCK                = %d\n", pyzlib.Z_BLOCK)
    printf("    Z_TREES                = %d\n", pyzlib.Z_TREES)
    printf("  Strategies:\n")
    printf("    Z_DEFAULT_STRATEGY     = %d\n", pyzlib.Z_DEFAULT_STRATEGY)
    printf("    Z_FILTERED             = %d\n", pyzlib.Z_FILTERED)
    printf("    Z_HUFFMAN_ONLY         = %d\n", pyzlib.Z_HUFFMAN_ONLY)
    printf("    Z_RLE                  = %d\n", pyzlib.Z_RLE)
    printf("    Z_FIXED                = %d\n", pyzlib.Z_FIXED)
    printf("  Return Codes:\n")
    printf("    Z_OK                   = %d\n", pyzlib.Z_OK)
    printf("    Z_STREAM_END           = %d\n", pyzlib.Z_STREAM_END)
    printf("    Z_NEED_DICT            = %d\n", pyzlib.Z_NEED_DICT)
    printf("    Z_ERRNO                = %d\n", pyzlib.Z_ERRNO)
    printf("    Z_STREAM_ERROR         = %d\n", pyzlib.Z_STREAM_ERROR)
    printf("    Z_DATA_ERROR           = %d\n", pyzlib.Z_DATA_ERROR)
    printf("    Z_MEM_ERROR            = %d\n", pyzlib.Z_MEM_ERROR)
    printf("    Z_BUF_ERROR            = %d\n", pyzlib.Z_BUF_ERROR)
    printf("    Z_VERSION_ERROR        = %d\n", pyzlib.Z_VERSION_ERROR)
    printf("  Window / Buffer:\n")
    printf("    MAX_WBITS              = %d\n", pyzlib.MAX_WBITS)
    printf("    DEF_BUF_SIZE           = %d\n", pyzlib.DEF_BUF_SIZE)
    printf("    DEF_MEM_LEVEL          = %d\n", pyzlib.DEF_MEM_LEVEL)
    section_footer()


def main123456() -> t.CInt:
    printf("==============================================\n")
    printf("  Python zlib - C Implementation Tests\n")
    printf("==============================================\n")

    test_constants()
    test_version()
    test_compress_decompress()
    test_compress_levels()
    test_compressobj()
    test_decompressobj()
    test_decompressobj_max_lengthgth()
    test_decompressobj_unused_data()
    test_compress_copy()
    test_decompress_copy()
    test_adler32()
    test_crc32()
    test_empty_compress()
    test_gzip_format()
    test_raw_deflate()
    test_zdict()
    test_huffman_tree()
    test_error_handling()

    printf("\n==============================================\n")
    printf("  Results: %d passed, %d failed\n", test_passed, test_failed)
    printf("==============================================\n")

    return 1 if test_failed > 0 else 0