package pdf

// FlateDecode -- RFC 1950 (zlib wrapper) and RFC 1951 (DEFLATE), ISO 32000-2
// 7.4.4.
//
// core:compress/zlib is a conforming decoder, which is exactly why it is not
// used here: it rejects a missing zlib header, a wrong Adler-32, a truncated
// stream and trailing garbage, and all four occur in PDFs that every shipping
// viewer opens. This decoder is deliberately permissive -- it returns whatever
// it managed to inflate and only reports failure when it got nothing at all.
//
// The Huffman decoder is table-driven: a flat 512-entry first-level table
// resolves every code of 9 bits or fewer in one load, and the rare longer
// codes fall back to a canonical bit-at-a-time walk. Every content stream in
// every document goes through here, so the fast path matters.

// A first level of 9 bits covers the whole fixed-literal alphabet and, in
// practice, well over 95% of dynamic codes.
@(private) FLATE_FAST_BITS   :: 9
@(private) FLATE_FAST_SIZE   :: 1 << FLATE_FAST_BITS
@(private) FLATE_FAST_MASK   :: FLATE_FAST_SIZE - 1
@(private) FLATE_MAX_BITS    :: 15
@(private) FLATE_MAX_SYMBOLS :: 288

// DEFLATE cannot expand its input by more than about 1032:1. The margin above
// that absorbs multi-block streams and the zlib wrapper; the floor keeps tiny
// streams (a two-token content stream, an empty stream) workable, and the
// ceiling is the backstop against a decompression bomb.
@(private) FLATE_MAX_RATIO   :: 1200
@(private) FLATE_MIN_OUTPUT  :: 1 << 16
@(private) FLATE_MAX_OUTPUT  :: 1 << 30

// How far into the stream a zlib header is still worth looking for. Producers
// that leak an extra EOL or a stray token before the stream data are common;
// producers that leak a kilobyte are not.
@(private) FLATE_JUNK_SCAN   :: 64

// flate_decode inflates a FlateDecode stream.
//
// The result is allocated with `allocator` and belongs to the caller -- and so
// do the attempts that lose. Damaged data is read several ways below and only
// the best reading survives; each rejected attempt is freed, which needs an
// `allocator` where freeing means something. Give it scratch, never a Document
// arena: an arena's `delete` is a silent no-op, so every attempt would sit
// there until `close`. `apply_filter` in filters.odin is where that is
// arranged, and the comment there is the long version.
//
// A truncated or otherwise damaged stream still returns .None together with
// everything that could be recovered, so callers should compare the length
// against what they expected rather than trusting the error alone; .None with
// a short result is the signal to record a .Bad_Filter_Data repair.
// .Bad_Flate_Data means not one byte was recoverable.
flate_decode :: proc(data: []byte, allocator := context.allocator) -> (out: []byte, err: Error) {
    if len(data) == 0 {
        return nil, .Bad_Flate_Data
    }

    // Leading EOL between "stream" and the data proper is required by the
    // spec but frequently mis-counted by the /Length, so it can survive into
    // the raw bytes.
    start := 0
    for start < len(data) && flate_is_space(data[start]) {
        start += 1
    }
    if start >= len(data) {
        return nil, .Bad_Flate_Data
    }

    // A zlib header is strong evidence: CM, CINFO and the FCHECK modulus all
    // have to agree. When one sits at the front of the data, the bytes behind
    // it are the stream, and guessing at other offsets would only invent
    // output for a file that is merely damaged.
    limit := flate_output_limit(len(data))
    header := -1
    if hlen, ok := flate_header_len(data[start:]); ok {
        header = start + hlen
        got, done := flate_inflate(data[header:], limit, allocator)
        if done || len(got) > 0 {
            return flate_finish(&got, done)
        }
        delete(got)
    }

    // Otherwise there is nothing to do but guess, in descending order of
    // evidence: a header a little way in (a producer that leaked a keyword or
    // an extra EOL before the stream data), raw DEFLATE with no wrapper at
    // all, Acrobat's long-standing bug of dropping the first header byte, and
    // a header that is present but corrupt.
    scan := -1
    if header < 0 {
        scan_end := min(start + FLATE_JUNK_SCAN, len(data) - 1)
        for i in start + 1 ..< scan_end {
            if n, found := flate_header_len(data[i:]); found {
                scan = i + n
                break
            }
        }
    }
    cands := [5]int{scan, start, start + 1, start + 2, 0}
    // If a header did validate and still yielded nothing, only a guess that
    // runs all the way to a final block is worth more than admitting defeat;
    // a partial decode from a guessed offset is noise dressed up as data.
    need_done := header >= 0

    best: [dynamic]byte
    best_done := false
    best_rank := -1

    for off, ci in cands {
        if off < 0 || off >= len(data) || off == header {
            continue
        }
        dup := false
        for j in 0 ..< ci {
            if cands[j] == off {
                dup = true
            }
        }
        if dup {
            continue
        }
        got, done := flate_inflate(data[off:], limit, allocator)
        rank := flate_rank(done, len(got))
        // Earlier candidates are the more likely reading of the bytes, so a
        // later one has to be strictly better to displace them.
        if (!need_done || done) && (rank > best_rank || (rank == best_rank && len(got) > len(best))) {
            if best_rank >= 0 {
                delete(best)
            }
            best = got
            best_done = done
            best_rank = rank
        } else {
            delete(got)
        }
        if best_rank == 3 {
            break
        }
    }

    if best_rank < 0 {
        return nil, .Bad_Flate_Data
    }
    return flate_finish(&best, best_done)
}

// flate_finish hands a completed attempt to the caller, trimming the growth
// slack off the buffer. An empty result is only an error when the stream never
// reached its end -- an empty stream really does inflate to nothing.
@(private)
flate_finish :: proc(buf: ^[dynamic]byte, done: bool) -> ([]byte, Error) {
    if len(buf) == 0 {
        delete(buf^)
        return nil, done ? .None : .Bad_Flate_Data
    }
    shrink(buf)
    return buf[:], .None
}

// flate_rank orders two attempts at the same bytes: a stream that finished and
// produced output beats one that only produced output, which beats one that
// finished empty, which beats nothing at all.
@(private)
flate_rank :: proc(done: bool, n: int) -> int {
    switch {
    case done && n > 0: return 3
    case n > 0:         return 2
    case done:          return 1
    }
    return 0
}

@(private)
flate_is_space :: proc(c: byte) -> bool {
    return c == 0x00 || c == 0x09 || c == 0x0a || c == 0x0c || c == 0x0d || c == 0x20
}

// flate_header_len recognises an RFC 1950 header and reports how many bytes it
// occupies (2, or 6 with a preset dictionary identifier).
@(private)
flate_header_len :: proc(d: []byte) -> (n: int, ok: bool) {
    if len(d) < 2 {
        return 0, false
    }
    cmf, flg := d[0], d[1]
    if cmf & 0x0f != 8 {
        return 0, false // CM must be DEFLATE
    }
    if cmf >> 4 > 7 {
        return 0, false // CINFO above 7 means a window larger than 32K
    }
    if (u16(cmf) << 8 | u16(flg)) % 31 != 0 {
        return 0, false // FCHECK
    }
    if flg & 0x20 != 0 {
        if len(d) < 6 {
            return 0, false
        }
        return 6, true
    }
    return 2, true
}

@(private)
flate_output_limit :: proc(input_size: int) -> int {
    if input_size >= FLATE_MAX_OUTPUT / FLATE_MAX_RATIO {
        return FLATE_MAX_OUTPUT // guard the multiply, not just its result
    }
    return clamp(input_size * FLATE_MAX_RATIO, FLATE_MIN_OUTPUT, FLATE_MAX_OUTPUT)
}

// ---------------------------------------------------------------------------
// bit reader
// ---------------------------------------------------------------------------

// DEFLATE packs bits least-significant-first within each byte, so the buffer
// fills from the low end and is consumed by shifting right.
@(private)
Flate_Bits :: struct {
    data: []byte,
    pos:  int,
    buf:  u64,
    cnt:  u32,
}

@(private)
flate_refill :: proc(br: ^Flate_Bits) {
    for br.cnt <= 56 {
        if br.pos >= len(br.data) {
            break
        }
        br.buf |= u64(br.data[br.pos]) << br.cnt
        br.pos += 1
        br.cnt += 8
    }
}

// flate_get returns n bits, or ok=false when the input ran out. n <= 32.
@(private)
flate_get :: proc(br: ^Flate_Bits, n: u32) -> (v: u32, ok: bool) {
    if n == 0 {
        return 0, true
    }
    if br.cnt < n {
        flate_refill(br)
        if br.cnt < n {
            return 0, false
        }
    }
    v = u32(br.buf & ((u64(1) << n) - 1))
    br.buf >>= n
    br.cnt -= n
    return v, true
}

// flate_align drops the partial byte and pushes whole buffered bytes back, so
// a stored block can be copied straight out of the input.
@(private)
flate_align :: proc(br: ^Flate_Bits) {
    whole := int(br.cnt / 8)
    br.pos -= whole
    br.buf = 0
    br.cnt = 0
}

// ---------------------------------------------------------------------------
// Huffman
// ---------------------------------------------------------------------------

// `fast` holds sym<<4 | len for every code of FLATE_FAST_BITS or fewer, keyed
// by the next 9 stream bits; 0 means "not a short code" and sends the caller
// to the canonical walk over `counts`/`symbols`. Zero can never be a real
// entry because a code length of 0 is not assigned.
@(private)
Flate_Huffman :: struct {
    fast:    [FLATE_FAST_SIZE]u16,
    counts:  [FLATE_MAX_BITS + 1]u16,
    symbols: [FLATE_MAX_SYMBOLS]u16,
}

@(private)
flate_reverse :: proc(code: u32, nbits: u32) -> u32 {
    v := code
    r := u32(0)
    for _ in 0 ..< nbits {
        r = (r << 1) | (v & 1)
        v >>= 1
    }
    return r
}

// flate_build fills a decoder from a code-length vector. An incomplete code is
// accepted -- a single-entry distance tree is legal and shows up in real
// streams -- but an over-subscribed one is not, because there is no sane
// interpretation of it.
@(private)
flate_build :: proc(h: ^Flate_Huffman, lengths: []u8) -> bool {
    h.fast = {}
    h.counts = {}
    if len(lengths) > FLATE_MAX_SYMBOLS {
        return false
    }

    used := 0
    for l in lengths {
        if l > FLATE_MAX_BITS {
            return false
        }
        if l != 0 {
            h.counts[l] += 1
            used += 1
        }
    }
    if used == 0 {
        return true // no codes: any decode against it fails, which is correct
    }

    left := 1
    for l in 1 ..= FLATE_MAX_BITS {
        left <<= 1
        left -= int(h.counts[l])
        if left < 0 {
            return false
        }
    }

    offsets: [FLATE_MAX_BITS + 1]u16
    next:    [FLATE_MAX_BITS + 1]u32
    off := u16(0)
    code := u32(0)
    for l in 1 ..= FLATE_MAX_BITS {
        offsets[l] = off
        off += h.counts[l]
        code = (code + u32(h.counts[l - 1])) << 1
        next[l] = code
    }

    for sym in 0 ..< len(lengths) {
        l := u32(lengths[sym])
        if l == 0 {
            continue
        }
        h.symbols[offsets[l]] = u16(sym)
        offsets[l] += 1
        cv := next[l]
        next[l] += 1
        if l <= FLATE_FAST_BITS {
            entry := u16(sym) << 4 | u16(l)
            step := u32(1) << l
            for i := flate_reverse(cv, l); i < FLATE_FAST_SIZE; i += step {
                h.fast[i] = entry
            }
        }
    }
    return true
}

// flate_slow walks the canonical code one bit at a time. Only codes longer
// than FLATE_FAST_BITS reach it.
@(private)
flate_slow :: proc(br: ^Flate_Bits, h: ^Flate_Huffman) -> (sym: u16, ok: bool) {
    code  := u32(0)
    first := u32(0)
    index := u32(0)
    for l in 1 ..= FLATE_MAX_BITS {
        b := flate_get(br, 1) or_return
        code |= b
        count := u32(h.counts[l])
        if code - first < count {
            at := index + (code - first)
            if at >= FLATE_MAX_SYMBOLS {
                return 0, false
            }
            return h.symbols[at], true
        }
        index += count
        first = (first + count) << 1
        code <<= 1
    }
    return 0, false
}

@(private)
flate_sym :: proc(br: ^Flate_Bits, h: ^Flate_Huffman) -> (sym: u16, ok: bool) {
    if br.cnt < FLATE_FAST_BITS {
        flate_refill(br)
    }
    // Near end of input the peek is zero-padded, but a hit is still sound:
    // the code is a prefix code, so if the first `l` bits form a complete code
    // that is the code, whatever follows.
    e := h.fast[u32(br.buf) & FLATE_FAST_MASK]
    if e != 0 {
        l := u32(e & 15)
        if l > br.cnt {
            return 0, false
        }
        br.buf >>= l
        br.cnt -= l
        return e >> 4, true
    }
    return flate_slow(br, h)
}

// ---------------------------------------------------------------------------
// inflate
// ---------------------------------------------------------------------------

// RFC 1951 3.2.5.
@(private) FLATE_LEN_BASE := [29]u16{
    3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59,
    67, 83, 99, 115, 131, 163, 195, 227, 258,
}
@(private) FLATE_LEN_EXTRA := [29]u8{
    0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3,
    4, 4, 4, 4, 5, 5, 5, 5, 0,
}
@(private) FLATE_DIST_BASE := [30]u16{
    1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193,
    257, 385, 513, 769, 1025, 1537, 2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577,
}
@(private) FLATE_DIST_EXTRA := [30]u8{
    0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6,
    7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13,
}
// RFC 1951 3.2.7: code-length codes are transmitted in this order.
@(private) FLATE_CLEN_ORDER := [19]u8{
    16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15,
}

@(private)
Flate_State :: struct {
    br:   Flate_Bits,
    lit:  Flate_Huffman,
    dist: Flate_Huffman,
    clen: Flate_Huffman,
}

// flate_inflate decodes a raw DEFLATE stream. `done` is true only when a final
// block was reached and decoded; otherwise `out` holds everything recovered
// before the stream went bad, ran out, or hit `limit`.
@(private)
flate_inflate :: proc(src: []byte, limit: int, allocator := context.allocator) -> (out: [dynamic]byte, done: bool) {
    // PDF streams typically expand three- or four-fold; guessing that up front
    // saves most of the regrowth without betting a large allocation on an
    // offset that may turn out to be nonsense.
    initial := min(max(1024, len(src) * 4), 1 << 20)
    initial = min(initial, limit)
    out = make([dynamic]byte, 0, max(initial, 16), allocator)
    if len(src) == 0 {
        return out, false
    }

    // Roughly 5 KiB of tables. On the stack rather than the caller's
    // allocator, because the caller's allocator is usually a Document arena
    // that is never freed until the document closes, and a long document
    // decodes thousands of streams.
    state: Flate_State
    st := &state
    st.br.data = src

    // Every block header consumes at least three bits, so this can only be
    // reached by input that is already exhausted; it is a belt-and-braces
    // bound, not the real terminator.
    max_blocks := len(src) * 8 / 3 + 8
    for _ in 0 ..< max_blocks {
        final := flate_get(&st.br, 1) or_break
        btype := flate_get(&st.br, 2) or_break
        switch btype {
        case 0:
            if !flate_stored(st, &out, limit) {
                return out, false
            }
        case 1:
            if !flate_fixed_tables(st) {
                return out, false
            }
            if !flate_block(st, &out, limit) {
                return out, false
            }
        case 2:
            if !flate_dynamic_tables(st) {
                return out, false
            }
            if !flate_block(st, &out, limit) {
                return out, false
            }
        case:
            return out, false // btype 3 is reserved
        }
        if final == 1 {
            // Anything after the final block is not our business; PDF streams
            // routinely carry padding, an EOL, or the next object's bytes.
            return out, true
        }
    }
    return out, false
}

@(private)
flate_grow :: proc(out: ^[dynamic]byte, extra: int, limit: int) -> bool {
    need := len(out) + extra
    if need > limit {
        return false
    }
    if need <= cap(out) {
        return true
    }
    new_cap := cap(out) * 2
    if new_cap < need {
        new_cap = need
    }
    if new_cap > limit {
        new_cap = limit
    }
    if reserve(out, new_cap) != nil {
        return false
    }
    return true
}

@(private)
flate_stored :: proc(st: ^Flate_State, out: ^[dynamic]byte, limit: int) -> bool {
    br := &st.br
    flate_align(br)
    if br.pos < 0 || br.pos + 4 > len(br.data) {
        return false
    }
    blen  := int(br.data[br.pos])     | int(br.data[br.pos + 1]) << 8
    bnlen := int(br.data[br.pos + 2]) | int(br.data[br.pos + 3]) << 8
    br.pos += 4
    // LEN/~LEN disagreeing means we are not looking at a block header at all,
    // usually because a fallback offset guessed wrong. Trusting LEN anyway
    // would copy megabytes of garbage.
    if blen != (~bnlen & 0xffff) {
        return false
    }
    avail := min(blen, len(br.data) - br.pos)
    if avail > 0 {
        if !flate_grow(out, avail, limit) {
            avail = min(avail, limit - len(out))
            if avail > 0 && !flate_grow(out, avail, limit) {
                avail = 0
            }
        }
        if avail > 0 {
            n := len(out)
            non_zero_resize(out, n + avail)
            copy(out[n:], br.data[br.pos:br.pos + avail])
            br.pos += avail
        }
    }
    return avail == blen
}

@(private)
flate_fixed_tables :: proc(st: ^Flate_State) -> bool {
    // RFC 1951 3.2.6.
    lengths: [FLATE_MAX_SYMBOLS]u8
    for i in 0 ..< 144   { lengths[i] = 8 }
    for i in 144 ..< 256 { lengths[i] = 9 }
    for i in 256 ..< 280 { lengths[i] = 7 }
    for i in 280 ..< 288 { lengths[i] = 8 }
    if !flate_build(&st.lit, lengths[:]) {
        return false
    }
    dlengths: [30]u8
    for i in 0 ..< 30 { dlengths[i] = 5 }
    return flate_build(&st.dist, dlengths[:])
}

@(private)
flate_dynamic_tables :: proc(st: ^Flate_State) -> bool {
    br := &st.br
    hlit  := int(flate_get(br, 5) or_return) + 257
    hdist := int(flate_get(br, 5) or_return) + 1
    hclen := int(flate_get(br, 4) or_return) + 4

    clen: [19]u8
    for i in 0 ..< hclen {
        clen[FLATE_CLEN_ORDER[i]] = u8(flate_get(br, 3) or_return)
    }
    if !flate_build(&st.clen, clen[:]) {
        return false
    }

    total := hlit + hdist // at most 288 + 32
    lengths: [FLATE_MAX_SYMBOLS + 32]u8
    i := 0
    for i < total {
        sym := flate_sym(br, &st.clen) or_return
        switch {
        case sym < 16:
            lengths[i] = u8(sym)
            i += 1
        case sym == 16:
            if i == 0 {
                return false // nothing to repeat
            }
            prev := lengths[i - 1]
            rep := 3 + int(flate_get(br, 2) or_return)
            rep = min(rep, total - i) // a run past the end is corruption; clamp
            for _ in 0 ..< rep {
                lengths[i] = prev
                i += 1
            }
        case sym == 17:
            rep := 3 + int(flate_get(br, 3) or_return)
            i = min(i + rep, total)
        case sym == 18:
            rep := 11 + int(flate_get(br, 7) or_return)
            i = min(i + rep, total)
        case:
            return false
        }
    }

    if !flate_build(&st.lit, lengths[:hlit]) {
        return false
    }
    return flate_build(&st.dist, lengths[hlit:total])
}

@(private)
flate_block :: proc(st: ^Flate_State, out: ^[dynamic]byte, limit: int) -> bool {
    br := &st.br
    // Each iteration consumes at least one bit, so the input already bounds
    // this loop; the counter only makes that guarantee local.
    fuel := len(br.data) * 8 + 64
    for _ in 0 ..< fuel {
        sym := flate_sym(br, &st.lit) or_return
        switch {
        case sym < 256:
            if !flate_grow(out, 1, limit) {
                return false
            }
            append(out, u8(sym))
        case sym == 256:
            return true
        case:
            si := int(sym) - 257
            if si >= 29 {
                return false // 286 and 287 are never assigned a meaning
            }
            length := int(FLATE_LEN_BASE[si]) + int(flate_get(br, u32(FLATE_LEN_EXTRA[si])) or_return)
            dsym := flate_sym(br, &st.dist) or_return
            if int(dsym) >= 30 {
                return false
            }
            dist := int(FLATE_DIST_BASE[dsym]) + int(flate_get(br, u32(FLATE_DIST_EXTRA[dsym])) or_return)
            if dist <= 0 || dist > len(out) {
                return false // reference before the start of the output
            }
            if !flate_copy(out, dist, length, limit) {
                return false
            }
        }
    }
    return false
}

// flate_copy performs the LZ77 back-reference. The copy must be byte at a
// time: runs overlap by design (distance 1, length 258 is a 258-byte run).
@(private)
flate_copy :: proc(out: ^[dynamic]byte, dist: int, length: int, limit: int) -> bool {
    n := len(out)
    take := min(length, limit - n)
    if take <= 0 {
        return false
    }
    if !flate_grow(out, take, limit) {
        return false
    }
    non_zero_resize(out, n + take)
    src := n - dist
    for i in 0 ..< take {
        out[n + i] = out[src + i]
    }
    return take == length
}