package pdf

import "core:mem"
import "core:mem/virtual"
import "core:os"
import "core:slice"
import "core:strings"

// ---------------------------------------------------------------------------
// Limits
//
// Every one of these bounds something an attacker (or a fuzzer, or a merely
// broken producer) controls. They are deliberately generous -- real documents
// stay far below them -- and exist only so that a malformed file fails instead
// of allocating the machine to death or spinning forever.
// ---------------------------------------------------------------------------

// Nested arrays and dictionaries, and form-XObject recursion.
MAX_DEPTH             :: 32
// Object numbers above this are rejected outright (ISO 32000-2, 7.5.4 caps
// them at 8388607 for a reason: the xref table is indexed by them).
MAX_OBJECT_NUMBER     :: 8388607
// /Prev hops through the cross-reference chain.
MAX_XREF_CHAIN        :: 64
// Operands stacked before a content operator; the format never needs more
// than 8 apart from `d0`/`d1` and the shading operators.
MAX_CONTENT_OPERANDS  :: 32
// 268M pixels, about 16k x 16k. Above this `render_page` refuses.
MAX_BITMAP_PIXELS     :: 1 << 28
// Ref chains: "1 0 R" pointing at "2 0 R" pointing at ... Legal, but only
// ever seen in files trying to make a reader loop.
MAX_REF_CHAIN         :: 32
// How deeply one object's definition may require another's before the library
// stops following. This is *nesting*, not chaining: it bounds only the cases
// where parsing object A cannot finish until object B is parsed -- an indirect
// /Length, an object stream's container, an indirect /Filter. Legitimate files
// nest two or three deep. Without this bound a file can name an acyclic chain
// of 100000 such objects and overflow the native stack, which the cycle
// guards cannot catch because no object number ever repeats.
MAX_LOAD_DEPTH        :: 64
// Bytes an object stream is allowed to expand to.
MAX_OBJSTM_SIZE       :: 1 << 30
// Repairs recorded before the library stops bothering; a file this broken is
// not going to be diagnosed by reading more of the list.
MAX_REPAIRS           :: 4096
// Pages one document may contain. The largest real documents run to a few
// tens of thousands; a /Kids graph that yields more than this is a file
// trying to make the page-tree walk allocate without end.
MAX_PAGES             :: 1 << 20

// ---------------------------------------------------------------------------
// Cross-reference table
//
// Declared here rather than in xref.odin because the table is the Document's
// own state; xref.odin holds the algorithms that fill it in.
// ---------------------------------------------------------------------------

Xref_Kind :: enum u8 {
    // Never defined, or on the free list. Reads as null.
    Free = 0,
    // Lives at `offset` in the file as "num gen obj ... endobj".
    Normal,
    // Lives inside object stream `container` at index `index`.
    Compressed,
}

Xref_Entry :: struct {
    offset:    i64,    // Normal: byte offset in the file
    container: u32,    // Compressed: object number of the /ObjStm
    index:     u32,    // Compressed: position within that stream
    gen:       u16,
    kind:      Xref_Kind,
    // Set once the entry has been read and its object cached, so a second
    // lookup skips straight to the cache.
    cached:    bool,
}

// ---------------------------------------------------------------------------
// Document
// ---------------------------------------------------------------------------

Open_Options :: struct {
    // User or owner password for an encrypted document. Empty is correct for
    // the common case of a file encrypted with an empty user password -- which
    // is what "restricted permissions, opens without prompting" means.
    password:      string,
    // When true the Document takes ownership of the byte slice handed to
    // `open_bytes` and frees it in `close`. `open_file` sets this.
    take_ownership: bool,
    // Skip the page-tree walk. Useful for tools that only want to inspect
    // objects, and for opening documents whose page tree is broken.
    skip_page_tree: bool,
    // Allocator for the Document handle itself, and for the file bytes when
    // `open_file` reads them.
    //
    // It is deliberately NOT where parsed objects come from: those live in a
    // growing virtual-memory arena that takes pages from the OS directly, so
    // that a large document costs one address-space reservation rather than
    // thousands of small allocations. A caller that needs to account for a
    // document's real footprint should read `document_memory_used`, not
    // instrument this allocator.
    allocator:     mem.Allocator,
}

Document :: struct {
    // The whole file. Object offsets index into this directly, and most
    // strings and stream payloads are sub-slices of it rather than copies.
    data:           []byte,
    owns_data:      bool,
    backing:        mem.Allocator,

    // Everything parsed out of the file lives here and dies with `close`.
    arena:          virtual.Arena,
    allocator:      mem.Allocator,

    version_major:  int,
    version_minor:  int,
    // Byte position of "%PDF-". Files with junk prepended (a mail header, an
    // HTTP response, a self-extracting wrapper) are common, and every offset
    // inside such a file is relative to this, not to zero.
    header_offset:  i64,

    xref:           [dynamic]Xref_Entry,
    trailer:        ^Dict,
    catalog:        ^Dict,

    // Resolved objects, keyed by object number. Generation is not part of the
    // key: a file never has two live generations of the same number, and the
    // ones that try are broken in ways the generation check makes worse.
    cache:          map[u32]Object,
    // Objects currently being parsed, to break reference cycles that would
    // otherwise recurse until the stack ends.
    loading:        map[u32]bool,
    // Decoded object streams, keyed by their object number.
    objstm_cache:   map[u32]^Objstm,

    // Colour spaces and functions built from this file, keyed by the address
    // of the object they were built from. A page that names the same space in
    // four hundred XObjects, or a shading that names one tint transform from
    // every pattern that uses it, must build each of them once. Both are
    // filled in lazily and, like everything else the Document owns, die with
    // its arena.
    cs_cache:       map[uintptr]^Colorspace,
    fn_cache:       map[uintptr]^Function,
    // Fonts, keyed the same way, by the address of the font dictionary. A
    // page that selects the same font two hundred times must build it once,
    // and building one means parsing an embedded font program.
    font_cache:     map[uintptr]^Font,

    // Images read from this file, and the decoded pixels of the ones most
    // recently painted. Unlike the two caches above this one has a byte
    // budget and evicts, because a decoded photograph is megabytes and a
    // catalogue has hundreds of them; see image.odin.
    img_cache:      Image_Cache,

    encrypt:        ^Encryption,

    // Page dictionaries in document order, flattened from the page tree.
    pages:          [dynamic]^Dict,
    page_refs:      [dynamic]Ref,

    repairs:        [dynamic]Repair,
    // True once anything was worked around. Cheaper to test than len(repairs).
    is_damaged:     bool,
}

// ---------------------------------------------------------------------------
// Opening and closing
// ---------------------------------------------------------------------------

// open_file reads a PDF from disk. The Document owns the bytes.
open_file :: proc(path: string, options := Open_Options{}) -> (doc: ^Document, err: Error) {
    opts := options
    if opts.allocator.procedure == nil {
        opts.allocator = context.allocator
    }
    data, read_err := os.read_entire_file(path, opts.allocator)
    if read_err != nil {
        return nil, .Cant_Open_File
    }
    opts.take_ownership = true
    return open_bytes(data, opts)
}

// open_bytes reads a PDF already in memory.
//
// Unless `options.take_ownership` is set the slice is borrowed, not copied, and
// must outlive the Document -- object offsets, string values and stream
// payloads all point into it.
open_bytes :: proc(data: []byte, options := Open_Options{}) -> (doc: ^Document, err: Error) {
    opts := options
    if opts.allocator.procedure == nil {
        opts.allocator = context.allocator
    }
    if len(data) == 0 {
        return nil, .Not_A_PDF
    }

    doc = new(Document, opts.allocator)
    doc.data = data
    doc.owns_data = opts.take_ownership
    doc.backing = opts.allocator

    if arena_err := virtual.arena_init_growing(&doc.arena); arena_err != nil {
        // `close` cannot run yet -- there is no arena for it to destroy -- so
        // the bytes we took ownership of have to be released here.
        if doc.owns_data {
            delete(doc.data, opts.allocator)
        }
        free(doc, opts.allocator)
        return nil, .Out_Of_Memory
    }
    doc.allocator = virtual.arena_allocator(&doc.arena)

    doc.xref         = make([dynamic]Xref_Entry, doc.allocator)
    doc.cache        = make(map[u32]Object, 64, doc.allocator)
    doc.loading      = make(map[u32]bool, 16, doc.allocator)
    doc.objstm_cache = make(map[u32]^Objstm, 8, doc.allocator)
    doc.pages        = make([dynamic]^Dict, doc.allocator)
    doc.page_refs    = make([dynamic]Ref, doc.allocator)
    doc.repairs      = make([dynamic]Repair, doc.allocator)

    // A missing header is not fatal on its own: enough files in the wild have
    // one that a reader which insists on it is less useful than one that goes
    // looking for objects.
    read_header(doc)

    if xerr := xref_load(doc); xerr != .None {
        close(doc)
        return nil, xerr
    }

    if cerr := crypt_init(doc, transmute([]byte)opts.password); cerr != .None {
        close(doc)
        return nil, cerr
    }

    if derr := find_catalog(doc); derr != .None {
        close(doc)
        return nil, derr
    }

    if !opts.skip_page_tree {
        if perr := pagetree_build(doc); perr != .None {
            close(doc)
            return nil, perr
        }
    }

    return doc, .None
}

// close releases everything the Document owns. The Document pointer, every
// object reachable from it, and (if it was taken) the file's byte slice are all
// invalid afterwards.
close :: proc(doc: ^Document) {
    if doc == nil { return }
    crypt_destroy(doc)
    // Decoded image pixels are the one thing the arena does not own: they have
    // to be individually freeable for the cache's budget to mean anything.
    image_cache_release(doc)
    virtual.arena_destroy(&doc.arena)
    if doc.owns_data && doc.data != nil {
        delete(doc.data, doc.backing)
    }
    backing := doc.backing
    free(doc, backing)
}

// read_header locates "%PDF-M.N" and records the version and the offset
// everything else in the file is measured from.
@(private)
read_header :: proc(doc: ^Document) {
    // The header is required to be at byte 0, but junk in front of it is
    // ordinary enough that viewers all scan a window for it. 1 KiB matches
    // what the format's own conformance guidance allows for.
    limit := min(len(doc.data), 4096)
    idx := strings.index(string(doc.data[:limit]), "%PDF-")
    if idx < 0 {
        // Last resort: some generators emit "%!PS-Adobe" wrappers, and some
        // emit nothing at all. Leave the offset at zero and let the xref
        // rebuild find objects by scanning.
        doc.version_major = 1
        doc.version_minor = 7
        return
    }

    doc.header_offset = i64(idx)
    v := doc.data[idx + 5:]
    doc.version_major = 1
    doc.version_minor = 7
    if len(v) >= 3 && v[0] >= '0' && v[0] <= '9' && v[1] == '.' && v[2] >= '0' && v[2] <= '9' {
        doc.version_major = int(v[0] - '0')
        doc.version_minor = int(v[2] - '0')
    }
}

// document_version returns the version from the header, upgraded by the
// catalog's /Version entry when that is newer (ISO 32000-2, 7.5.5: a file may
// be upgraded by an incremental update without rewriting the header).
document_version :: proc(doc: ^Document) -> (major, minor: int) {
    major, minor = doc.version_major, doc.version_minor
    if v := dict_name(doc, doc.catalog, "Version"); len(v) >= 3 && v[1] == '.' &&
       v[0] >= '0' && v[0] <= '9' && v[2] >= '0' && v[2] <= '9' {
        cm := int(v[0] - '0')
        cn := int(v[2] - '0')
        if cm > major || (cm == major && cn > minor) {
            major, minor = cm, cn
        }
    }
    return
}

// ---------------------------------------------------------------------------
// Object access
// ---------------------------------------------------------------------------

// resolve follows indirect references until it reaches a direct object.
//
// This is the single most-called procedure in the library: nearly every value
// read out of a dictionary is allowed to be a reference, and callers that skip
// it work on most files and mysteriously fail on the rest. A nil Document
// makes it a no-op, which is what the standalone object parser wants.
resolve :: proc(doc: ^Document, o: Object) -> Object {
    if doc == nil { return o }
    obj := o
    for i in 0 ..< MAX_REF_CHAIN {
        ref, is_ref := obj.(Ref)
        if !is_ref { return obj }
        obj = get_object(doc, ref)
        _ = i
    }
    // A chain this long is a loop. Null is what a viewer shows.
    return nil
}

// get_object loads one indirect object, from the cache if it is there.
//
// A reference to an object that does not exist is null, not an error: the
// format says so (ISO 32000-2, 7.3.10) and files rely on it.
get_object :: proc(doc: ^Document, ref: Ref) -> Object {
    if doc == nil || ref.num == 0 || ref.num > MAX_OBJECT_NUMBER {
        return nil
    }
    if cached, ok := doc.cache[ref.num]; ok {
        return cached
    }
    // Already being parsed further up the stack: an object whose own
    // definition refers back to it. Break the cycle with null.
    if doc.loading[ref.num] or_else false {
        add_repair(doc, .Missing_Object, 0, ref)
        return nil
    }
    // `loading` holds exactly the objects whose parses are in flight, so its
    // size is the current nesting depth. Bounding it stops an acyclic chain --
    // object 2's /Length pointing at object 3's, and so on for thousands of
    // distinct objects -- from walking the native stack off its end. The
    // cycle guard above cannot see such a chain, because no object number
    // ever repeats.
    if len(doc.loading) >= MAX_LOAD_DEPTH {
        add_repair(doc, .Recursion_Limit_Hit, 0, ref)
        return nil
    }
    if int(ref.num) >= len(doc.xref) {
        return nil
    }

    entry := doc.xref[ref.num]
    doc.loading[ref.num] = true
    defer delete_key(&doc.loading, ref.num)

    obj: Object
    switch entry.kind {
    case .Free:
        obj = nil
    case .Normal:
        parsed, got, err := parse_indirect_object_at(doc, entry.offset, Ref{ref.num, entry.gen})
        if err == .None {
            // An xref pointing at the wrong object is a classic symptom of a
            // file that was edited without its table being updated. Take what
            // is actually there, but only if the number matches -- taking a
            // different object's body would be worse than nothing.
            if got.num == ref.num {
                obj = parsed
            } else {
                add_repair(doc, .Xref_Entry_Mismatch, entry.offset, ref)
                obj = nil
            }
        } else {
            add_repair(doc, .Missing_Object, entry.offset, ref)
        }
    case .Compressed:
        parsed, err := objstm_get_object(doc, entry.container, entry.index, ref.num)
        if err == .None {
            obj = parsed
        } else {
            add_repair(doc, .Missing_Object, 0, ref)
        }
    }

    doc.cache[ref.num] = obj
    return obj
}

// get_dict, get_array, get_stream are the typed forms of get_object, for the
// many call sites that know what they are looking for.
get_dict :: proc(doc: ^Document, ref: Ref) -> ^Dict {
    d, _ := to_dict(get_object(doc, ref))
    return d
}

get_array :: proc(doc: ^Document, ref: Ref) -> ^Array {
    a, _ := to_array(get_object(doc, ref))
    return a
}

get_stream :: proc(doc: ^Document, ref: Ref) -> ^Stream {
    s, _ := to_stream(get_object(doc, ref))
    return s
}

// document_memory_used is how many bytes the Document's arena has handed out:
// every object, dictionary, stream payload and decoded buffer parsed so far.
// It grows as pages are visited, since streams decode lazily, and drops to
// zero only at `close`.
document_memory_used :: proc(doc: ^Document) -> int {
    if doc == nil { return 0 }
    return int(doc.arena.total_used)
}

// document_memory_reserved is the address space the arena holds. On a 64-bit
// system this is reserved, not committed, so it is much larger than the
// resident cost and is not the number to budget against.
document_memory_reserved :: proc(doc: ^Document) -> int {
    if doc == nil { return 0 }
    return int(doc.arena.total_reserved)
}

// object_count is the highest object number the file defines, plus one.
object_count :: proc(doc: ^Document) -> int {
    return len(doc.xref)
}

// ---------------------------------------------------------------------------
// Catalog
// ---------------------------------------------------------------------------

@(private)
find_catalog :: proc(doc: ^Document) -> Error {
    if root := dict_dict(doc, doc.trailer, "Root"); root != nil {
        // Trust it only if it looks like a catalog. Some files have a /Root
        // pointing at a page or at a free object after a bad incremental
        // update, and scanning finds the real one.
        if dict_type(doc, root) == "Catalog" || dict_has(root, "Pages") {
            doc.catalog = root
            return .None
        }
    }

    // Scan every object for /Type /Catalog. Costly, but only reached on files
    // whose trailer is unusable, and the alternative is refusing to open them.
    for num in 1 ..< len(doc.xref) {
        obj := get_object(doc, Ref{u32(num), doc.xref[num].gen})
        d, ok := to_dict(obj)
        if !ok { continue }
        if dict_type(doc, d) == "Catalog" {
            doc.catalog = d
            add_repair(doc, .Missing_Root, 0, Ref{u32(num), 0})
            return .None
        }
    }

    // A file with no catalog can still have pages; pagetree_build falls back to
    // scanning for /Type /Page. Report the damage and let it try.
    add_repair(doc, .Missing_Root, 0, {})
    return .None
}

// ---------------------------------------------------------------------------
// Repairs
// ---------------------------------------------------------------------------

add_repair :: proc(doc: ^Document, kind: Repair_Kind, offset: i64, ref := Ref{}) {
    if doc == nil { return }
    doc.is_damaged = true
    if len(doc.repairs) >= MAX_REPAIRS { return }
    append(&doc.repairs, Repair{kind = kind, offset = offset, ref = ref})
}

// ---------------------------------------------------------------------------
// Byte-slice helpers used across the reader
// ---------------------------------------------------------------------------

// slice_at returns doc.data[offset:] with the offset clamped, so a garbage
// offset out of a broken xref yields an empty slice instead of a panic.
@(private)
slice_at :: proc(doc: ^Document, offset: i64) -> []byte {
    if offset < 0 || offset >= i64(len(doc.data)) {
        return nil
    }
    return doc.data[offset:]
}

// arena_bytes copies `src` into the Document's arena. Used when a value cannot
// borrow from `doc.data` -- decrypted strings, decoded streams, anything
// synthesised.
@(private)
arena_bytes :: proc(doc: ^Document, src: []byte) -> []byte {
    if len(src) == 0 { return nil }
    dst := make([]byte, len(src), doc.allocator)
    copy(dst, src)
    return dst
}

@(private)
arena_string :: proc(doc: ^Document, src: string) -> string {
    return string(arena_bytes(doc, transmute([]byte)src))
}

// index_bytes finds `needle` in `haystack`, or -1.
@(private)
index_bytes :: proc(haystack, needle: []byte) -> int {
    return slice_indexof_bytes(haystack, needle)
}

@(private)
slice_indexof_bytes :: proc(haystack, needle: []byte) -> int {
    if len(needle) == 0 { return 0 }
    if len(needle) > len(haystack) { return -1 }
    first := needle[0]
    for i := 0; i + len(needle) <= len(haystack); i += 1 {
        if haystack[i] != first { continue }
        if slice.equal(haystack[i:i + len(needle)], needle) {
            return i
        }
    }
    return -1
}

// last_index_bytes finds the last occurrence of `needle`, or -1. Used to walk
// backwards to `startxref`.
@(private)
last_index_bytes :: proc(haystack, needle: []byte) -> int {
    if len(needle) == 0 { return len(haystack) }
    if len(needle) > len(haystack) { return -1 }
    for i := len(haystack) - len(needle); i >= 0; i -= 1 {
        if slice.equal(haystack[i:i + len(needle)], needle) {
            return i
        }
    }
    return -1
}