package pdf

import "core:slice"

// LZWDecode -- the TIFF flavour of LZW that PDF adopted (ISO 32000-2, 7.4.4).
//
// Codes are packed most-significant-bit first, start at 9 bits wide and grow
// to 12. 256 clears the table, 257 is end-of-data, and the string table is
// filled from 258 upwards. The only wrinkle is /EarlyChange: a TIFF encoder
// adds its table entry immediately after writing a code, so its "next free
// code" runs exactly one ahead of the decoder's. Encoders that were written
// against the TIFF text therefore switch to a wider code one code sooner than
// the textbook algorithm says, and a decoder has to follow suit. /EarlyChange
// 1 is the default and by far the common case, but 0 occurs in real files, so
// both are supported.

@(private) LZW_CLEAR      :: 256
@(private) LZW_EOD        :: 257
@(private) LZW_FIRST_CODE :: 258
@(private) LZW_MAX_CODE   :: 4095      // 12 bits is the widest code
@(private) LZW_TABLE_SIZE :: 4096
@(private) LZW_NO_PREFIX  :: 0xffff    // marks a single-byte root entry

// Bytes a single LZW stream is allowed to expand to. LZW's worst case is a
// factor of a few thousand -- a megabyte of hostile codes can name gigabytes
// of output -- so the ceiling has to be absolute rather than a multiple of the
// input. Matches MAX_OBJSTM_SIZE in document.odin; no honest PDF comes close,
// since LZW is a legacy filter used for small images and old content streams.
@(private) LZW_MAX_OUTPUT :: 1 << 30

// One string in the table, stored as "the string at `prefix`, plus `suffix`".
// `first` is cached because the decoder needs the first byte of a string on
// every single code and walking the prefix chain for it would be quadratic.
@(private)
Lzw_Entry :: struct {
    prefix: u16,
    length: u16,
    suffix: u8,
    first:  u8,
}

@(private)
Lzw_Bits :: struct {
    data: []byte,
    pos:  u64,    // index of the next bit to read
}

// lzw_code_width returns the width the *next* code is read at, given how many
// table entries exist. With early_change the switch happens at 511/1023/2047
// instead of 512/1024/2048.
@(private)
lzw_code_width :: proc(next_code: int, early_change: bool) -> uint {
    n := next_code + (1 if early_change else 0)
    switch {
    case n < 512:  return 9
    case n < 1024: return 10
    case n < 2048: return 11
    }
    return 12
}

// lzw_read_code pulls one big-endian code off the bit stream. It fails rather
// than reading past the end, which is how a stream truncated mid-code stops.
@(private)
lzw_read_code :: proc(b: ^Lzw_Bits, width: uint) -> (code: u16, ok: bool) {
    if width < 1 || width > 12 {
        return 0, false
    }
    total := u64(len(b.data)) * 8
    if b.pos + u64(width) > total {
        return 0, false
    }

    // A code spans at most 3 bytes (7 bits of offset + 12 of width = 19).
    i := int(b.pos >> 3)
    off := uint(b.pos & 7)
    acc := u32(b.data[i]) << 16
    if i + 1 < len(b.data) {
        acc |= u32(b.data[i + 1]) << 8
    }
    if i + 2 < len(b.data) {
        acc |= u32(b.data[i + 2])
    }
    shift := 24 - off - width
    b.pos += u64(width)
    return u16((acc >> shift) & ((u32(1) << width) - 1)), true
}

// lzw_emit appends the string named by `code`. The chain is walked backwards
// and the result reversed in place, which avoids a scratch buffer. It either
// appends the whole string or nothing at all, so a refusal never leaves half a
// string in the output.
@(private)
lzw_emit :: proc(buf: ^[dynamic]byte, table: ^[LZW_TABLE_SIZE]Lzw_Entry, code: u16) -> bool {
    if int(code) >= LZW_TABLE_SIZE {
        return false
    }
    n := int(table[code].length)
    if n <= 0 || n > LZW_TABLE_SIZE {
        return false
    }
    if len(buf^) > LZW_MAX_OUTPUT - n {
        return false
    }

    start := len(buf^)
    c := code
    // Prefixes are strictly decreasing by construction, so this terminates
    // after at most LZW_TABLE_SIZE steps even if the table were corrupt.
    for _ in 0 ..< LZW_TABLE_SIZE {
        append(buf, table[c].suffix)
        p := table[c].prefix
        if p == LZW_NO_PREFIX || int(p) >= LZW_TABLE_SIZE {
            break
        }
        c = p
    }
    if len(buf^) - start != n {
        resize(buf, start)
        return false
    }
    slice.reverse(buf^[start:])
    return true
}

// lzw_decode decodes an LZWDecode stream. `early_change` mirrors the
// /EarlyChange decode parameter (default 1, i.e. true).
//
// The returned slice is allocated with `allocator` and belongs to the caller,
// and so does the buffer it grew out of: LZW's expansion is not knowable up
// front, so the buffer is sized by guess and regrown, and a stream that
// decodes to nothing frees it again. Give it scratch, not a Document arena,
// whose `delete` is a silent no-op that would keep the guess and every
// regrowth until `close`; `apply_filter` in filters.odin is where that is
// arranged.
//
// Damaged input degrades instead of failing: a stream with no EOD marker, a
// stream truncated in the middle of a code, a table that overflows without an
// intervening clear, and a code that was never defined all stop the decode and
// hand back everything produced up to that point. `err` is therefore non-.None
// only when nothing at all could be decoded from a non-empty stream -- callers
// that get bytes back should use them.
lzw_decode :: proc(data: []byte, early_change := true, allocator := context.allocator) -> (out: []byte, err: Error) {
    if len(data) == 0 {
        return nil, .None
    }

    table: [LZW_TABLE_SIZE]Lzw_Entry
    for i in 0 ..< 256 {
        table[i] = {prefix = LZW_NO_PREFIX, length = 1, suffix = u8(i), first = u8(i)}
    }

    // A guess, not a bound: LZW on real data lands near 3x. The buffer grows
    // as needed and LZW_MAX_OUTPUT is what actually limits it.
    guess := clamp(len(data) * 3, 64, 1 << 20)
    buf, alloc_err := make([dynamic]byte, 0, guess, allocator)
    if alloc_err != nil {
        return nil, .Out_Of_Memory
    }

    bits := Lzw_Bits{data = data}
    next_code := LZW_FIRST_CODE
    prev := -1
    saw_eod := false

    // Every pass either consumes at least 9 bits or leaves the loop, and the
    // bit position is bounded by len(data)*8, so this cannot spin.
    decode: for {
        width := lzw_code_width(next_code, early_change)
        raw, ok := lzw_read_code(&bits, width)
        if !ok {
            break decode
        }
        code := int(raw)

        switch {
        case code == LZW_CLEAR:
            next_code = LZW_FIRST_CODE
            prev = -1
            continue decode
        case code == LZW_EOD:
            saw_eod = true
            break decode
        }

        start := len(buf)
        switch {
        case code < 256:
            if len(buf) >= LZW_MAX_OUTPUT {
                break decode
            }
            append(&buf, u8(code))

        case code >= LZW_FIRST_CODE && code < next_code:
            if !lzw_emit(&buf, &table, u16(code)) {
                break decode
            }

        case code == next_code && prev >= 0 && prev < LZW_TABLE_SIZE && next_code <= LZW_MAX_CODE:
            // KwKwK: the encoder used the entry it is defining with this very
            // code. The string is the previous one with its own first byte
            // appended. ISO 32000-2, 7.4.4.2.
            if !lzw_emit(&buf, &table, u16(prev)) {
                break decode
            }
            if len(buf) >= LZW_MAX_OUTPUT {
                break decode
            }
            append(&buf, table[prev].first)

        case:
            // A code that was never defined and is not the legal KwKwK case.
            // There is no way to resynchronise, so keep what we have.
            break decode
        }
        if len(buf) == start {
            break decode
        }

        // The new entry is the previous string extended by the first byte of
        // the one just emitted. Once the table is full we stop growing it and
        // carry on decoding with what is there: a stream that overflows
        // without sending a clear is broken, but usually only mildly.
        if prev >= 0 && prev < LZW_TABLE_SIZE && next_code <= LZW_MAX_CODE {
            table[next_code] = {
                prefix = u16(prev),
                length = u16(min(int(table[prev].length) + 1, LZW_TABLE_SIZE)),
                suffix = buf[start],
                first  = table[prev].first,
            }
            next_code += 1
        }
        prev = code
    }

    if len(buf) == 0 {
        delete(buf)
        // An empty stream that reached its EOD really did encode nothing.
        return nil, saw_eod ? .None : .Bad_Lzw_Data
    }
    _, _ = shrink(&buf)
    return buf[:], .None
}