sea-dev 1.0.0

package/libs/core/src/data-extraction/structuring/bank-statement.ts

@@ -0,0 +1,1182 @@
+import { google } from "@ai-sdk/google";
+import { formatISODate } from "@sea/util/fmt";
+import { generateText, Output } from "ai";
+import {
+  getMonth,
+  getYear,
+  isValid,
+  isWithinInterval,
+  parse,
+  setYear,
+  startOfDay,
+  endOfDay,
+} from "date-fns";
+import { z } from "zod/v4";
+import type {
+  ValueWithReference,
+  StructuredGroup,
+  ValidationData,
+  StructuredData,
+  BankStatementReconciliation,
+} from "../types";
+import type {
+  PretrainedEntity,
+  DocumentLayout,
+  ReferenceMap,
+  LayoutItemType,
+} from "@sea/schemas/core/layout";
+import { Table, Section } from "@sea/schemas/core/layout";
+
+const DATE_FORMATS: string[] = [
+  // Full Dates with 2-Digit Year (checked first to avoid ambiguity)
+  "dd MM yy",
+  "dd MMM yy",
+  "dd MMMM yy",
+  "d MM yy",
+  "d MMM yy",
+  "d MMMM yy",
+  "yy MM dd",
+  "yy MMM dd",
+  "yy MMMM dd",
+  "yy MM d",
+  "yy MMM d",
+  "yy MMMM d",
+  "MM dd yy",
+  "MMM dd yy",
+  "MMMM dd yy",
+  "MM d yy",
+  "MMM d yy",
+  "MMMM d yy",
+  // Full Dates with 4-Digit Year
+  "dd MM yyyy",
+  "dd MMM yyyy",
+  "dd MMMM yyyy",
+  "d MM yyyy",
+  "d MMM yyyy",
+  "d MMMM yyyy",
+  "yyyy MM dd",
+  "yyyy MMM dd",
+  "yyyy MMMM dd",
+  "yyyy MM d",
+  "yyyy MMM d",
+  "yyyy MMMM d",
+  "MM dd yyyy",
+  "MMM dd yyyy",
+  "MMMM dd yyyy",
+  "MM d yyyy",
+  "MMM d yyyy",
+  "MMMM d yyyy",
+  // Partial Dates with no Year
+  "dd MM",
+  "d M",
+  "dd MMM",
+  "d MMMM",
+  "MMM dd",
+  "MMMM dd",
+  "MMMM d",
+];
+
+const BANK_TYPES: readonly string[] = ["barclays", "metro", "other"];
+
+export const BankStatementBaseSchema = z.object({
+  bankName: z
+    .enum(BANK_TYPES)
+    .describe(
+      "Bank name from the available list. If the bank name is not in the list, return 'other'.",
+    ),
+  transactionDateFormat: z.enum(DATE_FORMATS).describe(
+    `Identify the exact date-fns format string used for transaction dates within the provided bank statement table.
+
+CRITICAL RULES:
+
+- Year Check: Look at the transaction dates. If there is NO year digit present, you are strictly forbidden from using yyyy or yy in the output.
+- Separator Check: You MUST replace all slashes, dashes, dots, or commas with a single space.
+- Context: Only use the format found in the transaction table rows.
+
+Step-by-Step Logic:
+
+Step 1: Does the date string contain a year (e.g., 2024 or 24)?
+Step 2: If NO, the format must only contain day and month tokens.
+Step 3: Convert the pattern to date-fns tokens using ONLY spaces.
+
+Examples:
+
+"31 Oct" -> "d MMM" (Correct: No year present)
+"31 Oct 2023" -> "d MMM yyyy" (Correct: Year present)
+"31/10" -> "dd MM" (Correct: No year present)
+
+Final Output: Return ONLY the format string. No extra words or years.
+`,
+  ),
+  statementPeriodStart: z.iso
+    .date()
+    .describe("Bank statement period start date in YYYY-MM-DD format"),
+  statementPeriodEnd: z.iso
+    .date()
+    .describe("Bank statement period end date in YYYY-MM-DD format"),
+  openingBalance: z
+    .number()
+    .nullable()
+    .describe(
+      "Opening balance for the period. If value contains D, it means the balance is negative.",
+    ),
+  closingBalance: z
+    .number()
+    .nullable()
+    .describe(
+      "Closing balance for the period. If value contains D, it means the balance is negative.",
+    ),
+  isReverseOrdered: z
+    .boolean()
+    .default(false)
+    .describe(
+      "Whether the transactions are reverse ordered, namely if the most recent transactions are at the top of the list. Default to false",
+    ),
+});
+
+/**
+ * Structure bank statement data.
+ *
+ * @param layout - Parsed document layout from the layout extractor
+ * @param getPdfBuffer - Optional lazy getter for the raw PDF bytes; required for VLM healing (step 3)
+ */
+export async function structureBankStatementData(
+  layout: DocumentLayout,
+): Promise<{
+  data: StructuredData;
+  validationData?: ValidationData;
+  reconciliationData?: BankStatementReconciliation;
+}> {
+  try {
+    const baseData = await getPretrainedBaseData(layout);
+    let { data, validationData } = parseAndValidateOutput(layout, baseData);
+
+    // Self-healing pipeline: attempt to fix transactions that failed balance validation
+    if (
+      validationData?.Transactions &&
+      validationData.Transactions.length > 0
+    ) {
+      const transactions = data as Record<string, ValueWithReference>[];
+      validationData = await runSelfHealingPipeline(
+        transactions,
+        validationData.Transactions,
+        layout,
+        baseData,
+      );
+      data = transactions as unknown as StructuredGroup;
+    } else {
+      console.info("All transaction balances validated successfully");
+    }
+
+    // Calculate reconciliation data
+    let totalCredits = 0;
+    let totalDebits = 0;
+
+    if (Array.isArray(data)) {
+      for (const transaction of data) {
+        const depositAmount =
+          (transaction.DepositAmount?.value as number | undefined) ?? 0;
+        const withdrawalAmount =
+          (transaction.WithdrawalAmount?.value as number | undefined) ?? 0;
+        totalCredits += depositAmount;
+        totalDebits += withdrawalAmount;
+      }
+    }
+
+    const openingBalance = baseData.openingBalance;
+    const closingBalance = baseData.closingBalance;
+    const calculatedClosingBalance =
+      (openingBalance ?? 0) + totalCredits - totalDebits;
+    const balanceDifference =
+      (closingBalance ?? calculatedClosingBalance) - calculatedClosingBalance;
+    const isReconciled = Math.abs(balanceDifference) < 0.01;
+
+    const reconciliationData: BankStatementReconciliation = {
+      statementPeriodStart: baseData.statementPeriodStart,
+      statementPeriodEnd: baseData.statementPeriodEnd,
+      openingBalance,
+      closingBalance,
+      totalCredits,
+      totalDebits,
+      calculatedClosingBalance,
+      balanceDifference,
+      isReconciled,
+    };
+
+    return {
+      data: { Transactions: data },
+      validationData,
+      reconciliationData,
+    };
+  } catch (error) {
+    console.error("Error structuring and storing bank statement data:", error);
+    throw new Error(
+      `Failed to structure and store bank statement data: ${error}`,
+    );
+  }
+}
+
+// ============================================================================
+// Self-healing helpers
+// ============================================================================
+
+/** A contiguous block of transaction indices to re-extract together. */
+type HealingBlock = { start: number; end: number };
+
+/** Schema for a single re-extracted transaction row. */
+const HealedTransactionSchema = z.object({
+  date: z.iso
+    .date()
+    .nullable()
+    .describe("Transaction date in YYYY-MM-DD format, or null if not found"),
+  description: z
+    .string()
+    .nullable()
+    .describe("Transaction description/narrative"),
+  type: z
+    .string()
+    .nullable()
+    .describe("Transaction type code, or null if absent"),
+  depositAmount: z
+    .number()
+    .nullable()
+    .describe("Credit/deposit amount (positive number), or null"),
+  withdrawalAmount: z
+    .number()
+    .nullable()
+    .describe("Debit/withdrawal amount (positive number), or null"),
+  balance: z
+    .number()
+    .nullable()
+    .describe("Running balance after this transaction, or null if not present"),
+});
+
+const HealedTransactionListSchema = z.object({
+  transactions: z
+    .array(HealedTransactionSchema)
+    .describe("Re-extracted transactions in the same order as the source rows"),
+});
+
+type HealedTransaction = z.infer<typeof HealedTransactionSchema>;
+
+// ─── Pipeline orchestration ──────────────────────────────────────────────────
+
+/**
+ * Logs a set of validation errors with their current transaction values.
+ */
+function logFailures(
+  label: string,
+  failures: { index: number; message: string }[],
+  transactions: Record<string, ValueWithReference>[],
+): void {
+  console.warn(`[Self-heal] ── ${label} (${failures.length} failing) ──`);
+  for (const { index, message } of failures) {
+    console.warn(
+      `[Self-heal]   [${index}] ${message} | ${formatTransactionSummary(index, transactions[index]!)}`,
+    );
+  }
+}
+
+/**
+ * Logs the contents of a set of healing blocks.
+ */
+function logBlocks(
+  label: string,
+  blocks: HealingBlock[],
+  transactions: Record<string, ValueWithReference>[],
+): void {
+  console.info(
+    `[Self-heal] ${label}: ${blocks.map((b) => `[${b.start}–${b.end}]`).join(", ")}`,
+  );
+  for (const block of blocks) {
+    console.info(
+      `[Self-heal]   Block [${block.start}–${block.end}] (${block.end - block.start + 1} rows):`,
+    );
+    for (let i = block.start; i <= block.end; i++) {
+      console.info(
+        `[Self-heal]     ${formatTransactionSummary(i, transactions[i]!)}`,
+      );
+    }
+  }
+}
+
+/**
+ * Runs the full self-healing pipeline on a mutable transactions array.
+ *
+ * Pass 1 — text:  builds blocks → calls Gemini with raw layout text → revalidates.
+ * Pass 2 — VLM:   builds blocks from remaining failures → calls Gemini with
+ *                  the relevant PDF page → final revalidation.
+ *
+ * Mutates `transactions` in-place.
+ *
+ * @returns Updated ValidationData (undefined if fully reconciled, populated if
+ *          some failures remain after both passes).
+ */
+async function runSelfHealingPipeline(
+  transactions: Record<string, ValueWithReference>[],
+  initialFailures: { index: number; message: string }[],
+  layout: DocumentLayout,
+  baseData: z.infer<typeof BankStatementBaseSchema>,
+): Promise<ValidationData | undefined> {
+  const MAX_HEALING_BLOCKS = 10;
+
+  console.warn(
+    `[Self-heal] ── PIPELINE START ─────────────────────────────────────`,
+  );
+  console.warn(
+    `[Self-heal] ${initialFailures.length} failing transaction(s) out of ${transactions.length} total`,
+  );
+  logFailures("Initial failures", initialFailures, transactions);
+
+  // ── Step 0: deposit/withdrawal flip pre-pass ─────────────────────────────
+  const afterFlip = tryFlipDepositWithdrawal(
+    transactions,
+    initialFailures,
+    baseData.openingBalance ?? 0,
+  );
+
+  if (afterFlip.length === 0) {
+    console.info(`[Self-heal] ── STEP 0 RESULT: all reconciled by flip ✓`);
+    console.info(
+      `[Self-heal] ── PIPELINE END ────────────────────────────────`,
+    );
+    return undefined;
+  }
+
+  logFailures("After flip — still failing", afterFlip, transactions);
+
+  // ── Step 1: text re-extraction ──────────────────────────────────────────
+  console.info(`[Self-heal] ── STEP 1: text re-extraction ──────────────────`);
+  const textBlocks = buildHealingBlocks(afterFlip, transactions.length);
+
+  if (textBlocks.length > MAX_HEALING_BLOCKS) {
+    console.warn(
+      `[Self-heal] ${textBlocks.length} blocks exceeds limit of ${MAX_HEALING_BLOCKS} — skipping self-heal`,
+    );
+    console.info(
+      `[Self-heal] ── PIPELINE END ────────────────────────────────`,
+    );
+    return { Transactions: initialFailures };
+  }
+  logBlocks(`${textBlocks.length} text block(s)`, textBlocks, transactions);
+
+  await reextractBlocksFromText(transactions, textBlocks, layout, baseData);
+
+  // ── Step 2: revalidate ──────────────────────────────────────────────────
+  const afterText = validateTransactionBalances(
+    transactions,
+    baseData.openingBalance ?? 0,
+  );
+
+  if (afterText.Transactions.length === 0) {
+    console.info(`[Self-heal] ── STEP 1 RESULT: all reconciled ✓`);
+    console.info(
+      `[Self-heal] ── PIPELINE END ────────────────────────────────`,
+    );
+    return undefined;
+  }
+
+  logFailures(
+    "STEP 1 RESULT — still failing",
+    afterText.Transactions,
+    transactions,
+  );
+
+  // ── Step 3: VLM re-extraction (disabled — mupdf/WASM dependency removed) ─
+  // console.info(`[Self-heal] ── STEP 3: VLM re-extraction ───────────────────`);
+  // const vlmBlocks = buildHealingBlocks(
+  //   afterText.Transactions,
+  //   transactions.length,
+  // );
+  // logBlocks(`${vlmBlocks.length} VLM block(s)`, vlmBlocks, transactions);
+  // if (getPdfBuffer) {
+  //   const pdfBuffer = await getPdfBuffer();
+  //   await reextractBlocksFromVLM(
+  //     transactions,
+  //     vlmBlocks,
+  //     layout.referenceMap ?? {},
+  //     pdfBuffer,
+  //     baseData,
+  //   );
+  // } else {
+  //   console.warn("[Self-heal] VLM skipped — no PDF buffer getter provided");
+  // }
+  // const afterVlm = validateTransactionBalances(
+  //   transactions,
+  //   baseData.openingBalance ?? 0,
+  // );
+  // if (afterVlm.Transactions.length === 0) {
+  //   console.info(`[Self-heal] ── STEP 3 RESULT: all reconciled ✓`);
+  //   console.info(`[Self-heal] ── PIPELINE END ────────────────────────────────`);
+  //   return undefined;
+  // }
+  // logFailures("STEP 3 RESULT — giving up", afterVlm.Transactions, transactions);
+
+  logFailures(
+    "STEP 1 RESULT — giving up",
+    afterText.Transactions,
+    transactions,
+  );
+  console.info(`[Self-heal] ── PIPELINE END ────────────────────────────────`);
+  return { Transactions: afterText.Transactions };
+}
+
+// ─── Deposit/withdrawal flip pre-pass ────────────────────────────────────────
+
+/**
+ * Clones a single transaction's DepositAmount and WithdrawalAmount, swapped.
+ */
+function flipTx(
+  tx: Record<string, ValueWithReference>,
+): Record<string, ValueWithReference> {
+  const dep = tx.DepositAmount ?? null;
+  const wdl = tx.WithdrawalAmount ?? null;
+  return {
+    ...tx,
+    DepositAmount: wdl ? { ...wdl } : null,
+    WithdrawalAmount: dep ? { ...dep } : null,
+  };
+}
+
+/**
+ * Returns true if tx at `index` reconciles given the balance of the preceding
+ * transaction (or openingBalance if index === 0).
+ */
+function txReconciles(
+  transactions: Record<string, ValueWithReference>[],
+  index: number,
+  openingBalance: number,
+): boolean {
+  const tx = transactions[index];
+  if (!tx) return false;
+  const prevBalance =
+    index === 0
+      ? openingBalance
+      : ((transactions[index - 1]?.Balance?.value ?? openingBalance) as number);
+  const balance = tx.Balance?.value as number | undefined;
+  if (balance === undefined || balance === null) return true; // null balance not our problem
+  const dep = (tx.DepositAmount?.value ?? 0) as number;
+  const wdl = (tx.WithdrawalAmount?.value ?? 0) as number;
+  return Math.abs(balance - (prevBalance + dep - wdl)) <= 0.01;
+}
+
+/**
+ * Pre-pass Step 0: for each failing transaction, try swapping its
+ * DepositAmount and WithdrawalAmount (and/or the previous transaction's)
+ * and check whether the pair(s) in question now reconcile locally.
+ *
+ * Candidates tried per failing index (in order, stopping at first fix):
+ *   1. Flip just the failing tx.
+ *   2. Flip the failing tx AND the previous tx.
+ *
+ * The check is local: does flipping make index `i` (and `i-1` if included)
+ * reconcile against their immediate neighbours? No full-list revalidation needed.
+ *
+ * Mutates `transactions` in-place for every individual fix found.
+ * Returns the remaining failures after all individual fixes have been applied.
+ */
+function tryFlipDepositWithdrawal(
+  transactions: Record<string, ValueWithReference>[],
+  failures: { index: number; message: string }[],
+  openingBalance: number,
+): { index: number; message: string }[] {
+  console.info(
+    `[Self-heal flip] ── STEP 0: deposit/withdrawal flip pre-pass ──────`,
+  );
+
+  const stillFailing: { index: number; message: string }[] = [];
+
+  for (const failure of failures) {
+    const { index } = failure;
+
+    // Candidates: flip just this tx, or this tx + previous
+    const candidateSets: number[][] = [
+      [index],
+      ...(index > 0 ? [[index, index - 1]] : []),
+    ];
+
+    let fixed = false;
+    for (const candidates of candidateSets) {
+      // Apply flip to scratch copies of the candidate rows
+      const saved = candidates.map((i) => transactions[i]!);
+      for (const i of candidates) {
+        transactions[i] = flipTx(transactions[i]!);
+      }
+
+      // Check local reconciliation for each candidate index
+      const allReconcile = candidates.every((i) =>
+        txReconciles(transactions, i, openingBalance),
+      );
+
+      if (allReconcile) {
+        console.info(
+          `[Self-heal flip] Flipping [${candidates.join(", ")}] fixes index ${index} locally ✓ — applying`,
+        );
+        fixed = true;
+        break;
+      }
+
+      // Restore originals — this candidate set didn't work
+      for (let k = 0; k < candidates.length; k++) {
+        transactions[candidates[k]!] = saved[k]!;
+      }
+    }
+
+    if (!fixed) {
+      console.info(
+        `[Self-heal flip] Flip did not fix index ${index} — leaving for LLM`,
+      );
+      stillFailing.push(failure);
+    }
+  }
+
+  console.info(
+    `[Self-heal flip] ${stillFailing.length} failure(s) remain after flip pre-pass`,
+  );
+  return stillFailing;
+}
+
+// ─── Block computation ───────────────────────────────────────────────────────
+
+/**
+ * Takes a list of validation errors and builds contiguous index blocks to heal.
+ *
+ * For each failing transaction we also pull in the one before it (to give the
+ * model context about the previous balance), then merge all overlapping or
+ * adjacent ranges into non-overlapping contiguous blocks.
+ *
+ * Example: failures at indices 2 and 4, totalCount = 7
+ *   → raw ranges [1,2] and [3,4]
+ *   → merged: [1,4]  (adjacent, so they collapse into one block)
+ */
+function buildHealingBlocks(
+  failures: { index: number }[],
+  totalCount: number,
+): HealingBlock[] {
+  if (failures.length === 0) return [];
+
+  // Expand each failing index: include the transaction before it for context
+  const ranges: HealingBlock[] = failures.map(({ index }) => ({
+    start: Math.max(0, index - 1),
+    end: Math.min(totalCount - 1, index),
+  }));
+
+  console.info(
+    `[Self-heal blocks] Raw ranges before merge: ${ranges.map((r) => `[${r.start}–${r.end}]`).join(", ")}`,
+  );
+
+  // Sort by start index so we can merge in a single pass
+  ranges.sort((a, b) => a.start - b.start);
+
+  const merged: HealingBlock[] = [];
+  let current = ranges[0]!;
+
+  for (let i = 1; i < ranges.length; i++) {
+    const next = ranges[i]!;
+    // Blocks overlap or are directly adjacent → merge
+    if (next.start <= current.end + 1) {
+      const merged_end = Math.max(current.end, next.end);
+      console.info(
+        `[Self-heal blocks] Merging [${current.start}–${current.end}] + [${next.start}–${next.end}] → [${current.start}–${merged_end}]`,
+      );
+      current = { start: current.start, end: merged_end };
+    } else {
+      merged.push(current);
+      current = next;
+    }
+  }
+  merged.push(current);
+
+  console.info(
+    `[Self-heal blocks] Final merged blocks (${merged.length}): ${merged.map((b) => `[${b.start}–${b.end}]`).join(", ")}`,
+  );
+
+  return merged;
+}
+
+// ─── Layout context ──────────────────────────────────────────────────────────
+
+const HEALED_SENTINELS = ["system", "none", "healed-text", "healed-vlm"];
+
+function getTxRefIds(tx: Record<string, ValueWithReference>): string[] {
+  return Object.values(tx)
+    .map((v) => v?.refId)
+    .filter(
+      (id): id is string =>
+        typeof id === "string" && !HEALED_SENTINELS.includes(id),
+    );
+}
+
+/**
+ * Serialises the raw layout text spanning the full contiguous block.
+ * Anchors on the first ref-text of the first transaction and the last
+ * ref-text of the last transaction, returning the lines between them
+ * with a small surrounding margin.
+ */
+function getBlockLayoutContext(
+  block: HealingBlock,
+  transactions: Record<string, ValueWithReference>[],
+  layout: DocumentLayout,
+): string {
+  const refMap = layout.referenceMap ?? {};
+  const layoutText = layout.serializeAsText();
+
+  const firstTx = transactions[block.start];
+  const lastTx = transactions[block.end];
+  if (!firstTx || !lastTx) return layoutText.slice(0, 6000);
+
+  const firstRefText =
+    getTxRefIds(firstTx)
+      .map((id) => refMap[id]?.text)
+      .find(Boolean) ?? "";
+
+  const lastRefText =
+    getTxRefIds(lastTx)
+      .map((id) => refMap[id]?.text)
+      .find(Boolean) ?? "";
+
+  const pageNumber = getTxRefIds(firstTx)
+    .map((id) => refMap[id]?.boundingRegions?.[0]?.pageNumber)
+    .find((n): n is number => n !== undefined);
+
+  const pagePrefix = pageNumber !== undefined ? `[Page ${pageNumber}]\n` : "";
+
+  const startIdx = firstRefText ? layoutText.indexOf(firstRefText) : -1;
+  if (startIdx === -1) return `${pagePrefix}${layoutText.slice(0, 6000)}`;
+
+  const endSearchFrom = startIdx + firstRefText.length;
+  const endIdx =
+    lastRefText && lastRefText !== firstRefText
+      ? layoutText.indexOf(lastRefText, endSearchFrom)
+      : startIdx;
+
+  const lines = layoutText.split("\n");
+  const startLine = Math.max(
+    0,
+    layoutText.slice(0, startIdx).split("\n").length - 2,
+  );
+  const endLine = Math.min(
+    lines.length - 1,
+    layoutText.slice(0, endIdx === -1 ? startIdx : endIdx).split("\n").length +
+      1,
+  );
+
+  return `${pagePrefix}${lines.slice(startLine, endLine + 1).join("\n")}`;
+}
+
+// ─── Apply healed rows ───────────────────────────────────────────────────────
+
+function formatTransactionSummary(
+  index: number,
+  tx: Record<string, ValueWithReference>,
+): string {
+  return (
+    `  [${index}] Date=${tx.Date?.value ?? "?"} ` +
+    `Dep=${tx.DepositAmount?.value ?? "?"} ` +
+    `Wdl=${tx.WithdrawalAmount?.value ?? "?"} ` +
+    `Bal=${tx.Balance?.value ?? "?"} ` +
+    `Desc="${tx.Description?.value ?? "?"}"`
+  );
+}
+
+/**
+ * Splices healed rows back into `transactions` for the given block.
+ * Iterates back-to-front within the block so earlier indices stay stable
+ * if the caller later processes blocks in reverse order.
+ */
+function spliceHealedBlock(
+  transactions: Record<string, ValueWithReference>[],
+  block: HealingBlock,
+  healed: HealedTransaction[],
+  refSuffix: string,
+): void {
+  const blockSize = block.end - block.start + 1;
+
+  if (healed.length !== blockSize) {
+    console.warn(
+      `[Self-heal] Block [${block.start}–${block.end}]: expected ${blockSize} rows, got ${healed.length} — applying field-level merge for available rows`,
+    );
+  }
+
+  // Apply from back to front within the block
+  for (let i = Math.min(healed.length, blockSize) - 1; i >= 0; i--) {
+    const txIndex = block.start + i;
+    const tx = transactions[txIndex];
+    const row = healed[i];
+    if (!tx || !row) continue;
+
+    console.info(
+      `[Self-heal splice] [${txIndex}] BEFORE: ${formatTransactionSummary(txIndex, tx)}`,
+    );
+    console.info(
+      `[Self-heal splice] [${txIndex}] MODEL OUTPUT: date=${row.date ?? "null"} dep=${row.depositAmount ?? "null"} wdl=${row.withdrawalAmount ?? "null"} bal=${row.balance ?? "null"} desc="${row.description ?? "null"}"`,
+    );
+
+    if (row.date !== null)
+      tx.Date = { value: row.date, refId: tx.Date?.refId ?? refSuffix };
+    if (row.description !== null)
+      tx.Description = {
+        value: row.description,
+        refId: tx.Description?.refId ?? refSuffix,
+      };
+    if (row.type !== null)
+      tx.Type = { value: row.type, refId: tx.Type?.refId ?? refSuffix };
+    if (row.depositAmount !== null)
+      tx.DepositAmount = {
+        value: Math.abs(row.depositAmount),
+        refId: tx.DepositAmount?.refId ?? refSuffix,
+      };
+    if (row.withdrawalAmount !== null)
+      tx.WithdrawalAmount = {
+        value: Math.abs(row.withdrawalAmount),
+        refId: tx.WithdrawalAmount?.refId ?? refSuffix,
+      };
+    if (row.balance !== null)
+      tx.Balance = {
+        value: row.balance,
+        refId: tx.Balance?.refId ?? refSuffix,
+      };
+
+    console.info(
+      `[Self-heal splice] [${txIndex}]  AFTER: ${formatTransactionSummary(txIndex, tx)}`,
+    );
+  }
+}
+
+// ─── Text-based healing ──────────────────────────────────────────────────────
+
+/**
+ * Walks the document layout (recursing into Sections) and returns all Table
+ * instances whose body cells include any of the given ref IDs.
+ */
+function findTablesContainingRefs(
+  items: LayoutItemType[],
+  refIds: Set<string>,
+): Table[] {
+  const found: Table[] = [];
+  for (const item of items) {
+    if (item instanceof Table) {
+      const hasRef = item.content.some((row) =>
+        row.cells.some((cell) => refIds.has(cell.id)),
+      );
+      if (hasRef) found.push(item);
+    } else if (item instanceof Section) {
+      found.push(...findTablesContainingRefs(item.content, refIds));
+    }
+  }
+  return found;
+}
+
+/**
+ * Returns the header row(s) of the transactions table as a comma-separated
+ * string, or null if no table / no header rows are found.
+ */
+function getTableHeaderForBlock(
+  block: HealingBlock,
+  transactions: Record<string, ValueWithReference>[],
+  layout: DocumentLayout,
+): string | null {
+  const firstTx = transactions[block.start];
+  if (!firstTx) return null;
+
+  const refIds = new Set(getTxRefIds(firstTx));
+  if (refIds.size === 0) return null;
+
+  const tables = findTablesContainingRefs(layout.documentLayout, refIds);
+  console.log("Tables found:", tables);
+  if (tables.length === 0) return null;
+
+  const headerRows = tables[0]!.content.filter((r) => r.type === "head");
+  if (headerRows.length === 0) return tables[0]!.content[0]!.serializeAsText();
+
+  return headerRows.map((r) => r.serializeAsText()).join("\n");
+}
+
+/**
+ * Step 1: Re-extract each contiguous block of failing transactions from the
+ * raw layout text using Gemini.  All blocks are processed in parallel; within
+ * each block all rows are extracted in a single LLM call.
+ *
+ * Blocks are spliced back in reverse order (last block first) so indices
+ * remain stable throughout.  Mutates `transactions` in-place.
+ */
+async function reextractBlocksFromText(
+  transactions: Record<string, ValueWithReference>[],
+  blocks: HealingBlock[],
+  layout: DocumentLayout,
+  baseData: z.infer<typeof BankStatementBaseSchema>,
+): Promise<void> {
+  const period = `${baseData.statementPeriodStart} to ${baseData.statementPeriodEnd}`;
+
+  // Heal all blocks in parallel, collecting results keyed by block
+  const results = await Promise.all(
+    blocks.map(async (block) => {
+      const blockSize = block.end - block.start + 1;
+      const context = getBlockLayoutContext(block, transactions, layout);
+      const tableHeader = getTableHeaderForBlock(block, transactions, layout);
+      const currentRows = Array.from({ length: blockSize }, (_, i) =>
+        formatTransactionSummary(
+          block.start + i,
+          transactions[block.start + i]!,
+        ),
+      ).join("\n");
+
+      const headerSection = tableHeader
+        ? `\nTable column headers:\n${tableHeader}\n`
+        : "";
+
+      const prompt = `You are correcting ${blockSize} bank statement transaction row(s) (indices ${block.start}–${block.end}) that failed balance reconciliation.
+      Don't forget to check if deposit and withdrawal amounts are flipped.
+
+Statement period: ${period}
+Opening balance: ${baseData.openingBalance ?? "unknown"}
+Date format: ${baseData.transactionDateFormat}
+${headerSection}
+Currently extracted values (may contain errors):
+${currentRows}
+
+Raw document text for this block:
+${context}
+
+Re-extract all ${blockSize} transaction row(s) in order from the raw text. Return exactly ${blockSize} items in the "transactions" array.`;
+
+      console.info(
+        `[Self-heal text] Block [${block.start}–${block.end}] — table header: ${tableHeader ?? "(none)"}`,
+      );
+      console.info(
+        `[Self-heal text] Block [${block.start}–${block.end}] — sending prompt:\n${"─".repeat(60)}\n${prompt}\n${"─".repeat(60)}`,
+      );
+
+      try {
+        const result = await generateText({
+          model: google("gemini-3-pro-preview"),
+          system:
+            "You are a precise financial data extractor. Re-extract the given bank statement rows in order. Return null for any field absent in the source.",
+          prompt,
+          output: Output.object({ schema: HealedTransactionListSchema }),
+          maxRetries: 2,
+          temperature: 0,
+        });
+        const healed = HealedTransactionListSchema.parse(
+          result.output,
+        ).transactions;
+        console.info(
+          `[Self-heal text] Block [${block.start}–${block.end}] — model returned ${healed.length} row(s):`,
+        );
+        for (const [i, row] of healed.entries()) {
+          console.info(
+            `[Self-heal text]   [${block.start + i}] date=${row.date ?? "null"} dep=${row.depositAmount ?? "null"} wdl=${row.withdrawalAmount ?? "null"} bal=${row.balance ?? "null"} desc="${row.description ?? "null"}"`,
+          );
+        }
+        return { block, healed };
+      } catch (err) {
+        console.error(
+          `[Self-heal text] Failed for block [${block.start}–${block.end}]:`,
+          err,
+        );
+        return null;
+      }
+    }),
+  );
+
+  // Splice back in reverse block order (last block first) to preserve indices
+  const sorted = results
+    .filter((r): r is NonNullable<typeof r> => r !== null)
+    .sort((a, b) => b.block.start - a.block.start);
+
+  for (const { block, healed } of sorted) {
+    console.info(
+      `[Self-heal text] Splicing block [${block.start}–${block.end}] (processing back-to-front)`,
+    );
+    spliceHealedBlock(transactions, block, healed, "healed-text");
+  }
+}
+
+function parseAndValidateOutput(
+  layout: DocumentLayout,
+  baseData: z.infer<typeof BankStatementBaseSchema>,
+): {
+  data: StructuredGroup;
+  validationData?: ValidationData;
+} {
+  const transactions: StructuredGroup = [];
+  if (!layout?.entities) {
+    return {
+      data: transactions,
+    };
+  }
+
+  // Get Transactions entity from the record
+  const transactionsEntity = layout.entities["Transactions"];
+  if (!transactionsEntity) {
+    return {
+      data: transactions,
+    };
+  }
+
+  // Process transaction array
+  const transactionsList = transactionsEntity.valueArray;
+  if (!transactionsList) {
+    return {
+      data: transactions,
+    };
+  }
+
+  const lastValidDate = new Date(baseData.statementPeriodStart);
+  if (baseData.bankName === "barclays") {
+    baseData.transactionDateFormat = "dd MMM";
+  }
+  for (const entity of transactionsList) {
+    const transaction: Record<string, ValueWithReference> = {
+      Balance: null,
+      Date: null,
+      Type: null,
+      Description: null,
+      DepositAmount: null,
+      WithdrawalAmount: null,
+    };
+
+    // Process properties in the object
+    if (entity.type === "object" && entity.valueObject) {
+      for (const [propType, prop] of Object.entries(entity.valueObject)) {
+        const sourceRef = prop.id;
+        switch (propType) {
+          case "Balance": {
+            const amount = parseTransactionAmount(prop);
+            transaction[propType] =
+              amount !== null ? { value: amount, refId: sourceRef } : null;
+            break;
+          }
+          case "DepositAmount":
+          case "WithdrawalAmount": {
+            const amount = parseTransactionAmount(prop);
+            transaction[propType] =
+              amount !== null
+                ? { value: Math.abs(amount), refId: sourceRef }
+                : null;
+            break;
+          }
+          case "Date": {
+            const dateParsed = parseCompleteTransactionDate(
+              prop.content?.trim() ?? "",
+              lastValidDate,
+              new Date(baseData.statementPeriodEnd),
+              baseData.transactionDateFormat,
+            );
+            transaction.Date = dateParsed
+              ? { value: formatISODate(dateParsed), refId: sourceRef }
+              : null;
+            break;
+          }
+          case "Description":
+          case "Type": {
+            transaction[propType] = prop.content
+              ? { value: prop.content, refId: sourceRef }
+              : null;
+            break;
+          }
+          default: {
+            console.warn(`Unknown transaction type: ${propType}`);
+            break;
+          }
+        }
+      }
+    }
+
+    // Filter transactions without transferred amount
+    if (
+      transaction.WithdrawalAmount ||
+      transaction.DepositAmount ||
+      transaction.Date
+    ) {
+      transactions.push(transaction);
+    }
+  }
+
+  // Order transactions by y-position on page
+  const refMap: ReferenceMap = layout.referenceMap ?? {};
+  transactions.sort((a, b) => {
+    const refIdA =
+      a.WithdrawalAmount?.refId ?? a.DepositAmount?.refId ?? a.Date?.refId;
+    const refIdB =
+      b.WithdrawalAmount?.refId ?? b.DepositAmount?.refId ?? b.Date?.refId;
+    if (!refIdA || !refIdB) return 0;
+    const aRegion = refMap[refIdA]?.boundingRegions?.[0];
+    const bRegion = refMap[refIdB]?.boundingRegions?.[0];
+    if (!aRegion || !bRegion) return 0;
+    const pageDiff = (aRegion.pageNumber ?? 0) - (bRegion.pageNumber ?? 0);
+    if (pageDiff !== 0) return pageDiff;
+    const aY = aRegion.polygon?.[1] ?? 0;
+    const bY = bRegion.polygon?.[1] ?? 0;
+    return aY - bY;
+  });
+
+  // Reverse the transaction order if the bank statement is reverse ordered
+  if (baseData.isReverseOrdered) {
+    transactions.reverse();
+  }
+
+  // Fill in potentially missing dates, assuming the previous one extends
+  const statementStartDate: ValueWithReference = {
+    value: baseData.statementPeriodStart,
+    refId: "system",
+  };
+  let lastDate: ValueWithReference = statementStartDate;
+  for (const transaction of transactions) {
+    if (!transaction.Date) {
+      transaction.Date = lastDate;
+    } else {
+      lastDate = transaction.Date;
+    }
+  }
+
+  // filter out transactions without deposit or withdrawal amount
+  const cleanedTransactions = transactions.filter(
+    (transaction) => transaction.WithdrawalAmount || transaction.DepositAmount,
+  );
+
+  // validate transaction balances
+  const validationData = validateTransactionBalances(
+    cleanedTransactions,
+    baseData.openingBalance ?? 0,
+  );
+
+  return {
+    data: cleanedTransactions,
+    validationData:
+      validationData.Transactions.length > 0 ? validationData : undefined,
+  };
+}
+
+function validateTransactionBalances(
+  transactions: Record<string, ValueWithReference>[],
+  openingBalance?: number,
+): { Transactions: { index: number; message: string }[] } {
+  const errors: { index: number; message: string }[] = [];
+  let prevBalance: number = openingBalance ?? 0;
+  for (const [index, transaction] of transactions.entries()) {
+    const currentBalance = (transaction.Balance?.value ?? null) as
+      | number
+      | null;
+    const deposit = (transaction.DepositAmount?.value ?? 0) as number;
+    const withdrawal = (transaction.WithdrawalAmount?.value ?? 0) as number;
+
+    if (currentBalance === null) {
+      prevBalance = prevBalance + deposit - withdrawal;
+      transaction.Balance = { refId: "none", value: prevBalance };
+      continue;
+    }
+
+    const tolerance = 0.01;
+    if (
+      Math.abs(currentBalance - (prevBalance + deposit - withdrawal)) >
+      tolerance
+    ) {
+      console.warn(
+        "Balance validation failed: ",
+        currentBalance,
+        prevBalance,
+        deposit,
+        withdrawal,
+      );
+      errors.push({
+        index,
+        message: `Balance validation failed: got ${currentBalance.toFixed(2)} (deposit: ${deposit}, withdrawal: ${withdrawal})`,
+      });
+    }
+
+    prevBalance = currentBalance;
+  }
+
+  return { Transactions: errors };
+}
+
+function parseCompleteTransactionDate(
+  rawDateStr: string,
+  startDate: Date,
+  endDate: Date,
+  fmtStr: string,
+): Date | null {
+  if (!rawDateStr) return null;
+
+  const cleanedDateStr = rawDateStr
+    .replace(/(\d+)(?:st|nd|rd|th)/g, "$1")
+    .replace(/[\/\-\.\,]/g, " ")
+    .replace(/[\n\r\s]+/g, " ")
+    .replace(/T\d{2}:.+/g, "")
+    .trim()
+    .replace(
+      /\b([A-Z]{3,})\b/g,
+      (match) => match.charAt(0) + match.slice(1).toLowerCase(),
+    );
+
+  const tempDate = parse(cleanedDateStr, fmtStr, startDate);
+  console.log(
+    "PARSING DATE:",
+    rawDateStr,
+    "->",
+    cleanedDateStr,
+    "->",
+    tempDate,
+    "FROM:",
+    startOfDay(startDate),
+    "TO:",
+    endOfDay(endDate),
+    "FORMAT:",
+    fmtStr,
+  );
+
+  if (!isValid(tempDate)) {
+    console.log("INVALID DATE:", tempDate);
+    return null;
+  }
+
+  let normalizedDate = startOfDay(tempDate);
+  if (getMonth(normalizedDate) < getMonth(startDate)) {
+    normalizedDate = setYear(normalizedDate, getYear(endDate));
+  }
+
+  if (
+    isWithinInterval(normalizedDate, {
+      start: startOfDay(startDate),
+      end: endOfDay(endDate),
+    })
+  ) {
+    console.log("IN INTERVAL DATE:", normalizedDate);
+    return normalizedDate;
+  }
+
+  console.log("NOT IN INTERVAL DATE:", normalizedDate);
+  return null;
+}
+
+function parseTransactionAmount(entity: PretrainedEntity): number | null {
+  // Use valueNumber if available (normalized numeric value)
+  if (entity.type === "number" && entity.valueNumber !== undefined) {
+    return entity.valueNumber;
+  }
+
+  // Fall back to parsing content string
+  const isDebit = entity.content?.toLowerCase().includes("d") ?? false;
+  const cleanedString = entity.content?.replace(/[^0-9\.]/g, "");
+  if (!cleanedString) return null;
+  const amount = parseFloat(cleanedString);
+  return isDebit ? -amount : amount;
+}
+
+export async function getPretrainedBaseData(
+  documentLayout: DocumentLayout,
+): Promise<z.infer<typeof BankStatementBaseSchema>> {
+  const serializedLayout = documentLayout.serialize();
+  const result = await generateText({
+    model: google("gemini-3-pro-preview"),
+    system:
+      "Please extract all the required fields from the following text. Return dates in the ISO format YYYY-MM-DD.",
+    prompt: serializedLayout,
+    output: Output.object({
+      schema: BankStatementBaseSchema,
+    }),
+  });
+
+  return BankStatementBaseSchema.parse(result.output);
+}