Extract the fields of an invoice

"""
Read the header fields of an invoice from text that has already been
extracted: keyword anchors, then regular expressions.

Rung N0. No model, no training set, no service. Two hours of work and you can
read every invoice from the supplier you wrote it for.

The method is the one anybody reaches for: find the line carrying the label,
then read the value that follows it on that line. The labels are ordered from
the most specific to the least, because « Total TTC » and « Total HT » are one
word apart and the wrong one is a plausible number.

That ordering is also where the approach ends. See the test.
"""

import re

# French amounts: a comma before the decimals, a space or a dot every three
# digits. Requiring exactly three digits per group is what keeps the pattern
# from swallowing a quantity and a unit price as one number.
AMOUNT = r"\d{1,3}(?:[\s.]\d{3})*[,.]\d{2}"
DATE = r"\d{1,2}/\d{1,2}/\d{2,4}"
REFERENCE = r"\b(?:[A-Za-z]{1,3}[-/])?\d[\dA-Za-z/-]{3,}"


def parse_amount(raw: str) -> float:
    """Turn a written amount into a number the caller can compute with."""
    cleaned = re.sub(r"[^\d,.]", "", raw)
    # A comma means French spelling: the dots left are thousands separators.
    if "," in cleaned:
        cleaned = cleaned.replace(".", "").replace(",", ".")
    return float(cleaned)


# Per field: the labels to look for, most specific first, the pattern the value
# must match after the label, and how to read the match.
FIELDS = {
    "invoice_number": (("facture n°", "facture no", "n° facture"), REFERENCE, str.strip),
    "date": (("date",), DATE, str.strip),
    "total": (("total ttc", "montant ttc", "total"), AMOUNT, parse_amount),
}


def find_after_label(text: str, labels: tuple[str, ...], pattern: str) -> str | None:
    """
    First value matching `pattern` after one of `labels`, on the same line.

    A label that appears on a line holding no value is skipped rather than
    accepted, because a column heading is a label too.
    """
    lines = text.splitlines()
    for label in labels:
        for line in lines:
            position = line.lower().find(label)
            if position == -1:
                continue
            match = re.search(pattern, line[position + len(label):])
            if match:
                return match.group(0)
    return None


def extract_fields(text: str) -> dict:
    """Read the invoice number, the date and the total from extracted text."""
    fields = {}
    for name, (labels, pattern, read) in FIELDS.items():
        raw = find_after_label(text, labels, pattern)
        fields[name] = read(raw) if raw is not None else None
    return fields

/**
 * Read the header fields of an invoice from text that has already been
 * extracted: keyword anchors, then regular expressions.
 *
 * Rung N0. No model, no training set, no service. Two hours of work and you
 * can read every invoice from the supplier you wrote it for.
 *
 * The method is the one anybody reaches for: find the line carrying the label,
 * then read the value that follows it on that line. The labels are ordered
 * from the most specific to the least, because « Total TTC » and « Total HT »
 * are one word apart and the wrong one is a plausible number.
 *
 * That ordering is also where the approach ends. See the test.
 */

// French amounts: a comma before the decimals, a space or a dot every three
// digits. Requiring exactly three digits per group is what keeps the pattern
// from swallowing a quantity and a unit price as one number.
export const AMOUNT = /\d{1,3}(?:[\s.]\d{3})*[,.]\d{2}/;
const DATE = /\d{1,2}\/\d{1,2}\/\d{2,4}/;
const REFERENCE = /\b(?:[A-Za-z]{1,3}[-/])?\d[\dA-Za-z/-]{3,}/;

/** Turn a written amount into a number the caller can compute with. */
export function parseAmount(raw) {
  let cleaned = raw.replace(/[^\d,.]/g, '');
  // A comma means French spelling: the dots left are thousands separators.
  if (cleaned.includes(',')) cleaned = cleaned.replaceAll('.', '').replace(',', '.');
  return Number(cleaned);
}

const trim = (raw) => raw.trim();

// Per field: the labels to look for, most specific first, the pattern the value
// must match after the label, and how to read the match.
const FIELDS = {
  invoice_number: [['facture n°', 'facture no', 'n° facture'], REFERENCE, trim],
  date: [['date'], DATE, trim],
  total: [['total ttc', 'montant ttc', 'total'], AMOUNT, parseAmount],
};

/**
 * First value matching `pattern` after one of `labels`, on the same line.
 *
 * A label that appears on a line holding no value is skipped rather than
 * accepted, because a column heading is a label too.
 */
export function findAfterLabel(text, labels, pattern) {
  const lines = text.split('\n');
  for (const label of labels) {
    for (const line of lines) {
      const position = line.toLowerCase().indexOf(label);
      if (position === -1) continue;
      const match = line.slice(position + label.length).match(pattern);
      if (match) return match[0];
    }
  }
  return null;
}

/** Read the invoice number, the date and the total from extracted text. */
export function extractFields(text) {
  const fields = {};
  for (const [name, [labels, pattern, read]] of Object.entries(FIELDS)) {
    const raw = findAfterLabel(text, labels, pattern);
    fields[name] = raw === null ? null : read(raw);
  }
  return fields;
}

"""
Label each line of the invoice, then read the value out of the line.

Rung N1. N0 asked « what does this line say ». This asks « where does this line
sit, and what does it look like »: how far down the page, how indented, how
wordy, how many amounts, and whether the value hangs on the right-hand side.

Those features survive a change of supplier, which is exactly what the labels
of N0 do not. Training is a few dozen labelled lines, the model is a few
kilobytes, and nothing is downloaded.
"""

import re

from sklearn.linear_model import LogisticRegression

AMOUNT = r"\d{1,3}(?:[\s.]\d{3})*[,.]\d{2}"
MONTHS = "janvier|février|mars|avril|mai|juin|juillet|août|septembre|octobre|novembre|décembre"
DATE = rf"\d{{1,2}}/\d{{1,2}}/\d{{2,4}}|\d{{1,2}}\s+(?:{MONTHS})\s+\d{{4}}"
REFERENCE = r"\b(?:[A-Za-z]{1,3}[-/])?\d[\dA-Za-z/-]{3,}"


def page_lines(text: str) -> list[str]:
    """The lines that carry something, indentation kept: it is a feature."""
    return [line for line in text.splitlines() if line.strip()]


def line_features(line: str, index: int, count: int) -> list[float]:
    """Where the line sits and what it looks like. Never what it says."""
    text = line.strip()
    amounts = re.findall(AMOUNT, text)
    letters = "".join(c for c in text if c.isalpha())
    return [
        index / max(count - 1, 1),                          # how far down the page
        1.0 if index == count - 1 else 0.0,                 # the very last line
        min(len(line) - len(line.lstrip()), 40) / 40,       # indentation
        min(len(text.split()), 12) / 12,                    # how wordy
        sum(c.isdigit() for c in text) / len(text),         # digit share
        min(len(amounts), 3) / 3,                           # how many amounts
        1.0 if re.search(DATE, text) else 0.0,
        1.0 if re.search(REFERENCE, text) else 0.0,
        # A value hanging on the right of the line, the way a totals block does.
        1.0 if amounts and text.rindex(amounts[-1]) > len(text) / 2 else 0.0,
        1.0 if letters and letters.isupper() else 0.0,
    ]


def train(documents: list[str], labels: list[list[str]]):
    """`labels` carries one label per non-blank line of each document."""
    rows, targets = [], []
    for document, document_labels in zip(documents, labels):
        lines = page_lines(document)
        rows += [line_features(line, i, len(lines)) for i, line in enumerate(lines)]
        targets += list(document_labels)
    model = LogisticRegression(class_weight="balanced", max_iter=1000)
    model.fit(rows, targets)
    return model


def _match(pattern: str):
    """Reads the first value of that shape out of a line, or nothing."""
    def read(line):
        found = re.search(pattern, line)
        return found.group(0) if found else None

    return read


def _amount(line):
    """The rightmost amount: an item line carries a quantity and a unit price."""
    amounts = re.findall(AMOUNT, line)
    if not amounts:
        return None
    cleaned = re.sub(r"[^\d,.]", "", amounts[-1])
    # A comma means French spelling: the dots left are thousands separators.
    if "," in cleaned:
        cleaned = cleaned.replace(".", "").replace(",", ".")
    return float(cleaned)


READERS = {"invoice_number": _match(REFERENCE), "date": _match(DATE), "total": _amount}


def extract_fields(model, text: str) -> dict:
    """
    For each field, walk the lines from the most likely down.

    The model points at a line; reading a date or an amount out of it is still
    ours to do, and a line the model likes but that holds no value is not an
    answer.
    """
    lines = page_lines(text)
    if not lines:
        return {name: None for name in READERS}
    rows = [line_features(line, i, len(lines)) for i, line in enumerate(lines)]
    scores = model.predict_proba(rows)
    classes = list(model.classes_)
    fields = {}
    for name, read in READERS.items():
        column = classes.index(name)
        ranked = sorted(range(len(lines)), key=lambda i: -scores[i][column])
        fields[name] = next((v for v in (read(lines[i]) for i in ranked) if v is not None), None)
    return fields

/**
 * Label each line of the invoice, then read the value out of the line.
 *
 * Rung N1. N0 asked « what does this line say ». This asks « where does this
 * line sit, and what does it look like »: how far down the page, how indented,
 * how wordy, how many amounts, and whether the value hangs on the right-hand
 * side.
 *
 * Those features survive a change of supplier, which is exactly what the
 * labels of N0 do not. Training is a few dozen labelled lines, and logistic
 * regression is short enough to be read rather than imported.
 */

const AMOUNT = /\d{1,3}(?:[\s.]\d{3})*[,.]\d{2}/g;
const MONTHS = 'janvier|février|mars|avril|mai|juin|juillet|août|septembre|octobre|novembre|décembre';
const DATE = new RegExp(`\\d{1,2}/\\d{1,2}/\\d{2,4}|\\d{1,2}\\s+(?:${MONTHS})\\s+\\d{4}`);
const REFERENCE = /\b(?:[A-Za-z]{1,3}[-/])?\d[\dA-Za-z/-]{3,}/;

/** The lines that carry something, indentation kept: it is a feature. */
export function pageLines(text) {
  return text.split('\n').filter((line) => line.trim() !== '');
}

/** Where the line sits and what it looks like. Never what it says. */
export function lineFeatures(line, index, count) {
  const text = line.trim();
  const amounts = text.match(AMOUNT) ?? [];
  const letters = text.replace(/[^\p{L}]/gu, '');
  const words = text.split(/\s+/).filter(Boolean);
  return [
    index / Math.max(count - 1, 1), //                    how far down the page
    index === count - 1 ? 1 : 0, //                       the very last line
    Math.min(line.length - line.trimStart().length, 40) / 40, // indentation
    Math.min(words.length, 12) / 12, //                   how wordy
    (text.match(/\d/g) ?? []).length / text.length, //    digit share
    Math.min(amounts.length, 3) / 3, //                   how many amounts
    DATE.test(text) ? 1 : 0,
    REFERENCE.test(text) ? 1 : 0,
    // A value hanging on the right of the line, the way a totals block does.
    amounts.length && text.lastIndexOf(amounts.at(-1)) > text.length / 2 ? 1 : 0,
    letters && letters === letters.toUpperCase() ? 1 : 0,
  ];
}

/**
 * One binary classifier per label, each trained on every line.
 *
 * The weighting matters more than the optimiser: three lines out of thirty
 * carry a field, and an unweighted fit answers « other » to everything and is
 * right nine times in ten.
 */
function trainOne(rows, targets, label, epochs, rate) {
  const wanted = targets.map((t) => (t === label ? 1 : 0));
  const positives = wanted.reduce((a, b) => a + b, 0);
  const weights = new Float64Array(rows[0].length);
  let bias = 0;
  for (let epoch = 0; epoch < epochs; epoch += 1) {
    for (let i = 0; i < rows.length; i += 1) {
      const balance = rows.length / (2 * (wanted[i] ? positives : rows.length - positives));
      let z = bias;
      for (let j = 0; j < weights.length; j += 1) z += weights[j] * rows[i][j];
      const error = (1 / (1 + Math.exp(-z)) - wanted[i]) * balance;
      for (let j = 0; j < weights.length; j += 1) weights[j] -= rate * error * rows[i][j];
      bias -= rate * error;
    }
  }
  return { weights, bias };
}

/** `labels` carries one label per non-blank line of each document. */
export function train(documents, labels, { epochs = 400, rate = 0.3 } = {}) {
  const rows = [];
  const targets = [];
  documents.forEach((document, d) => {
    const lines = pageLines(document);
    lines.forEach((line, i) => rows.push(lineFeatures(line, i, lines.length)));
    targets.push(...labels[d]);
  });
  const classes = [...new Set(targets)].sort();
  return { classes, models: classes.map((c) => trainOne(rows, targets, c, epochs, rate)) };
}

function score(model, row) {
  let z = model.bias;
  for (let j = 0; j < row.length; j += 1) z += model.weights[j] * row[j];
  return 1 / (1 + Math.exp(-z));
}

/** Reads the first value of that shape out of a line, or nothing. */
const matching = (pattern) => (line) => line.match(pattern)?.[0] ?? null;

/** The rightmost amount: an item line carries a quantity and a unit price. */
function amount(line) {
  const amounts = line.match(AMOUNT);
  if (!amounts) return null;
  let cleaned = amounts.at(-1).replace(/[^\d,.]/g, '');
  // A comma means French spelling: the dots left are thousands separators.
  if (cleaned.includes(',')) cleaned = cleaned.replaceAll('.', '').replace(',', '.');
  return Number(cleaned);
}

const READERS = { invoice_number: matching(REFERENCE), date: matching(DATE), total: amount };

/**
 * For each field, walk the lines from the most likely down.
 *
 * The model points at a line; reading a date or an amount out of it is still
 * ours to do, and a line the model likes but that holds no value is not an
 * answer.
 */
export function extractFields(model, text) {
  const lines = pageLines(text);
  const rows = lines.map((line, i) => lineFeatures(line, i, lines.length));
  const fields = {};
  for (const [name, read] of Object.entries(READERS)) {
    const column = model.classes.indexOf(name);
    const ranked = rows
      .map((row, i) => [score(model.models[column], row), i])
      .sort((a, b) => b[0] - a[0]);
    fields[name] = ranked.map(([, i]) => read(lines[i])).find((value) => value !== null) ?? null;
  }
  return fields;
}

"""
Tag the lines of the invoice with a self-hosted document model.

Rung N2. Same idea as N1 — label the lines, then read the value out of them —
except that the features are no longer yours. A document encoder fine-tuned on
invoices reads the line, its neighbours and its place on the page at once, and
it keeps reading them when a supplier moves its totals block.

What you own on this rung is not the model, it is everything around it: the
page geometry you hand it, the threshold under which a field goes to a human,
and the answer to « what does the code do when the model says nothing usable ».
The weights stay on your machine, which is why an invoice never leaves it.
"""

from __future__ import annotations

import re

# A base encoder is a starting point, not an extractor: this rung assumes the
# checkpoint was fine-tuned on invoices, yours or someone else's.
MODEL_NAME = "microsoft/layoutlmv3-base"

DEFAULT_THRESHOLD = 0.75

# One pass reads one page. Beyond that the model would truncate in silence.
MAX_LINES = 120

AMOUNT = r"\d{1,3}(?:[\s.]\d{3})*[,.]\d{2}"
MONTHS = "janvier|février|mars|avril|mai|juin|juillet|août|septembre|octobre|novembre|décembre"
DATE = rf"\d{{1,2}}/\d{{1,2}}/\d{{2,4}}|\d{{1,2}}\s+(?:{MONTHS})\s+\d{{4}}"
REFERENCE = r"\b(?:[A-Za-z]{1,3}[-/])?\d[\dA-Za-z/-]{3,}"


class ExtractionUnavailable(Exception):
    """The model answered something no caller can act on."""


class LayoutModel:
    """The real model, loaded once and kept in memory for the process."""

    def __init__(self, name: str = MODEL_NAME) -> None:
        from transformers import pipeline  # a large download, done once

        self._pipe = pipeline("token-classification", model=name)

    def predict(self, lines: list[str]) -> list[dict[str, float]]:
        """One label-to-score mapping per line, in the order given."""
        rows = self._pipe({"words": lines, "boxes": boxes_for(lines)})
        return [{r["entity_group"]: r["score"] for r in row} for row in rows]


def boxes_for(lines: list[str]) -> list[list[int]]:
    """
    A box per line, on the thousandth-of-a-page grid these models expect.

    Extracted text keeps its geometry in two places only: how far a line is
    indented, and how far down the page it sits. That is what the model gets,
    and it is already more than a bag of words has.
    """
    height = max(len(lines), 1)
    return [
        [min(len(line) - len(line.lstrip()), 80) * 12, i * 1000 // height, 1000, (i + 1) * 1000 // height]
        for i, line in enumerate(lines)
    ]


def extract_fields(text: str, model=None, *, threshold: float = DEFAULT_THRESHOLD, attempts: int = 2) -> dict:
    """
    Read the fields, and say for each one whether a human should look.

    `model` is injected so this can be tested without downloading the weights.
    In production it defaults to the real model above.
    """
    model = model or LayoutModel()
    lines = [line for line in text.splitlines() if line.strip()]
    if len(lines) > MAX_LINES:
        raise ValueError(f"document longer than {MAX_LINES} lines")
    tagged = _tag(model, lines, attempts) if lines else []
    if len(tagged) != len(lines):
        raise ExtractionUnavailable("the model owed one row per line, and did not")
    return {field: _read(field, lines, tagged, threshold) for field in READERS}


def _tag(model, lines: list[str], attempts: int) -> list[dict]:
    """The whole page in one pass, and a failed pass is retried, not swallowed."""
    last_error: Exception | None = None
    for _ in range(attempts):
        try:
            return model.predict(lines)
        except Exception as error:  # noqa: BLE001 - any model failure is retried
            last_error = error
    raise ExtractionUnavailable(str(last_error))


def _read(field: str, lines: list[str], tagged: list[dict], threshold: float) -> dict:
    """
    Keep the best-scoring line for the field, then read the value out of it.

    The model points at a line; turning that line into a date or an amount is
    still ours, and so is deciding what happens when it cannot be done.
    """
    score, index = max(((_score(row, field), i) for i, row in enumerate(tagged)), default=(0.0, -1))
    if score <= 0.0:
        return {"value": None, "score": 0.0, "review": True}
    value = READERS[field](lines[index])
    # A doubtful field is not thrown away: it goes to a human with the score
    # that earned the doubt.
    return {"value": value, "score": score, "review": value is None or score < threshold}


def _score(row, field: str) -> float:
    """A number the caller can act on, rather than whatever came back."""
    value = (row or {}).get(field)
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return 0.0
    return float(value) if 0.0 <= value <= 1.0 else 0.0


def _match(pattern: str):
    def read(line: str):
        found = re.search(pattern, line)
        return found.group(0) if found else None

    return read


def _amount(line: str):
    """The rightmost amount: an item line carries a quantity and a unit price."""
    amounts = re.findall(AMOUNT, line)
    if not amounts:
        return None
    cleaned = re.sub(r"[^\d,.]", "", amounts[-1])
    if "," in cleaned:
        cleaned = cleaned.replace(".", "").replace(",", ".")
    return float(cleaned)


READERS = {"invoice_number": _match(REFERENCE), "date": _match(DATE), "total": _amount}

/**
 * Tag the lines of the invoice with a self-hosted document model.
 *
 * Rung N2. Same idea as N1 — label the lines, then read the value out of them
 * — except that the features are no longer yours. A document encoder
 * fine-tuned on invoices reads the line, its neighbours and its place on the
 * page at once, and it keeps reading them when a supplier moves its totals
 * block.
 *
 * What you own on this rung is not the model, it is everything around it: the
 * page geometry you hand it, the threshold under which a field goes to a
 * human, and the answer to « what does the code do when the model says nothing
 * usable ». The weights stay on your machine, which is why an invoice never
 * leaves it.
 */

// A base encoder is a starting point, not an extractor: this rung assumes the
// checkpoint was fine-tuned on invoices, yours or someone else's.
export const MODEL_NAME = 'Xenova/layoutlmv3-base';

export const DEFAULT_THRESHOLD = 0.75;

// One pass reads one page. Beyond that the model would truncate in silence.
export const MAX_LINES = 120;

const AMOUNT = /\d{1,3}(?:[\s.]\d{3})*[,.]\d{2}/g;
const MONTHS = 'janvier|février|mars|avril|mai|juin|juillet|août|septembre|octobre|novembre|décembre';
const DATE = new RegExp(`\\d{1,2}/\\d{1,2}/\\d{2,4}|\\d{1,2}\\s+(?:${MONTHS})\\s+\\d{4}`);
const REFERENCE = /\b(?:[A-Za-z]{1,3}[-/])?\d[\dA-Za-z/-]{3,}/;

export class ExtractionUnavailable extends Error {}

/** The real model, loaded once and kept in memory for the process. */
export class LayoutModel {
  static async load(name = MODEL_NAME) {
    const { pipeline } = await import('@huggingface/transformers'); // a large download, done once
    return new LayoutModel(await pipeline('token-classification', name));
  }

  constructor(pipe) {
    this.pipe = pipe;
  }

  /** One label-to-score mapping per line, in the order given. */
  async predict(lines) {
    const rows = await this.pipe({ words: lines, boxes: boxesFor(lines) });
    return rows.map((row) => Object.fromEntries(row.map((r) => [r.entity_group, r.score])));
  }
}

/**
 * A box per line, on the thousandth-of-a-page grid these models expect.
 *
 * Extracted text keeps its geometry in two places only: how far a line is
 * indented, and how far down the page it sits. That is what the model gets,
 * and it is already more than a bag of words has.
 */
export function boxesFor(lines) {
  const height = Math.max(lines.length, 1);
  return lines.map((line, i) => [
    Math.min(line.length - line.trimStart().length, 80) * 12,
    Math.floor((i * 1000) / height),
    1000,
    Math.floor(((i + 1) * 1000) / height),
  ]);
}

/**
 * Read the fields, and say for each one whether a human should look.
 *
 * `model` is injected so this can be tested without downloading the weights.
 * In production it defaults to the real model above.
 */
export async function extractFields(text, model, { threshold = DEFAULT_THRESHOLD, attempts = 2 } = {}) {
  const tagger = model ?? (await LayoutModel.load());
  const lines = text.split('\n').filter((line) => line.trim() !== '');
  if (lines.length > MAX_LINES) throw new RangeError(`document longer than ${MAX_LINES} lines`);
  const tagged = lines.length ? await tag(tagger, lines, attempts) : [];
  if (tagged.length !== lines.length) {
    throw new ExtractionUnavailable('the model owed one row per line, and did not');
  }
  return Object.fromEntries(Object.keys(READERS).map((f) => [f, read(f, lines, tagged, threshold)]));
}

/** The whole page in one pass, and a failed pass is retried, not swallowed. */
async function tag(model, lines, attempts) {
  let lastError;
  for (let i = 0; i < attempts; i += 1) {
    try {
      return await model.predict(lines);
    } catch (error) {
      lastError = error;
    }
  }
  throw new ExtractionUnavailable(String(lastError));
}

/**
 * Keep the best-scoring line for the field, then read the value out of it.
 *
 * The model points at a line; turning that line into a date or an amount is
 * still ours, and so is deciding what happens when it cannot be done.
 */
function read(field, lines, tagged, threshold) {
  let best = { score: 0, index: -1 };
  tagged.forEach((row, index) => {
    const value = scoreOf(row, field);
    if (value >= best.score) best = { score: value, index };
  });
  if (best.score <= 0) return { value: null, score: 0, review: true };
  const value = READERS[field](lines[best.index]);
  // A doubtful field is not thrown away: it goes to a human with the score
  // that earned the doubt.
  return { value, score: best.score, review: value === null || best.score < threshold };
}

/** A number the caller can act on, rather than whatever came back. */
function scoreOf(row, field) {
  const value = (row ?? {})[field];
  return typeof value === 'number' && value >= 0 && value <= 1 ? value : 0;
}

const matching = (pattern) => (line) => line.match(pattern)?.[0] ?? null;

/** The rightmost amount: an item line carries a quantity and a unit price. */
function amount(line) {
  const amounts = line.match(AMOUNT);
  if (!amounts) return null;
  let cleaned = amounts.at(-1).replace(/[^\d,.]/g, '');
  if (cleaned.includes(',')) cleaned = cleaned.replaceAll('.', '').replace(',', '.');
  return Number(cleaned);
}

const READERS = { invoice_number: matching(REFERENCE), date: matching(DATE), total: amount };

"""
Ask a general-purpose multimodal model to read the invoice.

Rung N3. This is the option people reach for first, and on this task it has a
real argument: the model is given a picture of the page, so it sees the column
an amount sits in, which the extracted text has already lost.

Note what the code has to do that N0 did not: cap the size of what it sends,
retry on failure, parse an answer that is only probably JSON, and check the
shape of what came back. That plumbing is the real cost of this rung, and it is
the part your tests have to cover, because the model itself is not testable.
"""

from __future__ import annotations

import base64
import json

PROMPT = (
    "Read the invoice below and return its header fields.\n"
    "Answer with JSON only: an object with the keys `invoice_number`, `date`\n"
    "and `total`. `total` is the amount due, taxes included, as a number.\n"
    "Use null for a field the page does not carry.\n\n"
    "Extracted text:\n{text}"
)

# A model charges by the token, and a scanned page is a lot of them. Refusing
# an oversized image is not an optimisation, it is a cost control.
MAX_IMAGE_BYTES = 4_000_000

FIELDS = ("invoice_number", "date", "total")


class ExtractionUnavailable(Exception):
    """The provider could not be reached, or answered something unusable."""


def extract_fields(text: str, page_image: bytes, client=None, *, attempts: int = 3) -> dict:
    """
    Read the invoice fields from its text and a picture of the page.

    `client` is injected so this function can be tested without a network call.
    In production it defaults to a real provider client.
    """
    if client is None:  # pragma: no cover - needs a key and a network
        from openai import OpenAI

        client = OpenAI()

    if len(page_image) > MAX_IMAGE_BYTES:
        raise ValueError(f"page image larger than {MAX_IMAGE_BYTES} bytes")

    image_url = "data:image/png;base64," + base64.b64encode(page_image).decode()
    return _decode(_ask(client, text, image_url, attempts))


def _ask(client, text: str, image_url: str, attempts: int) -> str:
    last_error: Exception | None = None
    for _ in range(attempts):
        try:
            return client.complete(
                prompt=PROMPT.format(text=text),
                image_url=image_url,
                # Temperature zero, because an amount that changes between two
                # identical calls cannot be reconciled with anything.
                temperature=0,
            )
        except Exception as error:  # noqa: BLE001 - any provider failure is retried
            last_error = error
    raise ExtractionUnavailable(str(last_error))


def _decode(answer: str) -> dict:
    """
    Turn the answer into the three fields, or refuse it.

    Models like to wrap JSON in a code fence. That is noise, not an error, and
    stripping it is cheaper than another call.
    """
    stripped = answer.strip().removeprefix("```json").removeprefix("```").removesuffix("```")
    try:
        parsed = json.loads(stripped)
    except ValueError as error:
        raise ExtractionUnavailable(f"the model did not answer with JSON: {error}")
    if not isinstance(parsed, dict):
        raise ExtractionUnavailable("the model answered something that is not an object")

    fields = {field: parsed.get(field) for field in FIELDS}
    total = fields["total"]
    # A total nobody can compute with is worse than no total at all: it would
    # travel down the pipeline looking like a number.
    if total is not None and (isinstance(total, bool) or not isinstance(total, (int, float))):
        raise ExtractionUnavailable(f"the model answered a total that is not a number: {total!r}")
    return fields

/**
 * Ask a general-purpose multimodal model to read the invoice.
 *
 * Rung N3. This is the option people reach for first, and on this task it has
 * a real argument: the model is given a picture of the page, so it sees the
 * column an amount sits in, which the extracted text has already lost.
 *
 * Note what the code has to do that N0 did not: cap the size of what it sends,
 * retry on failure, parse an answer that is only probably JSON, and check the
 * shape of what came back. That plumbing is the real cost of this rung, and it
 * is the part your tests have to cover, because the model itself is not
 * testable.
 */

const PROMPT = [
  'Read the invoice below and return its header fields.',
  'Answer with JSON only: an object with the keys `invoice_number`, `date`',
  'and `total`. `total` is the amount due, taxes included, as a number.',
  'Use null for a field the page does not carry.',
  '',
  'Extracted text:',
].join('\n');

// A model charges by the token, and a scanned page is a lot of them. Refusing
// an oversized image is not an optimisation, it is a cost control.
export const MAX_IMAGE_BYTES = 4_000_000;

const FIELDS = ['invoice_number', 'date', 'total'];

export class ExtractionUnavailable extends Error {}

/**
 * Read the invoice fields from its text and a picture of the page.
 *
 * @param {string} text the invoice, already extracted
 * @param {Uint8Array} pageImage the rendered page; this snippet opens no file
 * @param {object} options
 * @param {{complete: Function}} [options.client] injected so this can be
 *   tested without a network call; defaults to a real provider client
 * @param {number} [options.attempts]
 */
export async function extractFields(text, pageImage, { client, attempts = 3 } = {}) {
  let provider = client;
  if (!provider) {
    // Needs a key and a network, so it is never reached in the tests.
    const { OpenAI } = await import('openai');
    provider = new OpenAI();
  }

  if (pageImage.length > MAX_IMAGE_BYTES) {
    throw new RangeError(`page image larger than ${MAX_IMAGE_BYTES} bytes`);
  }

  const imageUrl = `data:image/png;base64,${Buffer.from(pageImage).toString('base64')}`;
  return decode(await ask(provider, text, imageUrl, attempts));
}

async function ask(client, text, imageUrl, attempts) {
  let lastError;
  for (let i = 0; i < attempts; i += 1) {
    try {
      return await client.complete({
        prompt: `${PROMPT}\n${text}`,
        imageUrl,
        // Temperature zero, because an amount that changes between two
        // identical calls cannot be reconciled with anything.
        temperature: 0,
      });
    } catch (error) {
      lastError = error;
    }
  }
  throw new ExtractionUnavailable(String(lastError));
}

/**
 * Turn the answer into the three fields, or refuse it.
 *
 * Models like to wrap JSON in a code fence. That is noise, not an error, and
 * stripping it is cheaper than another call.
 */
function decode(answer) {
  const stripped = answer.trim().replace(/^```(?:json)?/, '').replace(/```$/, '');
  let parsed;
  try {
    parsed = JSON.parse(stripped);
  } catch (error) {
    throw new ExtractionUnavailable(`the model did not answer with JSON: ${error.message}`);
  }
  if (parsed === null || typeof parsed !== 'object' || Array.isArray(parsed)) {
    throw new ExtractionUnavailable('the model answered something that is not an object');
  }

  const fields = Object.fromEntries(FIELDS.map((field) => [field, parsed[field] ?? null]));
  // A total nobody can compute with is worse than no total at all: it would
  // travel down the pipeline looking like a number.
  if (fields.total !== null && typeof fields.total !== 'number') {
    throw new ExtractionUnavailable(`the model answered a total that is not a number: ${fields.total}`);
  }
  return fields;
}