/*
 * ZXC - High-performance lossless compression
 *
 * Copyright (c) 2025-2026 Bertrand Lebonnois and contributors.
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include "test_common.h"



/* Round-trip the Huffman codec over a few representative literal distributions. */
static int huf_roundtrip_case(const char* label, const uint8_t* literals, size_t n) {
    uint32_t freq[ZXC_HUF_NUM_SYMBOLS] = {0};
    for (size_t i = 0; i < n; i++) freq[literals[i]]++;

    uint8_t code_len[ZXC_HUF_NUM_SYMBOLS];
    if (zxc_huf_build_code_lengths(freq, code_len, NULL, ZXC_HUF_MAX_CODE_LEN_DENSITY) != ZXC_OK) {
        printf("Failed [%s]: build_code_lengths\n", label);
        return 0;
    }
    /* Validate the lengths-limit invariant. */
    for (int i = 0; i < ZXC_HUF_NUM_SYMBOLS; i++) {
        if (code_len[i] > ZXC_HUF_MAX_CODE_LEN_ULTRA) {
            printf("Failed [%s]: code_len[%d] = %d > %d\n", label, i, code_len[i],
                   ZXC_HUF_MAX_CODE_LEN_ULTRA);
            return 0;
        }
    }

    /* Worst-case payload size: 128-byte header + packed codes + per-node pad. */
    const size_t cap = ZXC_HUF_TABLE_SIZE + 2 * n + 4096;
    uint8_t* enc = (uint8_t*)malloc(cap);
    uint8_t* dec = (uint8_t*)malloc(n + ZXC_PAD_SIZE);
    uint8_t* scr = (uint8_t*)malloc(n + ZXC_PIVCO_SCRATCH_PAD);
    int ok = 0;
    if (!enc || !dec || !scr) {
        printf("Failed [%s]: alloc\n", label);
        goto done;
    }

    const int written = zxc_huf_encode_section(literals, n, freq, code_len, enc, cap);
    if (written < 0) {
        printf("Failed [%s]: encode_section -> %d\n", label, written);
        goto done;
    }

    /* The Lagrangian selector prices candidates with zxc_huf_calc_size; it must
     * predict the EXACT encoded size (with the 128-byte lengths header). */
    const size_t est = zxc_huf_calc_size(freq, code_len, 1);
    if (est != (size_t)written) {
        printf("Failed [%s]: calc_size %zu != encoded %d\n", label, est, written);
        goto done;
    }

    const int rc = zxc_huf_decode_section(enc, (size_t)written, dec, n, scr);
    if (rc != ZXC_OK) {
        printf("Failed [%s]: decode_section -> %d\n", label, rc);
        goto done;
    }

    if (memcmp(literals, dec, n) != 0) {
        printf("Failed [%s]: roundtrip mismatch\n", label);
        goto done;
    }

    printf("  [PASS] %s (n=%zu, encoded=%d B, ratio=%.1f%%)\n", label, n, written,
           100.0 * (double)written / (double)n);
    ok = 1;
done:
    free(enc);
    free(dec);
    free(scr);
    return ok;
}

int test_huffman_codec() {
    printf("=== TEST: Unit - Huffman Codec (build/encode/decode roundtrip) ===\n");

    const size_t N = 8192;
    uint8_t* buf = (uint8_t*)malloc(N);
    if (!buf) return 0;

    /* Case 1: heavily skewed (90% one byte, 10% noise). */
    for (size_t i = 0; i < N; i++) buf[i] = (zxc_test_rand() % 10 == 0) ? (uint8_t)(zxc_test_rand() & 0xFF) : 'A';
    if (!huf_roundtrip_case("Skewed (90% 'A')", buf, N)) {
        free(buf);
        return 0;
    }

    /* Case 2: uniform random - Huffman should be near no-op (~1 byte/sym). */
    for (size_t i = 0; i < N; i++) buf[i] = (uint8_t)(zxc_test_rand() & 0xFF);
    if (!huf_roundtrip_case("Uniform random", buf, N)) {
        free(buf);
        return 0;
    }

    /* Case 3: two-symbol alphabet - best case, ~1 bit/symbol. */
    for (size_t i = 0; i < N; i++) buf[i] = (zxc_test_rand() & 1) ? 'X' : 'Y';
    if (!huf_roundtrip_case("Two-symbol alphabet", buf, N)) {
        free(buf);
        return 0;
    }

    /* Case 4: single-symbol - degenerate but must still roundtrip. */
    for (size_t i = 0; i < N; i++) buf[i] = 'Z';
    if (!huf_roundtrip_case("Single-symbol", buf, N)) {
        free(buf);
        return 0;
    }

    /* Case 5: small block (just above the min-literals threshold). */
    for (size_t i = 0; i < ZXC_HUF_MIN_LITERALS; i++)
        buf[i] = (zxc_test_rand() % 4 == 0) ? (uint8_t)(zxc_test_rand() & 0xFF) : 'k';
    if (!huf_roundtrip_case("Small block at threshold", buf, ZXC_HUF_MIN_LITERALS)) {
        free(buf);
        return 0;
    }

    free(buf);
    printf("PASS\n\n");
    return 1;
}

/* --------------------------------------------------------------------------
 * Encoder-side flat/length nudge (zxc_huf_nudge_code_lengths)
 * -------------------------------------------------------------------------- */

/* Cross-check the nudge's buddy-decomposition cost model against the REAL
 * tree the decoder builds: bits and modeled level-touches must match exactly,
 * with the same flat/leaf-pair/deep-flat weighting as zxc_pivco_decode_core.
 * This pins the "flat coverage from bl_count[] alone" math to the shipping
 * flat detector (zxc_pivco_tree_build) for any valid length vector. */
static int nudge_cost_matches_tree(const char* label, const uint8_t* code_len,
                                   const uint32_t* freq) {
    uint64_t bits = 0;
    uint64_t touches = 0;
    zxc_huf_nudge_cost(code_len, freq, &bits, &touches);

    uint8_t packed[ZXC_HUF_TABLE_SIZE];
    zxc_huf_pack_lengths(code_len, packed);
    static zxc_pivco_tree_t tree;
    static zxc_pivco_decode_aux_t aux;
    static uint32_t codes[ZXC_HUF_NUM_SYMBOLS];
    uint8_t len2[ZXC_HUF_NUM_SYMBOLS];
    if (zxc_huf_dict_tree_build(packed, &tree, codes, len2, &aux) != ZXC_OK) {
        printf("Failed [%s]: reference tree build rejected the lengths\n", label);
        return 0;
    }

    uint32_t count[ZXC_PIVCO_MAX_NODES];
    for (int i = tree.n_nodes - 1; i >= 0; i--) {
        const int nid = tree.bfs[i];
        if (tree.nd[nid].sym >= 0) {
            count[nid] = freq[tree.nd[nid].sym];
        } else {
            uint32_t c = 0;
            if (tree.nd[nid].child[0] >= 0) c += count[tree.nd[nid].child[0]];
            if (tree.nd[nid].child[1] >= 0) c += count[tree.nd[nid].child[1]];
            count[nid] = c;
        }
    }

    uint64_t rbits = 0;
    for (int s = 0; s < ZXC_HUF_NUM_SYMBOLS; s++) rbits += (uint64_t)code_len[s] * freq[s];
    uint64_t rtouches = 0;
    for (int i = 0; i < tree.n_nodes; i++) {
        const int nid = tree.bfs[i];
        if (tree.covered[nid]) continue;
        if (tree.nd[nid].sym >= 0) {
            /* Lone leaf memset; leaf-pair children are emitted by the parent. */
            if (!aux.skip[nid]) rtouches += count[nid];
        } else if (tree.flat_d[nid]) {
            uint64_t t = 1;
            if (tree.flat_d[nid] > ZXC_HUF_NUDGE_FLAT_SIMD_MAX) t += ZXC_HUF_NUDGE_DEEP_FLAT_PENALTY;
            rtouches += (uint64_t)count[nid] * t;
        } else {
            rtouches += count[nid]; /* merge node (leaf-pair parents included) */
        }
    }
    rtouches += (uint64_t)ZXC_HUF_NUDGE_LEVEL_COST * (uint64_t)(tree.max_depth + 1);

    if (bits != rbits || touches != rtouches) {
        printf("Failed [%s]: model (bits=%llu touches=%llu) != tree (bits=%llu touches=%llu)\n",
               label, (unsigned long long)bits, (unsigned long long)touches,
               (unsigned long long)rbits, (unsigned long long)rtouches);
        return 0;
    }
    return 1;
}

/* Full nudge pipeline on one distribution and cap: build -> nudge -> validate
 * structure, the calc_size == encode invariant, the decode roundtrip, the
 * adoption guard arithmetic, determinism, and the cost model on both the
 * baseline and the (possibly) adopted vector. */
static int huf_nudge_case(const char* label, const uint8_t* literals, size_t n_lit,
                          const int max_code_len) {
    uint32_t freq[ZXC_HUF_NUM_SYMBOLS] = {0};
    for (size_t i = 0; i < n_lit; i++) freq[literals[i]]++;

    uint8_t base_len[ZXC_HUF_NUM_SYMBOLS];
    if (zxc_huf_build_code_lengths(freq, base_len, NULL, max_code_len) != ZXC_OK) {
        printf("Failed [%s]: build_code_lengths\n", label);
        return 0;
    }
    if (!nudge_cost_matches_tree(label, base_len, freq)) return 0;

    uint8_t nudged[ZXC_HUF_NUM_SYMBOLS];
    memcpy(nudged, base_len, sizeof(nudged));
    const int adopted = zxc_huf_nudge_code_lengths(freq, nudged, NULL, max_code_len);

    /* Determinism: a second independent run must reproduce the result. */
    uint8_t nudged2[ZXC_HUF_NUM_SYMBOLS];
    memcpy(nudged2, base_len, sizeof(nudged2));
    const int adopted2 = zxc_huf_nudge_code_lengths(freq, nudged2, NULL, max_code_len);
    if (adopted != adopted2 || memcmp(nudged, nudged2, sizeof(nudged)) != 0) {
        printf("Failed [%s]: nudge is not deterministic\n", label);
        return 0;
    }

    if (!adopted) {
        if (memcmp(nudged, base_len, sizeof(nudged)) != 0) {
            printf("Failed [%s]: rejected nudge modified the lengths\n", label);
            return 0;
        }
    } else {
        /* Structural rails: cap respected, every live symbol keeps a code.
         * (Coarse-DP candidates may add zero-freq ghost leaves, so a length
         * on a freq == 0 symbol is legal; the reverse is not.) */
        for (int s = 0; s < ZXC_HUF_NUM_SYMBOLS; s++) {
            if (nudged[s] > max_code_len || (freq[s] != 0 && nudged[s] == 0)) {
                printf("Failed [%s]: adopted lengths invalid at sym %d (len=%d)\n", label, s,
                       nudged[s]);
                return 0;
            }
        }
        /* The adoption guard must hold on the exact model costs. */
        uint64_t b0, t0, b1, t1;
        zxc_huf_nudge_cost(base_len, freq, &b0, &t0);
        zxc_huf_nudge_cost(nudged, freq, &b1, &t1);
        if (b1 * 1000 > b0 * ZXC_HUF_NUDGE_BITS_PERMIL || t1 * 256 > t0 * ZXC_HUF_NUDGE_MERGE_Q8) {
            printf("Failed [%s]: adopted candidate violates the guard "
                   "(bits %llu->%llu, touches %llu->%llu)\n",
                   label, (unsigned long long)b0, (unsigned long long)b1, (unsigned long long)t0,
                   (unsigned long long)t1);
            return 0;
        }
        if (!nudge_cost_matches_tree(label, nudged, freq)) return 0;
    }

    /* Selector/encoder invariant + decode roundtrip on the final vector. */
    const size_t cap = ZXC_HUF_TABLE_SIZE + 2 * n_lit + 4096;
    uint8_t* enc = (uint8_t*)malloc(cap);
    uint8_t* dec = (uint8_t*)malloc(n_lit + ZXC_PAD_SIZE);
    uint8_t* scr = (uint8_t*)malloc(n_lit + ZXC_PIVCO_SCRATCH_PAD);
    int ok = 0;
    if (!enc || !dec || !scr) {
        printf("Failed [%s]: alloc\n", label);
        goto done;
    }
    {
        const int written = zxc_huf_encode_section(literals, n_lit, freq, nudged, enc, cap);
        if (written < 0) {
            printf("Failed [%s]: encode_section -> %d\n", label, written);
            goto done;
        }
        const size_t est = zxc_huf_calc_size(freq, nudged, 1);
        if (est != (size_t)written) {
            printf("Failed [%s]: calc_size %zu != encoded %d\n", label, est, written);
            goto done;
        }
        if (zxc_huf_decode_section(enc, (size_t)written, dec, n_lit, scr) != ZXC_OK ||
            memcmp(literals, dec, n_lit) != 0) {
            printf("Failed [%s]: nudged roundtrip mismatch\n", label);
            goto done;
        }
    }
    printf("  [PASS] %s (cap=%d, %s)\n", label, max_code_len, adopted ? "adopted" : "kept");
    ok = 1;
done:
    free(enc);
    free(dec);
    free(scr);
    return ok;
}

int test_huffman_nudge() {
    printf("=== TEST: Unit - Huffman flat/length nudge (model + guard + roundtrip) ===\n");

    const size_t N = 16384;
    uint8_t* buf = (uint8_t*)malloc(N);
    if (!buf) return 0;
    int ok = 1;

    /* Geometric: deep 11-bit trees at ULTRA cap, the nudge's main target. */
    for (size_t i = 0; i < N; i++) {
        uint32_t r = zxc_test_rand();
        int b = 0;
        while (b < 17 && (r & 1)) {
            b++;
            r >>= 1;
        }
        buf[i] = (uint8_t)b;
    }
    ok &= huf_nudge_case("Geometric", buf, N, ZXC_HUF_MAX_CODE_LEN_ULTRA);
    ok &= huf_nudge_case("Geometric cap8", buf, N, ZXC_HUF_MAX_CODE_LEN_DENSITY);

    /* Zipf over the full alphabet: ragged class counts, boundary rounding. */
    for (size_t i = 0; i < N; i++) {
        const uint32_t r = zxc_test_rand() % 6000;
        uint32_t s = 0;
        uint32_t acc = 0;
        while (s < 255) {
            acc += 1000 / (s + 1);
            if (r < acc) break;
            s++;
        }
        buf[i] = (uint8_t)s;
    }
    ok &= huf_nudge_case("Zipf", buf, N, ZXC_HUF_MAX_CODE_LEN_ULTRA);
    ok &= huf_nudge_case("Zipf cap8", buf, N, ZXC_HUF_MAX_CODE_LEN_DENSITY);

    /* Text-like: few dozen symbols, mild skew (the common literal section). */
    for (size_t i = 0; i < N; i++) {
        const uint32_t r = zxc_test_rand();
        buf[i] = (uint8_t)('a' + (((r & 0xFF) * ((r >> 8) & 0xFF)) >> 11));
    }
    ok &= huf_nudge_case("Text-like", buf, N, ZXC_HUF_MAX_CODE_LEN_ULTRA);

    /* Uniform 256: already a single flat root; the nudge must keep baseline. */
    for (size_t i = 0; i < N; i++) buf[i] = (uint8_t)(i & 0xFF);
    {
        uint32_t freq[ZXC_HUF_NUM_SYMBOLS] = {0};
        for (size_t i = 0; i < N; i++) freq[buf[i]]++;
        uint8_t len[ZXC_HUF_NUM_SYMBOLS];
        if (zxc_huf_build_code_lengths(freq, len, NULL, ZXC_HUF_MAX_CODE_LEN_DENSITY) != ZXC_OK) {
            ok = 0;
        } else {
            uint8_t kept[ZXC_HUF_NUM_SYMBOLS];
            memcpy(kept, len, sizeof(kept));
            if (zxc_huf_nudge_code_lengths(freq, len, NULL, ZXC_HUF_MAX_CODE_LEN_DENSITY) != 0 ||
                memcmp(kept, len, sizeof(kept)) != 0) {
                printf("Failed [Uniform-256]: expected a no-op nudge\n");
                ok = 0;
            } else {
                printf("  [PASS] Uniform-256 no-op\n");
            }
        }
    }

    /* Degenerate alphabets (n = 1, 2, 3): must never be touched. */
    for (int nsym = 1; nsym <= 3; nsym++) {
        uint32_t freq[ZXC_HUF_NUM_SYMBOLS] = {0};
        for (int s = 0; s < nsym; s++) freq[(uint8_t)('A' + s)] = (uint32_t)(100 * (s + 1));
        uint8_t len[ZXC_HUF_NUM_SYMBOLS];
        uint8_t kept[ZXC_HUF_NUM_SYMBOLS];
        if (zxc_huf_build_code_lengths(freq, len, NULL, ZXC_HUF_MAX_CODE_LEN_DENSITY) != ZXC_OK) {
            ok = 0;
            continue;
        }
        memcpy(kept, len, sizeof(kept));
        if (zxc_huf_nudge_code_lengths(freq, len, NULL, ZXC_HUF_MAX_CODE_LEN_DENSITY) != 0 ||
            memcmp(kept, len, sizeof(kept)) != 0) {
            printf("Failed [Degenerate n=%d]: expected untouched lengths\n", nsym);
            ok = 0;
        }
    }
    if (ok) printf("  [PASS] Degenerate n=1..3 untouched\n");

    /* Fuzz: random alphabets/frequencies; the cost model must match the real
     * tree for every baseline AND every adopted vector, at both caps (large
     * alphabets exercise the coarse-DP tiers and their ghost padding). */
    {
        int checked = 0;
        int adopted_cnt = 0;
        for (int it = 0; it < 300 && ok; it++) {
            uint32_t freq[ZXC_HUF_NUM_SYMBOLS] = {0};
            const int nsym = 4 + (int)(zxc_test_rand() % 253);
            for (int s = 0; s < nsym; s++)
                freq[(uint8_t)(zxc_test_rand() & 0xFF)] = 1 + (zxc_test_rand() & 0xFFFFF);
            const int cap = (it & 1) ? ZXC_HUF_MAX_CODE_LEN_ULTRA : ZXC_HUF_MAX_CODE_LEN_DENSITY;
            uint8_t len[ZXC_HUF_NUM_SYMBOLS];
            if (zxc_huf_build_code_lengths(freq, len, NULL, cap) != ZXC_OK) {
                ok = 0;
                break;
            }
            if (!nudge_cost_matches_tree("Fuzz baseline", len, freq)) {
                ok = 0;
                break;
            }
            checked++;
            if (zxc_huf_nudge_code_lengths(freq, len, NULL, cap)) {
                adopted_cnt++;
                if (!nudge_cost_matches_tree("Fuzz nudged", len, freq)) {
                    ok = 0;
                    break;
                }
            }
        }
        if (ok)
            printf("  [PASS] Fuzz: %d histograms cross-checked, %d nudges adopted\n", checked,
                   adopted_cnt);
    }

    free(buf);
    if (!ok) return 0;
    printf("PASS\n\n");
    return 1;
}

/* Round-trip the shared-table (dictionary) Huffman section codec: encode with
 * external code lengths and NO 128-byte header, decode through a prebuilt
 * table -- the enc_lit == 3 wire path (FORMAT.md Sec 5.2.2). */
static int huf_dict_roundtrip_case(const char* label, const uint8_t* literals, size_t n) {
    uint32_t freq[ZXC_HUF_NUM_SYMBOLS] = {0};
    for (size_t i = 0; i < n; i++) freq[literals[i]]++;

    uint8_t code_len[ZXC_HUF_NUM_SYMBOLS];
    if (zxc_huf_build_code_lengths(freq, code_len, NULL, ZXC_HUF_MAX_CODE_LEN_DENSITY) != ZXC_OK) {
        printf("Failed [%s]: build_code_lengths\n", label);
        return 0;
    }

    /* Round the lengths through the 128-byte packed form, as a .zxd does. */
    uint8_t packed[ZXC_HUF_TABLE_SIZE];
    uint8_t unpacked[ZXC_HUF_NUM_SYMBOLS];
    zxc_huf_pack_lengths(code_len, packed);
    if (zxc_huf_unpack_lengths(packed, unpacked) != ZXC_OK ||
        memcmp(code_len, unpacked, sizeof(code_len)) != 0) {
        printf("Failed [%s]: pack/unpack lengths roundtrip\n", label);
        return 0;
    }

    /* Tree-at-attach: prebuild the shared table's tree/codes/decoder tables
     * once, as zxc_cctx_attach_dict_huf does; the dict codec entry points
     * take them. */
    zxc_pivco_tree_t tree;
    zxc_pivco_decode_aux_t aux;
    uint32_t codes[ZXC_HUF_NUM_SYMBOLS];
    uint8_t tree_len[ZXC_HUF_NUM_SYMBOLS];
    if (zxc_huf_dict_tree_build(packed, &tree, codes, tree_len, &aux) != ZXC_OK ||
        memcmp(tree_len, code_len, sizeof(tree_len)) != 0) {
        printf("Failed [%s]: dict_tree_build\n", label);
        return 0;
    }

    const size_t cap = ZXC_HUF_TABLE_SIZE + 2 * n + 4096;
    uint8_t* enc = (uint8_t*)malloc(cap);
    uint8_t* enc_blk = (uint8_t*)malloc(cap);
    uint8_t* dec = (uint8_t*)malloc(n + ZXC_PAD_SIZE);
    uint8_t* scr = (uint8_t*)malloc(n + ZXC_PIVCO_SCRATCH_PAD);
    if (!enc || !enc_blk || !dec || !scr) {
        printf("Failed [%s]: alloc\n", label);
        goto fail;
    }

    const int written =
        zxc_huf_encode_section_dict(literals, n, freq, code_len, &tree, codes, enc, cap);
    if (written < 0) {
        printf("Failed [%s]: encode_section_dict -> %d\n", label, written);
        goto fail;
    }

    /* Header-less variant of the size estimator (with_header = 0) must match. */
    if (zxc_huf_calc_size(freq, code_len, 0) != (size_t)written) {
        printf("Failed [%s]: dict calc_size != encoded\n", label);
        goto fail;
    }

    /* Same lengths, same bitstreams: the dict section must be exactly the
     * per-block section minus its 128-byte lengths header. */
    const int written_blk = zxc_huf_encode_section(literals, n, freq, code_len, enc_blk, cap);
    if (written_blk != written + (int)ZXC_HUF_TABLE_SIZE ||
        memcmp(enc, enc_blk + ZXC_HUF_TABLE_SIZE, (size_t)written) != 0) {
        printf("Failed [%s]: dict section != per-block section minus header\n", label);
        goto fail;
    }

    if (zxc_huf_decode_section_dict(enc, (size_t)written, dec, n, &tree, &aux, scr) != ZXC_OK ||
        memcmp(literals, dec, n) != 0) {
        printf("Failed [%s]: decode_section_dict roundtrip mismatch\n", label);
        goto fail;
    }

    /* Error paths: truncated payload, undersized dst_cap. */
    if (zxc_huf_decode_section_dict(enc, 0, dec, n, &tree, &aux, scr) == ZXC_OK) {
        printf("Failed [%s]: truncated payload accepted\n", label);
        goto fail;
    }
    if (zxc_huf_encode_section_dict(literals, n, freq, code_len, &tree, codes, enc, 4) !=
        ZXC_ERROR_DST_TOO_SMALL) {
        printf("Failed [%s]: undersized dst_cap not rejected\n", label);
        goto fail;
    }

    free(enc);
    free(enc_blk);
    free(dec);
    free(scr);
    printf("  [PASS] %s (n=%zu, encoded=%d B, header saved=%d B)\n", label, n, written,
           (int)ZXC_HUF_TABLE_SIZE);
    return 1;

fail:
    free(enc);
    free(enc_blk);
    free(dec);
    free(scr);
    return 0;
}

int test_huffman_codec_dict() {
    printf("=== TEST: Unit - Huffman Codec, shared dictionary table (enc_lit == 3) ===\n");

    const size_t N = 8192;
    uint8_t* buf = (uint8_t*)malloc(N);
    if (!buf) return 0;

    /* Skewed text-like distribution: the shared-table sweet spot. */
    for (size_t i = 0; i < N; i++) buf[i] = (zxc_test_rand() % 10 == 0) ? (uint8_t)(zxc_test_rand() & 0x7F) : 'A';
    if (!huf_dict_roundtrip_case("Skewed (90% 'A')", buf, N)) {
        free(buf);
        return 0;
    }

    /* Two-symbol alphabet: ~1 bit/symbol, headerless gain is maximal. */
    for (size_t i = 0; i < N; i++) buf[i] = (zxc_test_rand() & 1) ? 'X' : 'Y';
    if (!huf_dict_roundtrip_case("Two-symbol alphabet", buf, N)) {
        free(buf);
        return 0;
    }

    /* Small block: where the 128-byte header would dominate per-block cost. */
    for (size_t i = 0; i < ZXC_HUF_MIN_LITERALS; i++)
        buf[i] = (zxc_test_rand() % 4 == 0) ? (uint8_t)('a' + (zxc_test_rand() % 26)) : 'k';
    if (!huf_dict_roundtrip_case("Small block at threshold", buf, ZXC_HUF_MIN_LITERALS)) {
        free(buf);
        return 0;
    }

    /* A literal with NO code in the shared table must be refused by the
     * encoder (the validity check that triggers the per-block fallback). */
    {
        uint32_t freq[ZXC_HUF_NUM_SYMBOLS] = {0};
        for (size_t i = 0; i < 256; i++) buf[i] = (zxc_test_rand() & 1) ? 'X' : 'Y';
        for (size_t i = 0; i < 256; i++) freq[buf[i]]++;
        uint8_t code_len[ZXC_HUF_NUM_SYMBOLS];
        if (zxc_huf_build_code_lengths(freq, code_len, NULL, ZXC_HUF_MAX_CODE_LEN_DENSITY) != ZXC_OK) {
            free(buf);
            return 0;
        }
        buf[100] = '!'; /* unseen in training: no code assigned */
        freq['!']++;    /* keep the histogram in sync with the mutated buffer */
        uint8_t enc[1024];
        zxc_pivco_tree_t tree;
        zxc_pivco_decode_aux_t aux;
        uint32_t codes[ZXC_HUF_NUM_SYMBOLS];
        uint8_t packed[ZXC_HUF_TABLE_SIZE];
        uint8_t tree_len[ZXC_HUF_NUM_SYMBOLS];
        zxc_huf_pack_lengths(code_len, packed);
        if (zxc_huf_dict_tree_build(packed, &tree, codes, tree_len, &aux) != ZXC_OK) {
            printf("Failed: dict_tree_build (code-less literal case)\n");
            free(buf);
            return 0;
        }
        if (zxc_huf_encode_section_dict(buf, 256, freq, code_len, &tree, codes, enc,
                                        sizeof(enc)) != ZXC_ERROR_CORRUPT_DATA) {
            printf("Failed: code-less literal not rejected by encode_section_dict\n");
            free(buf);
            return 0;
        }
        printf("  [PASS] code-less literal rejected (per-block fallback trigger)\n");
    }

    free(buf);
    printf("PASS\n\n");
    return 1;
}

/* Regression: a degenerate single-symbol table must carry code_len == 1
 * (FORMAT.md, decoder validation requirements). The v6 decoder rejected a
 * lone symbol with a longer length; the v7 rewrite briefly accepted it. */
int test_huffman_single_symbol_validation() {
    printf("=== TEST: Unit - Huffman single-symbol table validation ===\n");

    zxc_pivco_tree_t tree;
    zxc_pivco_decode_aux_t aux;
    uint32_t codes[ZXC_HUF_NUM_SYMBOLS];
    uint8_t tree_len[ZXC_HUF_NUM_SYMBOLS];
    uint8_t code_len[ZXC_HUF_NUM_SYMBOLS];
    uint8_t packed[ZXC_HUF_TABLE_SIZE];

    /* Lone symbol with code length 1: the only legal degenerate form. */
    memset(code_len, 0, sizeof(code_len));
    code_len['A'] = 1;
    zxc_huf_pack_lengths(code_len, packed);
    if (zxc_huf_dict_tree_build(packed, &tree, codes, tree_len, &aux) != ZXC_OK) {
        printf("Failed: single symbol with code_len=1 must be accepted\n");
        return 0;
    }
    printf("  [PASS] single symbol, code_len=1 accepted\n");

    /* Lone symbol with any longer length is declared corrupt by the format. */
    for (int len = 2; len <= ZXC_HUF_MAX_CODE_LEN_ULTRA; len++) {
        memset(code_len, 0, sizeof(code_len));
        code_len['A'] = (uint8_t)len;
        zxc_huf_pack_lengths(code_len, packed);
        if (zxc_huf_dict_tree_build(packed, &tree, codes, tree_len, &aux) !=
            ZXC_ERROR_CORRUPT_DATA) {
            printf("Failed: single symbol with code_len=%d must be rejected\n", len);
            return 0;
        }
    }
    printf("  [PASS] single symbol, code_len 2..%d rejected\n", ZXC_HUF_MAX_CODE_LEN_ULTRA);

    printf("PASS\n\n");
    return 1;
}

// Checks that the EOF block is correctly appended
int test_eof_block_structure() {
    printf("=== TEST: Unit - EOF Block Structure ===\n");

    const char* input = "test";
    size_t src_size = 4;
    size_t max_dst_size = (size_t)zxc_compress_bound(src_size);
    uint8_t* compressed = malloc(max_dst_size);
    if (!compressed) return 0;

    zxc_compress_opts_t _co26 = {.level = 1, .checksum_enabled = 0};
    int64_t comp_size = zxc_compress(input, src_size, compressed, max_dst_size, &_co26);
    if (comp_size <= 0) {
        printf("Failed: Compression returned 0\n");
        free(compressed);
        return 0;
    }

    // Validating Footer and EOF Block
    // Total Overhead: 12 bytes (Footer) + 8 bytes (EOF Header) = 20 bytes
    if (comp_size < 20) {
        printf("Failed: Compressed size too small for Footer + EOF (%lld)\n", (long long)comp_size);
        free(compressed);
        return 0;
    }

    // 1. Verify 12-byte Footer
    // Structure: [SrcSize (8)] + [Hash (4)]
    const uint8_t* footer_ptr = compressed + comp_size - 12;
    uint32_t f_src_low = zxc_le32(footer_ptr);       // Should be 4
    uint32_t f_src_high = zxc_le32(footer_ptr + 4);  // Should be 0
    uint32_t f_hash = zxc_le32(footer_ptr + 8);      // Should be 0 (checksum disabled)

    if (f_src_low != 4 || f_src_high != 0 || f_hash != 0) {
        printf("Failed: Footer mismatch. Src: %u, Hash: %u\n", f_src_low, f_hash);
        free(compressed);
        return 0;
    }

    // 2. Verify EOF Block Header (8 bytes)
    // Should be immediately before the footer
    const uint8_t* eof_ptr = compressed + comp_size - 20;
    uint8_t expected[8] = {0xFF, 0, 0, 0, 0, 0, 0, 0};
    expected[7] = zxc_hash8(expected);

    if (memcmp(eof_ptr, expected, 8) != 0) {
        printf(
            "Failed: EOF block mismatch.\nExpected: %02X %02X %02X ... %02X\nGot:      %02X %02X "
            "%02X ... %02X\n",
            expected[0], expected[1], expected[2], expected[7], eof_ptr[0], eof_ptr[1], eof_ptr[2],
            eof_ptr[7]);
        free(compressed);
        return 0;
    }

    printf("PASS\n\n");
    free(compressed);
    return 1;
}

int test_header_checksum() {
    printf("Running test_header_checksum...\n");

    uint8_t header_buf[ZXC_BLOCK_HEADER_SIZE];
    zxc_block_header_t bh_in = {.block_type = ZXC_BLOCK_GLO,
                                .block_flags = 0,
                                .reserved = 0,
                                .header_crc = 0,
                                .comp_size = 1024};

    // 1. Write Header
    if (zxc_write_block_header(header_buf, ZXC_BLOCK_HEADER_SIZE, &bh_in) !=
        ZXC_BLOCK_HEADER_SIZE) {
        printf("  [FAIL] zxc_write_block_header failed\n");
        return 0;
    }

    // Verify manually that checksum byte is non-zero (highly likely)
    if (header_buf[7] == 0) {
        // It's technically possible but very unlikely with a good hash
        printf("  [WARN] Checksum is 0 (unlikely but possible)\n");
    }

    // 2. Read Header (Valid)
    zxc_block_header_t bh_out;
    if (zxc_read_block_header(header_buf, ZXC_BLOCK_HEADER_SIZE, &bh_out) != 0) {
        printf("  [FAIL] zxc_read_block_header failed on valid input\n");
        return 0;
    }

    if (bh_out.block_type != bh_in.block_type || bh_out.comp_size != bh_in.comp_size ||
        bh_out.header_crc != header_buf[7]) {
        printf("  [FAIL] Read data mismatch\n");
        return 0;
    }

    // 3. Corrupt Header Checksum
    uint8_t original_crc = header_buf[7];
    header_buf[7] = ~original_crc;  // Flip bits
    if (zxc_read_block_header(header_buf, ZXC_BLOCK_HEADER_SIZE, &bh_out) == 0) {
        printf("  [FAIL] zxc_read_block_header should have failed on corrupted CRC\n");
        return 0;
    }
    header_buf[7] = original_crc;  // Restore

    // 4. Corrupt Header Content
    header_buf[0] = ZXC_BLOCK_RAW;  // Change type
    if (zxc_read_block_header(header_buf, ZXC_BLOCK_HEADER_SIZE, &bh_out) == 0) {
        printf("  [FAIL] zxc_read_block_header should have failed on corrupted content\n");
        return 0;
    }

    printf("PASS\n\n");
    return 1;
}

// 5. Test Global Checksum Order Sensitivity
// Ensures that swapping two blocks (even if valid individually) triggers a global checksum failure.
int test_global_checksum_order() {
    printf("TEST: Global Checksum Order Sensitivity... ");

    // 1. Create input data withDISTINCT patterns for 2 blocks (so blocks are different)
    // ZXC_BLOCK_SIZE_DEFAULT is 256KB. We need > 256KB. Let's use 600KB.
    size_t input_sz = 600 * 1024;
    uint8_t* val_buf = malloc(input_sz);
    if (!val_buf) return 0;

    // Fill Block 1 with 0xAA, Block 2 with 0xBB, Block 3 with 0xCC...
    memset(val_buf, 0xAA, 256 * 1024);
    memset(val_buf + 256 * 1024, 0xBB, 256 * 1024);
    memset(val_buf + 512 * 1024, 0xCC, input_sz - 512 * 1024);

    FILE* f_in = tmpfile();
    FILE* f_comp = tmpfile();
    fwrite(val_buf, 1, input_sz, f_in);
    rewind(f_in);

    // 2. Compress with Checksum Enabled
    zxc_compress_opts_t _sco27 = {.n_threads = 1, .level = 1, .checksum_enabled = 1};
    zxc_stream_compress(f_in, f_comp, &_sco27);

    // 3. Read compressed data to memory
    long comp_sz = ftell(f_comp);
    rewind(f_comp);
    uint8_t* comp_buf = malloc((size_t)comp_sz);
    if (fread(comp_buf, 1, comp_sz, f_comp) != (size_t)comp_sz) {
        printf("[FAIL] Failed to read compressed data\n");
        free(val_buf);
        free(comp_buf);
        fclose(f_in);
        fclose(f_comp);
        return 0;
    }

    // 4. Parse Blocks to identify Block 1 and Block 2
    // File Header: ZXC_FILE_HEADER_SIZE bytes
    size_t off1 = ZXC_FILE_HEADER_SIZE;
    // Parse Block 1 Header
    zxc_block_header_t bh1;
    zxc_read_block_header(comp_buf + off1, ZXC_BLOCK_HEADER_SIZE, &bh1);
    size_t len1 = ZXC_BLOCK_HEADER_SIZE + bh1.comp_size + ZXC_BLOCK_CHECKSUM_SIZE;

    size_t off2 = off1 + len1;
    // Parse Block 2 Header
    zxc_block_header_t bh2;
    zxc_read_block_header(comp_buf + off2, ZXC_BLOCK_HEADER_SIZE, &bh2);
    size_t len2 = ZXC_BLOCK_HEADER_SIZE + bh2.comp_size + ZXC_BLOCK_CHECKSUM_SIZE;

    // Ensure we have at least 2 full blocks + EOF + Global Checksum
    if (off2 + len2 > (size_t)comp_sz) {
        printf("[FAIL] Compressed size too small for test\n");
        free(val_buf);
        free(comp_buf);
        fclose(f_in);
        fclose(f_comp);
        return 0;
    }

    // 5. Swap Block 1 and Block 2
    // To safely swap, we need a new buffer
    uint8_t* swapped_buf = malloc((size_t)comp_sz);

    // Copy File Header
    // Copy File Header
    memcpy(swapped_buf, comp_buf, ZXC_FILE_HEADER_SIZE);
    size_t w_off = ZXC_FILE_HEADER_SIZE;

    // Write Block 2 first
    memcpy(swapped_buf + w_off, comp_buf + off2, len2);
    w_off += len2;

    // Write Block 1 second
    memcpy(swapped_buf + w_off, comp_buf + off1, len1);
    w_off += len1;

    // Write remaining data (EOF block + Global Checksum)
    size_t remaining_off = off2 + len2;
    size_t remaining_len = comp_sz - remaining_off;
    memcpy(swapped_buf + w_off, comp_buf + remaining_off, remaining_len);

    // 6. Write to File for Decompression
    FILE* f_bad = tmpfile();
    fwrite(swapped_buf, 1, (size_t)comp_sz, f_bad);
    rewind(f_bad);

    // 7. Attempt Decompression
    FILE* f_out = tmpfile();
    zxc_decompress_opts_t _sdo28 = {.n_threads = 1, .checksum_enabled = 1};
    int64_t res = zxc_stream_decompress(f_bad, f_out, &_sdo28);

    fclose(f_in);
    fclose(f_comp);
    fclose(f_bad);
    fclose(f_out);
    free(val_buf);
    free(comp_buf);
    free(swapped_buf);

    if (res >= 0) {
        printf("  [FAIL] zxc_stream_decompress unexpectedly succeeded on swapped blocks\n");
        return 0;
    }

    printf("PASS\n\n");
    return 1;
}

/* Builds a header with the given chunk-size code, fixes the CRC16, and returns
 * zxc_read_file_header's verdict (block_size out via *bs). */
static int chunk_code_verdict(uint8_t code, size_t *bs) {
    uint8_t hdr[ZXC_FILE_HEADER_SIZE];
    memset(hdr, 0, sizeof(hdr));
    hdr[0] = 0xF5; hdr[1] = 0x2E; hdr[2] = 0xB0; hdr[3] = 0x9C;  // magic (LE)
    hdr[4] = ZXC_FILE_FORMAT_VERSION;                            // version
    hdr[5] = code;                                              // chunk-size code
    hdr[6] = 0;                                                 // flags: no checksum
    uint16_t crc = zxc_hash16(hdr);                            // bytes 14-15 already 0
    hdr[14] = (uint8_t)(crc & 0xFF);
    hdr[15] = (uint8_t)(crc >> 8);
    int has_checksum = -1;
    *bs = 0;
    return zxc_read_file_header(hdr, sizeof(hdr), bs, &has_checksum, NULL);
}

int test_chunk_size_code() {
    printf("=== TEST: Chunk-size code validation ===\n");

    size_t bs = 0;

    // Valid exponent code 19 -> 512 KB.
    int rc = chunk_code_verdict(19, &bs);
    if (rc != ZXC_OK || bs != 512 * 1024) {
        printf("  [FAIL] code 19: rc=%d (%s), block_size=%zu\n", rc, zxc_error_name(rc), bs);
        return 0;
    }
    printf("  [PASS] Code 19 -> 512 KB\n");

    // Legacy code 64 was accepted in v5 (mapped to 256 KB) but is removed in v6.
    rc = chunk_code_verdict(64, &bs);
    if (rc != ZXC_ERROR_BAD_BLOCK_SIZE) {
        printf("  [FAIL] legacy code 64: expected %d, got %d\n", ZXC_ERROR_BAD_BLOCK_SIZE, rc);
        return 0;
    }
    printf("  [PASS] Legacy code 64 rejected (ZXC_ERROR_BAD_BLOCK_SIZE)\n");

    // Out-of-range code 99 -> rejected.
    rc = chunk_code_verdict(99, &bs);
    if (rc != ZXC_ERROR_BAD_BLOCK_SIZE) {
        printf("  [FAIL] code 99: expected %d, got %d\n", ZXC_ERROR_BAD_BLOCK_SIZE, rc);
        return 0;
    }
    printf("  [PASS] Code 99 rejected\n");

    printf("PASS\n\n");
    return 1;
}
