"""Byte-level structural features used by pdf-damage-classifier.

The same function is used for training and inference, so the ONNX model's
input vector is defined by FEATURE_NAMES (in this exact order).
Only the PDF's raw bytes are needed; no PDF parser is involved, so it works
on files that are too broken to open.
"""
import math
import re

FEATURE_NAMES = [
    "has_pdf_header",
    "eof_in_last_1k",
    "ends_with_eof",
    "startxref_present",
    "startxref_numeric",
    "startxref_points_to_xref",
    "xref_valid_fraction",
    "has_trailer",
    "root_object_exists",
    "obj_endobj_gap",
    "stream_endstream_gap",
    "log10_size",
]

_XREF_ENTRY = re.compile(rb"(\d{10}) (\d{5}) n")
_STARTXREF = re.compile(rb"startxref\s+(\S+)")
_ROOT = re.compile(rb"/Root\s+(\d+)\s+(\d+)\s+R")


def extract_features(data: bytes):
    n = len(data)
    tail = data[-1024:]

    has_header = float(data[:5] == b"%PDF-")
    eof_in_tail = float(b"%%EOF" in tail)
    ends_eof = float(data.rstrip()[-5:] == b"%%EOF")

    m = None
    for m in _STARTXREF.finditer(tail):
        pass
    sx_present = float(m is not None)
    sx_numeric = 0.0
    sx_points = 0.0
    if m is not None:
        tok = m.group(1)
        if tok.isdigit():
            sx_numeric = 1.0
            off = int(tok)
            if off < n and data[off : off + 4] == b"xref":
                sx_points = 1.0

    entries = _XREF_ENTRY.findall(data)
    if entries:
        ok = 0
        for off_s, _gen in entries:
            off = int(off_s)
            if off < n and re.match(rb"\d+ \d+ obj", data[off : off + 24]):
                ok += 1
        xref_frac = ok / len(entries)
    else:
        xref_frac = 0.0

    has_trailer = float(b"trailer" in data)

    rm = _ROOT.search(data)
    root_exists = 0.0
    if rm is not None:
        root_exists = float(re.search(rb"(?:^|\s)" + rm.group(1) + b" " + rm.group(2) + b" obj", data) is not None)

    objs = len(re.findall(rb"\d+ \d+ obj\b", data))
    endobjs = data.count(b"endobj")
    # Count stream starts, ignoring the "stream\n" inside "endstream\n".
    stream_starts = len(re.findall(rb"(?<!end)stream\r?\n", data))
    endstreams = data.count(b"endstream")

    obj_gap = (objs - endobjs) / max(objs, 1)
    stream_gap = (stream_starts - endstreams) / max(stream_starts, 1)

    return [
        has_header,
        eof_in_tail,
        ends_eof,
        sx_present,
        sx_numeric,
        sx_points,
        xref_frac,
        has_trailer,
        root_exists,
        obj_gap,
        stream_gap,
        math.log10(max(n, 1)),
    ]