package pdf

import "core:os"
import "core:strings"
import "core:testing"
import "core:time"

// ---------------------------------------------------------------------------
// Golden files
//
// The rest of the suite builds its PDFs in memory with the pb_* helpers, which
// is the right way to test a structure: the builder records the offsets, so a
// test never asserts a hand-counted number. But a builder is part of the same
// codebase as the reader, and three things it therefore cannot show:
//
//   1. that `open_file` works on an actual file, not just on a byte slice;
//   2. that byte offsets some *other* tool wrote are the ones we parse;
//   3. that a refactor did not quietly change what the library accepts by
//      changing the builder at the same time.
//
// So testdata/ holds fourteen committed files, one per structural case, and
// this table pins down exactly what each of them must produce -- including the
// exact set of repairs, so that a file which starts opening "more cleanly"
// than it should fails here rather than passing silently.
//
// testdata/README.md says how each file was made.
// ---------------------------------------------------------------------------

@(private = "file")
GOLDEN_DIR :: #directory + "testdata/"

@(private = "file")
Golden :: struct {
    name:   string,
    // Flattened page count.
    pages:  int,
    // First page's /MediaBox, after inheritance.
    width:  f64,
    height: f64,
    // First page's /Rotate, normalised.
    rotate: int,
    // What the first page's decoded content stream must begin with. Every
    // golden page draws the same filled rectangle and stroked diagonal behind
    // a comment naming the file, so this doubles as "the stream decoded".
    marker: string,
    // The repair kinds this file must record -- exactly, no more and no less,
    // once every page's content has been read.
    repairs: bit_set[Repair_Kind],
    // encryption_method_name, or "" for a file that is not encrypted.
    method: string,
    // The /Info /Title, which in the encrypted files is a string that only
    // decrypts correctly when the file key is right.
    title:  string,
}

@(private = "file")
GOLDEN := []Golden{
    // -- structure --------------------------------------------------------
    {"classic-xref.pdf",        1, 200, 200,   0, "% classic-xref",
     {}, "", ""},
    {"xref-stream.pdf",         1, 200, 200,   0, "% xref-stream",
     {}, "", ""},
    {"object-stream.pdf",       1, 200, 200,   0, "% object-stream",
     {}, "", ""},
    // The second generation replaces the page, so 300x300 -- not the 200x200
    // the first generation wrote -- is what proves the update was applied.
    {"incremental-update.pdf",  1, 300, 300,   0, "% incremental-update generation 2",
     {}, "", ""},
    {"junk-before-header.pdf",  1, 200, 200,   0, "% junk-before-header",
     {}, "", ""},

    // -- damage that must be repaired, not rejected ------------------------
    {"broken-xref.pdf",         1, 200, 200,   0, "% broken-xref",
     {.Bad_Startxref, .Rebuilt_Xref}, "", ""},
    {"bad-length.pdf",          1, 200, 200,   0, "% bad-length",
     {.Bad_Stream_Length}, "", ""},
    {"no-xref.pdf",             1, 200, 200,   0, "% no-xref",
     {.Bad_Startxref, .Rebuilt_Xref, .Missing_Root}, "", ""},
    {"page-tree-broken.pdf",    1, 200, 200,   0, "% page-tree-broken",
     {.Cyclic_Page_Tree}, "", ""},

    // -- encryption, every one with an empty user password -----------------
    {"encrypted-rc4-40.pdf",    1, 200, 200,   0, "% encrypted-rc4-40",
     {}, "RC4 40-bit", "Golden File"},
    {"encrypted-rc4-128.pdf",   1, 200, 200,   0, "% encrypted-rc4-128",
     {}, "RC4 128-bit", "Golden File"},
    {"encrypted-aes128.pdf",    1, 200, 200,   0, "% encrypted-aes128",
     {}, "AES-128", "Golden File"},
    {"encrypted-aes256.pdf",    1, 200, 200,   0, "% encrypted-aes256",
     {}, "AES-256", "Golden File"},

    // -- page tree ---------------------------------------------------------
    // /MediaBox and /Rotate come from the root, three levels up.
    {"page-tree-deep.pdf",      2, 300, 400,  90, "% page-tree-deep page 1",
     {}, "", ""},
}

// The password every encrypted golden file restricts itself with. The user
// password is empty, which is what "opens without prompting" means.
@(private = "file")
GOLDEN_OWNER_PW :: "owner"

@(private = "file")
golden_path :: proc(name: string) -> string {
    return strings.concatenate({GOLDEN_DIR, name}, context.temp_allocator)
}

// golden_repairs collects the kinds recorded so far, ignoring how many times
// each occurred: a repair that fires twice is the same fact as one that fires
// once, and the count is an implementation detail this table should not pin.
@(private = "file")
golden_repairs :: proc(doc: ^Document) -> (out: bit_set[Repair_Kind]) {
    if doc == nil { return }
    for r in doc.repairs {
        out += {r.kind}
    }
    return
}

// golden_content decodes one page's content stream.
@(private = "file")
golden_content :: proc(doc: ^Document, page: Page) -> (string, Error) {
    s := dict_stream(doc, page.dict, "Contents")
    data, err := stream_data(doc, s)
    return string(data), err
}

// ---------------------------------------------------------------------------
// The table
// ---------------------------------------------------------------------------

@(test)
test_golden_files :: proc(t: ^testing.T) {
    defer free_all(context.temp_allocator)

    for g in GOLDEN {
        path := golden_path(g.name)

        doc, err := open_file(path)
        defer close(doc)
        if !testing.expectf(t, err == .None, "%s: open_file failed: %s", g.name, error_string(err)) {
            continue
        }
        if !testing.expectf(t, doc != nil, "%s: no document", g.name) {
            continue
        }

        testing.expectf(t, page_count(doc) == g.pages,
            "%s: %d pages, wanted %d", g.name, page_count(doc), g.pages)

        // Every page's content must decode, not just the first: a file whose
        // second page is unreadable is not a file this suite should pass.
        for i in 0 ..< page_count(doc) {
            page, perr := page_get(doc, i)
            if !testing.expectf(t, perr == .None, "%s: page %d: %s", g.name, i, error_string(perr)) {
                continue
            }
            body, cerr := golden_content(doc, page)
            testing.expectf(t, cerr == .None,
                "%s: page %d content: %s", g.name, i, error_string(cerr))
            testing.expectf(t, len(body) > 0, "%s: page %d content is empty", g.name, i)
            testing.expectf(t, strings.contains(body, "160 100 re"),
                "%s: page %d content does not draw the rectangle", g.name, i)

            if i != 0 { continue }
            testing.expectf(t, rect_width(page.media_box) == g.width &&
                               rect_height(page.media_box) == g.height,
                "%s: page 0 is %vx%v, wanted %vx%v", g.name,
                rect_width(page.media_box), rect_height(page.media_box), g.width, g.height)
            testing.expectf(t, page.rotate == g.rotate,
                "%s: page 0 /Rotate %d, wanted %d", g.name, page.rotate, g.rotate)
            testing.expectf(t, strings.has_prefix(body, g.marker),
                "%s: page 0 content starts %q, wanted %q",
                g.name, body[:min(len(body), len(g.marker))], g.marker)
        }

        // Checked after the content streams have been read, because a wrong
        // /Length is only discovered when the stream itself is parsed.
        got := golden_repairs(doc)
        testing.expectf(t, got == g.repairs,
            "%s: repairs %v, wanted %v", g.name, got, g.repairs)
        testing.expectf(t, doc.is_damaged == (g.repairs != {}),
            "%s: is_damaged %v with repairs %v", g.name, doc.is_damaged, got)

        method := document_is_encrypted(doc) ? encryption_method_name(doc) : ""
        testing.expectf(t, method == g.method,
            "%s: encryption %q, wanted %q", g.name, method, g.method)

        info := dict_dict(doc, doc.trailer, "Info")
        title := string(dict_bytes(doc, info, "Title"))
        testing.expectf(t, title == g.title,
            "%s: /Info /Title %q, wanted %q", g.name, title, g.title)
    }
}

// open_bytes and open_file must agree. They share everything below the read,
// but only this says so.
@(test)
test_golden_open_bytes_agrees :: proc(t: ^testing.T) {
    defer free_all(context.temp_allocator)

    for g in GOLDEN {
        data, read_err := os.read_entire_file(golden_path(g.name), context.allocator)
        if !testing.expectf(t, read_err == nil, "%s: cannot read: %v", g.name, read_err) { continue }
        defer delete(data)

        doc, err := open_bytes(data)
        defer close(doc)
        if !testing.expectf(t, err == .None, "%s: open_bytes failed: %s", g.name, error_string(err)) {
            continue
        }
        testing.expectf(t, page_count(doc) == g.pages,
            "%s: open_bytes gave %d pages, wanted %d", g.name, page_count(doc), g.pages)

        page, perr := page_get(doc, 0)
        testing.expectf(t, perr == .None, "%s: page 0: %s", g.name, error_string(perr))
        testing.expectf(t, rect_width(page.media_box) == g.width,
            "%s: open_bytes gave width %v, wanted %v", g.name, rect_width(page.media_box), g.width)
    }
}

// ---------------------------------------------------------------------------
// Cases the table cannot express in one line
// ---------------------------------------------------------------------------

// page-tree-deep.pdf has its /MediaBox and /Rotate on the root, three /Pages
// levels above the leaves, and its second page overrides /Rotate alone. That
// the override does not also drop the inherited box is the whole point of the
// file.
@(test)
test_golden_page_tree_inheritance :: proc(t: ^testing.T) {
    defer free_all(context.temp_allocator)

    doc, err := open_file(golden_path("page-tree-deep.pdf"))
    defer close(doc)
    if !testing.expect(t, err == .None, "open failed") { return }

    p0, e0 := page_get(doc, 0)
    testing.expect(t, e0 == .None, "page 0 missing")
    testing.expectf(t, p0.rotate == 90, "page 0 /Rotate %d, wanted the inherited 90", p0.rotate)

    p1, e1 := page_get(doc, 1)
    if !testing.expect(t, e1 == .None, "page 1 missing") { return }
    testing.expectf(t, p1.rotate == 180, "page 1 /Rotate %d, wanted its own 180", p1.rotate)
    testing.expectf(t, rect_width(p1.media_box) == 300 && rect_height(p1.media_box) == 400,
        "page 1 is %vx%v; overriding /Rotate must not drop the inherited /MediaBox",
        rect_width(p1.media_box), rect_height(p1.media_box))

    // Inherited two levels down from the node that declares it.
    testing.expect(t, p0.resources != nil, "page 0 did not inherit /Resources")
    testing.expect(t, p1.resources != nil, "page 1 did not inherit /Resources")
}

// Each encrypted file also has an owner password, so the owner branch of the
// key derivation is exercised by a real file and not only by a unit vector.
@(test)
test_golden_owner_password :: proc(t: ^testing.T) {
    defer free_all(context.temp_allocator)

    for g in GOLDEN {
        if g.method == "" { continue }

        doc, err := open_file(golden_path(g.name), Open_Options{password = GOLDEN_OWNER_PW})
        defer close(doc)
        if !testing.expectf(t, err == .None, "%s: owner password rejected: %s",
            g.name, error_string(err)) {
            continue
        }
        testing.expectf(t, doc.encrypt != nil && doc.encrypt.authenticated_as == .Owner,
            "%s: authenticated as %v, wanted Owner", g.name,
            doc.encrypt != nil ? doc.encrypt.authenticated_as : Password_Kind.None)

        // The owner is not subject to the restrictions they set.
        all := Permissions{.Print, .Modify, .Copy, .Annotate, .Fill_Forms, .Extract, .Assemble, .Print_High_Res}
        testing.expectf(t, encryption_permissions(doc) == all,
            "%s: owner permissions %v, wanted everything", g.name, encryption_permissions(doc))

        page, perr := page_get(doc, 0)
        testing.expectf(t, perr == .None, "%s: page 0: %s", g.name, error_string(perr))
        body, cerr := golden_content(doc, page)
        testing.expectf(t, cerr == .None && strings.has_prefix(body, g.marker),
            "%s: content did not decrypt under the owner password", g.name)
    }
}

// The files restrict what a user may do, so /P must survive the round trip.
// -44 clears the modify and annotate bits and leaves the rest set; revision 2
// only ever defined four of them, which is why it reports two.
@(test)
test_golden_user_permissions :: proc(t: ^testing.T) {
    defer free_all(context.temp_allocator)

    Case :: struct {
        name: string,
        want: Permissions,
    }
    r3_up := Permissions{.Print, .Copy, .Fill_Forms, .Extract, .Assemble, .Print_High_Res}

    cases := []Case{
        {"encrypted-rc4-40.pdf",  {.Print, .Copy}},
        {"encrypted-rc4-128.pdf", r3_up},
        {"encrypted-aes128.pdf",  r3_up},
        {"encrypted-aes256.pdf",  r3_up},
    }

    for c in cases {
        doc, err := open_file(golden_path(c.name))
        defer close(doc)
        if !testing.expectf(t, err == .None, "%s: %s", c.name, error_string(err)) { continue }

        testing.expectf(t, encryption_permissions(doc) == c.want,
            "%s: permissions %v, wanted %v", c.name, encryption_permissions(doc), c.want)
        testing.expectf(t, doc.encrypt != nil && doc.encrypt.authenticated_as == .User,
            "%s: the empty password must authenticate as the user", c.name)
    }
}

// The table and the directory must not drift apart. Without this, a golden file
// added to testdata/ and forgotten in GOLDEN would be tested by nothing at all,
// and a row whose file was never committed would be a silent skip.
@(test)
test_golden_directory_matches_table :: proc(t: ^testing.T) {
    defer free_all(context.temp_allocator)

    entries, derr := os.read_directory_by_path(GOLDEN_DIR, -1, context.allocator)
    if !testing.expectf(t, derr == nil, "cannot list %s: %v", GOLDEN_DIR, derr) { return }
    defer os.file_info_slice_delete(entries, context.allocator)

    for e in entries {
        if e.type != .Regular || !strings.has_suffix(e.name, ".pdf") { continue }
        found := false
        for g in GOLDEN {
            if g.name == e.name { found = true; break }
        }
        testing.expectf(t, found, "testdata/%s is not in the GOLDEN table", e.name)
    }

    for g in GOLDEN {
        present := false
        for e in entries {
            if e.name == g.name { present = true; break }
        }
        testing.expectf(t, present, "GOLDEN lists %s but testdata/ does not have it", g.name)
    }
}

// A password that is neither the user's nor the owner's is refused rather than
// silently producing garbage.
@(test)
test_golden_wrong_password :: proc(t: ^testing.T) {
    defer free_all(context.temp_allocator)

    for g in GOLDEN {
        if g.method == "" { continue }
        doc, err := open_file(golden_path(g.name), Open_Options{password = "not the password"})
        defer close(doc)
        testing.expectf(t, err == .Wrong_Password,
            "%s: opened with a wrong password: %s", g.name, error_string(err))
        testing.expectf(t, doc == nil, "%s: a failed open must not return a document", g.name)
    }
}

// ---------------------------------------------------------------------------
// Fuzz-shaped: truncations and bit flips
//
// Not a fuzzer -- there is no feedback and no corpus growth -- but the same
// shape, and enough to catch the class of bug that matters here: an index
// derived from a file byte that is no longer in range once that byte changes.
// The mutations are driven by a fixed seed so that a failure is reproducible;
// a clock-seeded run that fails once and never again is worse than no test.
// ---------------------------------------------------------------------------

@(private = "file")
GOLDEN_FUZZ_TRUNCATIONS :: 40
@(private = "file")
GOLDEN_FUZZ_FLIPS :: 200
@(private = "file")
GOLDEN_FUZZ_SEED :: u64(0x9E3779B97F4A7C15)

// A single mutated file that takes this long has found a loop bound that is
// too generous. These are all under 1.5 KB; nothing legitimate comes close.
@(private = "file")
GOLDEN_FUZZ_SLOW_MS :: 5000

@(private = "file")
golden_rand :: proc(state: ^u64) -> u64 {
    state^ += 0x9E37_79B9_7F4A_7C15
    z := state^
    z = (z ~ (z >> 30)) * 0xBF58_476D_1CE4_E5B9
    z = (z ~ (z >> 27)) * 0x94D0_49BB_1331_11EB
    return z ~ (z >> 31)
}

// golden_probe opens one mutated file and drives it as far as a caller would.
// The assertion is only that this returns at all, with either a Document or an
// Error and never both: what a damaged file *means* is not defined, but
// crashing, hanging and handing back a half-built Document are all excluded.
@(private = "file")
golden_probe :: proc(t: ^testing.T, data: []byte, what: string) {
    start := time.now()

    doc, err := open_bytes(data)
    defer close(doc)

    if err == .None {
        testing.expectf(t, doc != nil, "%s: no error but no document", what)
        if doc != nil {
            for i in 0 ..< page_count(doc) {
                page, perr := page_get(doc, i)
                if perr != .None { continue }
                // Reading the stream is where a bad /Length or a corrupt Flate
                // payload would be dereferenced, so it belongs inside the probe.
                s := dict_stream(doc, page.dict, "Contents")
                _, _ = stream_data(doc, s)
            }
        }
    } else {
        testing.expectf(t, doc == nil, "%s: error %s but a document too", what, error_string(err))
    }

    ms := time.duration_milliseconds(time.since(start))
    testing.expectf(t, ms < GOLDEN_FUZZ_SLOW_MS, "%s: took %v ms", what, ms)
}

@(test)
test_golden_fuzz_truncated :: proc(t: ^testing.T) {
    defer free_all(context.temp_allocator)

    for g in GOLDEN {
        data, read_err := os.read_entire_file(golden_path(g.name), context.allocator)
        if !testing.expectf(t, read_err == nil, "%s: cannot read: %v", g.name, read_err) { continue }
        defer delete(data)

        // Evenly spaced, and never the whole file: 0 bytes and every prefix in
        // between, which is what a download that stopped early looks like.
        for i in 0 ..< GOLDEN_FUZZ_TRUNCATIONS {
            n := len(data) * i / GOLDEN_FUZZ_TRUNCATIONS
            golden_probe(t, data[:n], g.name)
        }
    }
}

@(test)
test_golden_fuzz_bit_flips :: proc(t: ^testing.T) {
    defer free_all(context.temp_allocator)

    for g, index in GOLDEN {
        data, read_err := os.read_entire_file(golden_path(g.name), context.allocator)
        if !testing.expectf(t, read_err == nil, "%s: cannot read: %v", g.name, read_err) { continue }
        defer delete(data)
        if len(data) == 0 { continue }

        scratch := make([]byte, len(data))
        defer delete(scratch)

        // Seeded from the file's position in the table, so every file gets a
        // different sequence and every run gets the same one.
        state := GOLDEN_FUZZ_SEED + u64(index) * 0x0100_0193
        for _ in 0 ..< GOLDEN_FUZZ_FLIPS {
            copy(scratch, data)
            off := int(golden_rand(&state) % u64(len(scratch)))
            bit := u8(golden_rand(&state) & 7)
            scratch[off] ~= 1 << bit
            golden_probe(t, scratch, g.name)
        }
    }
}