package pdf

import "core:mem"
import "core:strings"

// The COS tokeniser (ISO 32000-2, 7.2).
//
// It is deliberately total: every call to `lexer_next` either advances `pos`
// or returns an EOF token, so no input -- valid, truncated or random -- can
// make a caller spin. Nothing here rejects input either; the malformations
// this layer meets (a string that never closes, a hex digit missing its pair,
// a number that is three signs and a dot) are all common enough in real files
// that failing on them would mean refusing documents every viewer opens. The
// damage is reported through `Token.truncated` and through the values
// themselves, and the parser turns that into a `Repair`.

// ---------------------------------------------------------------------------
// Character classes (ISO 32000-2, 7.2.3, tables 1 and 2)
// ---------------------------------------------------------------------------

is_pdf_whitespace :: proc "contextless" (c: byte) -> bool {
    switch c {
    case 0x00, 0x09, 0x0A, 0x0C, 0x0D, 0x20: return true
    }
    return false
}

is_pdf_delimiter :: proc "contextless" (c: byte) -> bool {
    switch c {
    case '(', ')', '<', '>', '[', ']', '{', '}', '/', '%': return true
    }
    return false
}

// A regular character is anything that is neither whitespace nor a delimiter;
// names, numbers and keywords are runs of them.
is_pdf_regular :: proc "contextless" (c: byte) -> bool {
    return !is_pdf_whitespace(c) && !is_pdf_delimiter(c)
}

// ---------------------------------------------------------------------------
// Tokens
// ---------------------------------------------------------------------------

Token_Kind :: enum u8 {
    EOF,
    Integer,
    Real,
    String,
    Name,
    Array_Open,
    Array_Close,
    Dict_Open,
    Dict_Close,
    Brace_Open,     // only meaningful inside a type 4 (PostScript) function
    Brace_Close,
    Keyword,        // obj, endobj, stream, R, true, false, null, operators
    Error,          // a stray delimiter; one byte was consumed
}

Token :: struct {
    kind:  Token_Kind,
    // Byte range of the token in the lexer's input: [start, end).
    start: int,
    end:   int,
    // Both numeric fields are filled for both numeric kinds, so a caller that
    // wants "a number" never has to switch on the kind. A value that does not
    // fit, or that parsed to an infinity, reads as 0 -- see `parse_number`.
    int_value:  i64,
    real_value: f64,
    // Names (with #xx resolved) and keywords. Borrows from the input unless
    // the name contained an escape, in which case it was allocated with the
    // allocator passed to `lexer_next`.
    text:  string,
    // String objects, unescaped and hex-unpacked. Always allocated with the
    // allocator passed to `lexer_next` (never a view into the input): a
    // decryption pass may have to rewrite these bytes in place, and the input
    // is the file itself.
    bytes: []byte,
    // The token ran into the end of the input before its closing delimiter.
    // The value is still usable -- this is a repair, not an error.
    truncated: bool,
}

Lexer :: struct {
    data: []byte,
    pos:  int,
}

lexer_make :: proc(data: []byte) -> Lexer {
    return Lexer{data = data, pos = 0}
}

// lexer_skip_whitespace steps over whitespace and comments. A '%' outside of a
// string runs to the end of the line and is whitespace for every purpose
// (ISO 32000-2, 7.2.4).
lexer_skip_whitespace :: proc(lex: ^Lexer) {
    for lex.pos < len(lex.data) {
        c := lex.data[lex.pos]
        if is_pdf_whitespace(c) {
            lex.pos += 1
            continue
        }
        if c == '%' {
            lex.pos += 1
            for lex.pos < len(lex.data) {
                d := lex.data[lex.pos]
                if d == '\r' || d == '\n' { break }
                lex.pos += 1
            }
            continue
        }
        return
    }
}

// lexer_next returns the next token, advancing past it.
//
// `allocator` owns any bytes the token had to materialise: the payload of a
// string, and the text of a name that contained #xx escapes. Callers holding a
// Document pass `doc.allocator` so the token's memory dies with the Document.
lexer_next :: proc(lex: ^Lexer, allocator := context.allocator) -> Token {
    lexer_skip_whitespace(lex)
    if lex.pos >= len(lex.data) {
        return Token{kind = .EOF, start = lex.pos, end = lex.pos}
    }

    start := lex.pos
    c := lex.data[start]

    switch c {
    case '[':
        lex.pos += 1
        return Token{kind = .Array_Open, start = start, end = lex.pos}
    case ']':
        lex.pos += 1
        return Token{kind = .Array_Close, start = start, end = lex.pos}
    case '{':
        lex.pos += 1
        return Token{kind = .Brace_Open, start = start, end = lex.pos}
    case '}':
        lex.pos += 1
        return Token{kind = .Brace_Close, start = start, end = lex.pos}
    case '(':
        return lex_literal_string(lex, allocator)
    case '/':
        return lex_name(lex, allocator)
    case '<':
        if start + 1 < len(lex.data) && lex.data[start + 1] == '<' {
            lex.pos += 2
            return Token{kind = .Dict_Open, start = start, end = lex.pos}
        }
        return lex_hex_string(lex, allocator)
    case '>':
        if start + 1 < len(lex.data) && lex.data[start + 1] == '>' {
            lex.pos += 2
            return Token{kind = .Dict_Close, start = start, end = lex.pos}
        }
        // A lone '>' is junk. Consume it so the caller makes progress.
        lex.pos += 1
        return Token{kind = .Error, start = start, end = lex.pos}
    case ')':
        lex.pos += 1
        return Token{kind = .Error, start = start, end = lex.pos}
    case '+', '-', '.', '0' ..= '9':
        return lex_number(lex)
    }

    if is_pdf_regular(c) {
        return lex_keyword(lex)
    }

    // Unreachable for the delimiters above, but a delimiter added to the class
    // table later must still consume a byte rather than stall.
    lex.pos += 1
    return Token{kind = .Error, start = start, end = lex.pos}
}

// lexer_next_structural returns the next token without ever allocating: the
// kinds whose value has to be materialised -- strings, and names carrying #xx
// escapes -- come back with an empty payload. Their `kind`, `start` and `end`
// are still right, which is all a structural probe looks at.
//
// The probes that hunt for "num gen obj" or an xref subsection line accept
// integers and keywords only, and in a damaged file they run at every
// plausible offset and reject nearly all of them. Lexing those rejects with a
// real allocator copies the string bodies they step over -- a file reading
// "1 (" at a hundred probed offsets would copy its own tail a hundred times,
// and whichever of those copies is not attached to a Document leaks outright.
lexer_next_structural :: proc(lex: ^Lexer) -> Token {
    return lexer_next(lex, mem.nil_allocator())
}

// lexer_peek returns the next token without consuming it. It may allocate
// exactly as `lexer_next` does, so peeking at a string costs a copy; the
// parser keeps its own token queue instead of peeking for that reason.
lexer_peek :: proc(lex: ^Lexer, allocator := context.allocator) -> Token {
    save := lex.pos
    tok := lexer_next(lex, allocator)
    lex.pos = save
    return tok
}

// lexer_at_keyword reports whether the next token is exactly `kw`, without
// consuming anything. Used for the keywords that are found by position rather
// than parsed -- `stream`, `endstream`, `endobj`.
lexer_at_keyword :: proc(lex: ^Lexer, kw: string) -> bool {
    save := lex.pos
    defer lex.pos = save
    lexer_skip_whitespace(lex)
    return lexer_match_here(lex, kw)
}

// lexer_accept_keyword consumes `kw` if it is next, and reports whether it did.
lexer_accept_keyword :: proc(lex: ^Lexer, kw: string) -> bool {
    save := lex.pos
    lexer_skip_whitespace(lex)
    if lexer_match_here(lex, kw) {
        lex.pos += len(kw)
        return true
    }
    lex.pos = save
    return false
}

@(private = "file")
lexer_match_here :: proc(lex: ^Lexer, kw: string) -> bool {
    end := lex.pos + len(kw)
    if end > len(lex.data) { return false }
    if string(lex.data[lex.pos:end]) != kw { return false }
    // "stream" must not match the first six bytes of "streamable".
    if end < len(lex.data) && is_pdf_regular(lex.data[end]) { return false }
    return true
}

// lexer_skip_eol steps over one end-of-line marker: CRLF, LF or a bare CR.
// The bare CR is not legal after `stream` but is written by enough producers
// that rejecting it would fail real documents (ISO 32000-2, 7.3.8.1).
lexer_skip_eol :: proc(lex: ^Lexer) {
    if lex.pos < len(lex.data) && lex.data[lex.pos] == '\r' {
        lex.pos += 1
    }
    if lex.pos < len(lex.data) && lex.data[lex.pos] == '\n' {
        lex.pos += 1
    }
}

// ---------------------------------------------------------------------------
// Numbers
// ---------------------------------------------------------------------------

@(private = "file")
lex_number :: proc(lex: ^Lexer) -> Token {
    start := lex.pos
    for lex.pos < len(lex.data) && is_pdf_regular(lex.data[lex.pos]) {
        lex.pos += 1
    }
    text := string(lex.data[start:lex.pos])

    tok := Token{start = start, end = lex.pos}
    // parse_number never reports failure as an exception: a number a viewer
    // would read as 0 must not abort the object around it.
    value, _ := parse_number(text)
    #partial switch v in value {
    case i64:
        tok.kind = .Integer
        tok.int_value = v
        tok.real_value = f64(v)
    case f64:
        tok.kind = .Real
        tok.real_value = clamp_finite(v, 0)
        tok.int_value = f64_to_i64_saturating(tok.real_value)
    case:
        tok.kind = .Integer
    }
    return tok
}

// f64_to_i64_saturating clamps before converting: an out-of-range float to
// integer conversion is undefined behaviour, and /Length values of 1e300 exist.
@(private)
f64_to_i64_saturating :: proc "contextless" (v: f64) -> i64 {
    if !is_finite(v)      { return 0 }
    if v >= 9.2e18        { return max(i64) }
    if v <= -9.2e18       { return min(i64) }
    return i64(v)
}

// ---------------------------------------------------------------------------
// Names (ISO 32000-2, 7.3.5)
// ---------------------------------------------------------------------------

@(private = "file")
lex_name :: proc(lex: ^Lexer, allocator: mem.Allocator) -> Token {
    start := lex.pos
    lex.pos += 1                    // the '/'
    body := lex.pos
    for lex.pos < len(lex.data) && is_pdf_regular(lex.data[lex.pos]) {
        lex.pos += 1
    }
    raw := string(lex.data[body:lex.pos])
    // A bare '/' is the empty name, which is a legal object.
    return Token{
        kind  = .Name,
        start = start,
        end   = lex.pos,
        text  = decode_name_escapes(raw, allocator),
    }
}

// ---------------------------------------------------------------------------
// Keywords
// ---------------------------------------------------------------------------

@(private = "file")
lex_keyword :: proc(lex: ^Lexer) -> Token {
    start := lex.pos
    for lex.pos < len(lex.data) && is_pdf_regular(lex.data[lex.pos]) {
        lex.pos += 1
    }
    return Token{
        kind  = .Keyword,
        start = start,
        end   = lex.pos,
        text  = string(lex.data[start:lex.pos]),
    }
}

// ---------------------------------------------------------------------------
// Literal strings (ISO 32000-2, 7.3.4.2)
// ---------------------------------------------------------------------------

@(private = "file")
lex_literal_string :: proc(lex: ^Lexer, allocator: mem.Allocator) -> Token {
    start := lex.pos
    lex.pos += 1                    // the '('

    // The unescaped form is never longer than the escaped one, so one
    // allocation of the remaining length is always enough and never wrong.
    // Size the buffer from the raw body once, up front: the unescaped value is
    // never longer than the bytes it came from, so this both fits and cannot
    // over-allocate on a file that claims a huge string it does not contain.
    b := strings.builder_make(0, literal_body_len(lex.data, lex.pos), allocator)
    depth := 1
    truncated := true

    scan: for lex.pos < len(lex.data) {
        c := lex.data[lex.pos]
        lex.pos += 1
        switch c {
        case '\\':
            if lex.pos >= len(lex.data) { break scan }
            e := lex.data[lex.pos]
            lex.pos += 1
            switch e {
            case 'n':  strings.write_byte(&b, '\n')
            case 'r':  strings.write_byte(&b, '\r')
            case 't':  strings.write_byte(&b, '\t')
            case 'b':  strings.write_byte(&b, '\b')
            case 'f':  strings.write_byte(&b, '\f')
            case '(':  strings.write_byte(&b, '(')
            case ')':  strings.write_byte(&b, ')')
            case '\\': strings.write_byte(&b, '\\')
            case '\r':
                // Backslash before an EOL is a line continuation: both the
                // backslash and the EOL disappear.
                if lex.pos < len(lex.data) && lex.data[lex.pos] == '\n' {
                    lex.pos += 1
                }
            case '\n':
                // continuation
            case '0' ..= '7':
                value := int(e - '0')
                for _ in 0 ..< 2 {
                    if lex.pos >= len(lex.data) { break }
                    d := lex.data[lex.pos]
                    if d < '0' || d > '7' { break }
                    value = value * 8 + int(d - '0')
                    lex.pos += 1
                }
                strings.write_byte(&b, byte(value & 0xFF))
            case:
                // "\q" is just "q": the backslash is dropped, the byte kept.
                strings.write_byte(&b, e)
            }
        case '(':
            depth += 1
            strings.write_byte(&b, '(')
        case ')':
            depth -= 1
            if depth == 0 {
                truncated = false
                break scan
            }
            strings.write_byte(&b, ')')
        case '\r':
            // A raw CR, LF or CRLF in the source is one LF in the value.
            if lex.pos < len(lex.data) && lex.data[lex.pos] == '\n' {
                lex.pos += 1
            }
            strings.write_byte(&b, '\n')
        case:
            strings.write_byte(&b, c)
        }
    }

    return Token{
        kind      = .String,
        start     = start,
        end       = lex.pos,
        bytes     = transmute([]byte)strings.to_string(b),
        truncated = truncated,
    }
}

// ---------------------------------------------------------------------------
// Hex strings (ISO 32000-2, 7.3.4.3)
// ---------------------------------------------------------------------------

@(private = "file")
lex_hex_string :: proc(lex: ^Lexer, allocator: mem.Allocator) -> Token {
    start := lex.pos
    lex.pos += 1                    // the '<'

    b := strings.builder_make(0, hex_body_len(lex.data, lex.pos) / 2 + 1, allocator)
    nibble: byte
    have_nibble := false
    truncated := true

    for lex.pos < len(lex.data) {
        c := lex.data[lex.pos]
        lex.pos += 1
        if c == '>' {
            truncated = false
            break
        }
        v, ok := hex_value(c)
        // Whitespace is defined to be skipped here, and other junk is skipped
        // too: files with a stray '<' or a UTF-8 quote inside a hex string are
        // repaired by every viewer rather than rejected.
        if !ok { continue }
        if have_nibble {
            strings.write_byte(&b, nibble << 4 | v)
            have_nibble = false
        } else {
            nibble = v
            have_nibble = true
        }
    }
    // "An odd number of digits: the final digit is assumed to be followed by
    // a zero." (7.3.4.3)
    if have_nibble {
        strings.write_byte(&b, nibble << 4)
    }

    return Token{
        kind      = .String,
        start     = start,
        end       = lex.pos,
        bytes     = transmute([]byte)strings.to_string(b),
        truncated = truncated,
    }
}

// literal_body_len counts the raw bytes of a literal string body starting at
// `pos` (just past the '('), stopping at the balancing ')' or at EOF. It is an
// upper bound on the decoded length, which is what the value buffer is sized
// from.
@(private = "file")
literal_body_len :: proc(data: []byte, pos: int) -> int {
    depth := 1
    i := pos
    for i < len(data) {
        switch data[i] {
        case '\\':
            i += 1              // whatever follows is data, never a delimiter
        case '(':
            depth += 1
        case ')':
            depth -= 1
            if depth == 0 {
                return i - pos
            }
        }
        i += 1
    }
    return len(data) - pos
}

@(private = "file")
hex_body_len :: proc(data: []byte, pos: int) -> int {
    for i := pos; i < len(data); i += 1 {
        if data[i] == '>' {
            return i - pos
        }
    }
    return len(data) - pos
}