package pdf

import "core:slice"
import "core:testing"

// ---------------------------------------------------------------------------
// Test-side encoder
//
// Two pieces, deliberately kept apart. lzw_test_codes runs the textbook LZW
// encoder and produces a list of codes, which depends on nothing but the input
// bytes. lzw_test_pack turns that list into a bit stream, and *that* is where
// the code-width schedule lives. Splitting them means the round-trip tests
// exercise the table logic without the packer's width rule being taken on
// faith: the hand-computed vectors below pin the widths independently.
// ---------------------------------------------------------------------------

@(private = "file")
Bit_Writer :: struct {
    bytes: [dynamic]byte,
    acc:   u32,
    nbits: uint,
}

@(private = "file")
bw_put :: proc(w: ^Bit_Writer, code: u16, width: uint) {
    w.acc = (w.acc << width) | u32(code)
    w.nbits += width
    for w.nbits >= 8 {
        w.nbits -= 8
        append(&w.bytes, u8((w.acc >> w.nbits) & 0xff))
    }
}

@(private = "file")
bw_flush :: proc(w: ^Bit_Writer) {
    if w.nbits > 0 {
        append(&w.bytes, u8((w.acc << (8 - w.nbits)) & 0xff))
        w.nbits = 0
    }
}

// lzw_test_pack packs codes at the widths a decoder will read them at: the
// table size is advanced exactly the way the decoder advances it (one entry
// per code after the first following a clear).
@(private = "file")
lzw_test_pack :: proc(codes: []u16, early_change: bool, allocator := context.allocator) -> [dynamic]byte {
    w := Bit_Writer{bytes = make([dynamic]byte, allocator)}
    next := LZW_FIRST_CODE
    have_prev := false
    for c in codes {
        bw_put(&w, c, lzw_code_width(next, early_change))
        switch c {
        case LZW_CLEAR:
            next = LZW_FIRST_CODE
            have_prev = false
        case LZW_EOD:
            // no table change
        case:
            if have_prev && next <= LZW_MAX_CODE {
                next += 1
            }
            have_prev = true
        }
    }
    bw_flush(&w)
    return w.bytes
}

// lzw_test_codes is the plain TIFF encoder: it adds a table entry immediately
// after emitting a code, which is why its table runs one ahead of a decoder's.
@(private = "file")
lzw_test_codes :: proc(input: []byte, allocator := context.allocator) -> [dynamic]u16 {
    codes := make([dynamic]u16, allocator)
    append(&codes, u16(LZW_CLEAR))
    if len(input) == 0 {
        append(&codes, u16(LZW_EOD))
        return codes
    }

    dict := make(map[u32]u16, 1024, context.allocator)
    defer delete(dict)

    next := u16(LZW_FIRST_CODE)
    omega := u16(input[0])
    for i in 1 ..< len(input) {
        k := input[i]
        key := (u32(omega) << 8) | u32(k)
        if c, found := dict[key]; found {
            omega = c
        } else {
            append(&codes, omega)
            if int(next) <= LZW_MAX_CODE {
                dict[key] = next
                next += 1
            }
            omega = u16(k)
        }
    }
    append(&codes, omega)
    append(&codes, u16(LZW_EOD))
    return codes
}

@(private = "file")
lzw_test_encode :: proc(input: []byte, early_change: bool, allocator := context.allocator) -> [dynamic]byte {
    codes := lzw_test_codes(input, context.allocator)
    defer delete(codes)
    return lzw_test_pack(codes[:], early_change, allocator)
}

// ---------------------------------------------------------------------------
// Published vector: TIFF 6.0 section 13 encodes "-----A---B" as the codes
// 256 45 258 258 65 259 66 257, packed at 9 bits into these nine bytes.
// ---------------------------------------------------------------------------

@(private = "file")
TIFF_CODES  := [?]u16{256, 45, 258, 258, 65, 259, 66, 257}
@(private = "file")
TIFF_PACKED := [?]byte{0x80, 0x0b, 0x60, 0x50, 0x22, 0x0c, 0x0c, 0x85, 0x01}
@(private = "file")
TIFF_PLAIN  :: "-----A---B"

@(test)
test_lzw_tiff_vector :: proc(t: ^testing.T) {
    for ec in ([?]bool{true, false}) {
        out, err := lzw_decode(TIFF_PACKED[:], ec)
        defer delete(out)
        testing.expect_value(t, err, Error.None)
        testing.expect_value(t, string(out), TIFF_PLAIN)
    }
}

// The packer has to reproduce the spec's bytes exactly, or every round-trip
// test below would only be proving the decoder consistent with itself.
@(test)
test_lzw_packer_matches_tiff_vector :: proc(t: ^testing.T) {
    packed := lzw_test_pack(TIFF_CODES[:], true)
    defer delete(packed)
    testing.expect(t, slice.equal(packed[:], TIFF_PACKED[:]), "packer disagrees with TIFF 6.0 vector")

    codes := lzw_test_codes(transmute([]byte)string(TIFF_PLAIN))
    defer delete(codes)
    testing.expect(t, slice.equal(codes[:], TIFF_CODES[:]), "encoder disagrees with TIFF 6.0 vector")
}

// ---------------------------------------------------------------------------
// KwKwK: "AAAAAA" encodes to 256 65 258 259 257, where 258 and 259 are both
// used by the code that defines them. Packed at 9 bits: 80 10 60 50 38 08.
// ---------------------------------------------------------------------------

@(test)
test_lzw_kwkwk :: proc(t: ^testing.T) {
    packed := [?]byte{0x80, 0x10, 0x60, 0x50, 0x38, 0x08}
    out, err := lzw_decode(packed[:])
    defer delete(out)
    testing.expect_value(t, err, Error.None)
    testing.expect_value(t, string(out), "AAAAAA")
}

@(test)
test_lzw_kwkwk_without_predecessor :: proc(t: ^testing.T) {
    // 258 immediately after a clear names a string that cannot exist yet.
    codes := [?]u16{256, 258, 257}
    packed := lzw_test_pack(codes[:], true)
    defer delete(packed)
    out, err := lzw_decode(packed[:])
    defer delete(out)
    testing.expect_value(t, err, Error.Bad_Lzw_Data)
    testing.expect_value(t, len(out), 0)
}

// ---------------------------------------------------------------------------
// Degenerate and empty streams
// ---------------------------------------------------------------------------

@(test)
test_lzw_empty_input :: proc(t: ^testing.T) {
    out, err := lzw_decode(nil)
    defer delete(out)
    testing.expect_value(t, err, Error.None)
    testing.expect_value(t, len(out), 0)
}

@(test)
test_lzw_eod_only :: proc(t: ^testing.T) {
    for ec in ([?]bool{true, false}) {
        codes := [?]u16{256, 257}
        packed := lzw_test_pack(codes[:], ec)
        defer delete(packed)
        out, err := lzw_decode(packed[:], ec)
        defer delete(out)
        // An empty stream is legal; it must not be reported as damaged.
        testing.expect_value(t, err, Error.None)
        testing.expect_value(t, len(out), 0)
    }
}

@(test)
test_lzw_too_short_for_a_code :: proc(t: ^testing.T) {
    // One byte cannot hold a 9-bit code, so nothing at all comes out.
    short := [?]byte{0x80}
    out, err := lzw_decode(short[:])
    defer delete(out)
    testing.expect_value(t, err, Error.Bad_Lzw_Data)
    testing.expect_value(t, len(out), 0)
}

@(test)
test_lzw_undefined_code :: proc(t: ^testing.T) {
    // 300 is far beyond the two entries that exist at that point.
    codes := [?]u16{256, u16('X'), 300, 257}
    packed := lzw_test_pack(codes[:], true)
    defer delete(packed)
    out, err := lzw_decode(packed[:])
    defer delete(out)
    testing.expect_value(t, err, Error.None)
    testing.expect_value(t, string(out), "X")
}

@(test)
test_lzw_missing_eod :: proc(t: ^testing.T) {
    // Everything but the trailing EOD; the decoder must stop at end of input
    // rather than demanding the marker.
    codes := TIFF_CODES[:len(TIFF_CODES) - 1]
    packed := lzw_test_pack(codes, true)
    defer delete(packed)
    out, err := lzw_decode(packed[:])
    defer delete(out)
    testing.expect_value(t, err, Error.None)
    testing.expect_value(t, string(out), TIFF_PLAIN)
}

@(test)
test_lzw_truncated_at_every_offset :: proc(t: ^testing.T) {
    // Truncation drops whole codes off the end, so whatever survives has to be
    // a prefix of the full output -- and nothing may crash or hang.
    whole := TIFF_PLAIN
    for n in 0 ..= len(TIFF_PACKED) {
        out, err := lzw_decode(TIFF_PACKED[:n])
        defer delete(out)
        got := string(out)
        testing.expect(t, len(got) <= len(whole), "truncated input produced too much output")
        testing.expect(t, got == whole[:len(got)], "truncated output is not a prefix of the whole")
        if len(got) > 0 {
            testing.expect_value(t, err, Error.None)
        }
    }
}

@(test)
test_lzw_clear_mid_stream :: proc(t: ^testing.T) {
    // A clear resets the table, so the codes after it mean what they would at
    // the start of a stream.
    codes := [?]u16{256, u16('a'), u16('b'), 258, 256, u16('a'), u16('b'), 258, 257}
    packed := lzw_test_pack(codes[:], true)
    defer delete(packed)
    out, err := lzw_decode(packed[:])
    defer delete(out)
    testing.expect_value(t, err, Error.None)
    // 258 is "ab" both times; without the reset the second 258 would still be
    // "ab" but the entry numbering afterwards would differ.
    testing.expect_value(t, string(out), "abababab")
}

// ---------------------------------------------------------------------------
// Code width and /EarlyChange
// ---------------------------------------------------------------------------

@(test)
test_lzw_code_width_boundaries :: proc(t: ^testing.T) {
    Case :: struct {
        next:  int,
        early: bool,
        width: uint,
    }
    cases := [?]Case{
        {258, true, 9}, {510, true, 9}, {511, true, 10}, {512, true, 10},
        {1022, true, 10}, {1023, true, 11}, {2046, true, 11}, {2047, true, 12},
        {4095, true, 12}, {4096, true, 12},
        {258, false, 9}, {511, false, 9}, {512, false, 10}, {1023, false, 10},
        {1024, false, 11}, {2047, false, 11}, {2048, false, 12}, {4096, false, 12},
    }
    for c in cases {
        testing.expect_value(t, lzw_code_width(c.next, c.early), c.width)
    }
}

// 254 literal codes leave the table holding exactly 511 entries, which is the
// one place /EarlyChange 0 and 1 disagree about how wide the next code is.
@(private = "file")
lzw_width_boundary_codes :: proc(allocator := context.allocator) -> [dynamic]u16 {
    codes := make([dynamic]u16, allocator)
    append(&codes, u16(LZW_CLEAR))
    for i in 0 ..< 255 {
        append(&codes, u16(i))
    }
    append(&codes, u16(LZW_EOD))
    return codes
}

@(test)
test_lzw_early_change_boundary :: proc(t: ^testing.T) {
    codes := lzw_width_boundary_codes()
    defer delete(codes)

    expect := make([]byte, 255)
    defer delete(expect)
    for i in 0 ..< 255 {
        expect[i] = u8(i)
    }

    for ec in ([?]bool{true, false}) {
        packed := lzw_test_pack(codes[:], ec)
        defer delete(packed)
        out, err := lzw_decode(packed[:], ec)
        defer delete(out)
        testing.expect_value(t, err, Error.None)
        testing.expect(t, slice.equal(out, expect), "width boundary round-trip failed")
    }

    // The two settings really do produce different bytes, and reading one with
    // the other's rule desynchronises rather than silently agreeing.
    early := lzw_test_pack(codes[:], true)
    defer delete(early)
    late := lzw_test_pack(codes[:], false)
    defer delete(late)
    testing.expect(t, !slice.equal(early[:], late[:]), "EarlyChange made no difference to the bit stream")

    wrong, wrong_err := lzw_decode(early[:], false)
    defer delete(wrong)
    testing.expect_value(t, wrong_err, Error.None)
    testing.expect(t, !slice.equal(wrong, expect), "decoding with the wrong EarlyChange should not match")
}

// ---------------------------------------------------------------------------
// Round trips, including a table that fills up
// ---------------------------------------------------------------------------

@(private = "file")
lzw_round_trip :: proc(t: ^testing.T, input: []byte, label: string, loc := #caller_location) {
    for ec in ([?]bool{true, false}) {
        packed := lzw_test_encode(input, ec)
        defer delete(packed)
        out, err := lzw_decode(packed[:], ec)
        defer delete(out)
        testing.expect_value(t, err, Error.None, loc = loc)
        testing.expectf(t, slice.equal(out, input), "%s round trip failed (early_change=%v)", label, ec, loc = loc)
    }
}

@(test)
test_lzw_round_trip_small :: proc(t: ^testing.T) {
    inputs := [?]string{
        "-----A---B",
        "A",
        "AA",
        "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
        "ABABABABABABABABABABABABA",
        "the quick brown fox jumps over the lazy dog",
        "\x00\x00\x00\xff\xff\xff\x00\x00\x00\xff\xff\xff",
    }
    for s in inputs {
        lzw_round_trip(t, transmute([]byte)s, s)
    }
}

@(test)
test_lzw_round_trip_all_byte_values :: proc(t: ^testing.T) {
    input := make([]byte, 256 * 8)
    defer delete(input)
    for i in 0 ..< len(input) {
        input[i] = u8(i & 0xff)
    }
    lzw_round_trip(t, input, "all byte values")
}

// A stream long enough to exhaust all 3838 table entries without ever sending
// a clear: the decoder has to stop growing the table and keep going.
@(test)
test_lzw_table_overflow_without_clear :: proc(t: ^testing.T) {
    input := make([]byte, 40000)
    defer delete(input)
    state: u32 = 0x1234_5678
    for i in 0 ..< len(input) {
        state = state * 1664525 + 1013904223
        input[i] = u8((state >> 19) & 0xff)
    }

    codes := lzw_test_codes(input)
    defer delete(codes)
    // Confirms the premise: many more codes than the table can hold entries.
    testing.expect(t, len(codes) > LZW_MAX_CODE + 500, "test input did not overflow the table")

    lzw_round_trip(t, input, "table overflow")
}

@(test)
test_lzw_long_run_expands :: proc(t: ^testing.T) {
    // Highly compressible input exercises the growth path of the output
    // buffer and the long prefix chains that go with it.
    input := make([]byte, 100_000)
    defer delete(input)
    slice.fill(input, u8('Z'))
    lzw_round_trip(t, input, "long run")
}

// ---------------------------------------------------------------------------
// Nothing may crash, hang or leak on arbitrary bytes
// ---------------------------------------------------------------------------

@(test)
test_lzw_arbitrary_bytes :: proc(t: ^testing.T) {
    buf := make([]byte, 512)
    defer delete(buf)
    state: u32 = 0xdead_beef
    for _ in 0 ..< 200 {
        n := (int(state >> 7) % len(buf)) + 1
        for i in 0 ..< n {
            state = state * 1103515245 + 12345
            buf[i] = u8(state >> 16)
        }
        for ec in ([?]bool{true, false}) {
            out, err := lzw_decode(buf[:n], ec)
            defer delete(out)
            if err != .None {
                testing.expect_value(t, len(out), 0)
            }
        }
    }
}

@(test)
test_lzw_all_zero_bytes :: proc(t: ^testing.T) {
    // Nine zero bits is code 0, so this is a run of NUL bytes plus whatever the
    // growing table makes of it -- a cheap check that a pathological but legal
    // stream terminates.
    zeros := make([]byte, 4096)
    defer delete(zeros)
    out, err := lzw_decode(zeros)
    defer delete(out)
    testing.expect_value(t, err, Error.None)
    testing.expect(t, len(out) > 0, "a stream of zero bits should decode to something")
    for b in out {
        testing.expect_value(t, b, u8(0))
    }
}

@(test)
test_lzw_all_ff_bytes :: proc(t: ^testing.T) {
    ones := make([]byte, 4096)
    defer delete(ones)
    slice.fill(ones, u8(0xff))
    out, err := lzw_decode(ones)
    defer delete(out)
    // 0x1ff is 511, an undefined code straight after the start, so nothing is
    // decodable -- but the call must still return cleanly.
    testing.expect_value(t, err, Error.Bad_Lzw_Data)
    testing.expect_value(t, len(out), 0)
}

// LZW's expansion is not knowable in advance, so lzw_decode sizes its buffer
// from the input length and grows it as it goes. Against a Document arena the
// first guess and every superseded growth buffer would be permanent -- an
// arena's `delete` is a silent no-op -- and a stream that decodes to nothing
// would leave the whole guess behind. apply_filter is where that is prevented,
// by decoding into scratch and copying the result across, so what the arena is
// charged is the output and nothing else.
@(test)
test_lzw_growth_stays_out_of_the_document_arena :: proc(t: ^testing.T) {
    doc := filter_test_doc()
    defer filter_test_doc_free(doc)

    // 0x82 0x00 is the 9-bit code 260, which no table defines at the start of
    // a stream, so not one byte is decoded -- after 300 KiB have been reserved
    // for output that will never be written.
    junk := make([]byte, 100 * 1024)
    defer delete(junk)
    junk[0] = 0x82

    before := document_memory_used(doc)
    out, err := apply_filter(doc, "LZWDecode", junk, nil, doc.allocator)
    testing.expect_value(t, err, Error.Bad_Lzw_Data)
    testing.expect_value(t, len(out), 0)
    testing.expectf(t, document_memory_used(doc) == before,
        "%d bytes stranded in the arena by a stream that decoded to nothing",
        document_memory_used(doc) - before)

    // 16 KiB out of a buffer that starts at a few kilobytes: a handful of
    // regrowths, every one of which would be permanent if it came from here.
    plain := make([]byte, 16 * 1024)
    defer delete(plain)
    for i in 0 ..< len(plain) {
        plain[i] = byte((i / 64) % 251)
    }
    packed := lzw_test_encode(plain, true)
    defer delete(packed)

    before = document_memory_used(doc)
    out2, err2 := apply_filter(doc, "LZWDecode", packed[:], nil, doc.allocator)
    testing.expect_value(t, err2, Error.None)
    testing.expect(t, slice.equal(out2, plain), "the stream did not round trip")
    growth := document_memory_used(doc) - before
    testing.expectf(t, growth <= len(out2) + FT_ARENA_SLACK,
        "the arena grew by %d bytes for %d bytes of output", growth, len(out2))
}