app/lib/utils/heat_code.py

"""Heat code extraction for MTR (Material Test Report) PDFs.

Public API:
    extract_heat_codes(file_bytes, filename=None) -> list[str]
        Returns a deduplicated, order-preserving list of heat codes found
        in the document. Empty list if none found — caller's responsibility
        to route to the no-heat-code fallback.

Strategy (in order; results unioned, then filtered + deduped):
    1. pdfplumber table extraction — look for columns whose header matches
       a heat-code label and pull the column's non-empty values.
    2. Label-anchored regex on extracted text — handles inline "Heat No.: ABC123"
       and similar patterns.
    3. Filename hint — if the filename contains a plausible heat-code-shaped
       token AND that token also appears in document text, include it.
    4. OCR fallback for image-only PDFs — runs at 300 DPI with psm 6, then
       reruns strategies 2+3 on the OCR text.

Filtering rejects:
    - ASTM/SA/EN/ISO/ASME spec numbers ("A105", "SA-312", "EN 10204", "ASME B16.5")
    - Standalone 4-digit years (2020–2030)
    - Common ALL-CAPS technical words (MAX, MIN, PASS, FAIL, etc.)
    - Template placeholders ("__", "___", "—", "-")
    - Single characters or 1–3 char tokens (too short to be heat codes)

Designed to drop into `app/lib/utils/heat_code.py` in the email automation
repo with no changes — public API and imports are repo-compatible.
"""
from __future__ import annotations

import io
import logging
import re
from typing import Iterable

logger = logging.getLogger(__name__)


# ─── Label patterns ────────────────────────────────────────────────────────────
# Synonyms for "heat code" that may appear as labels in MTRs. Case-insensitive.
# Order matters only for logging; matching is union.
_LABEL_PATTERNS: tuple[str, ...] = (
    r"heat\s*code",
    r"heat\s*no\.?",
    r"heat\s*nr\.?",
    r"heat\s*#",
    r"heat\s*number",
    r"lot\s*no\.?",      # Bonney Forge and other small-mill MTRs
    r"lot\s*number",
)

# Matches a single label occurrence, capturing label and trailing token.
# Heat code token: 4-15 chars, alphanumeric with optional dashes.
_LABEL_ANCHORED_RE = re.compile(
    r"(?i)(?P<label>" + "|".join(_LABEL_PATTERNS) + r")"
    r"\s*[:\-]?\s*"
    r"(?P<code>[A-Z0-9][A-Z0-9\-]{2,14})",
)

# Standalone heat-code shape (used for filename/table validation).
# Must start with alphanumeric, 4-15 chars total, alphanumeric + dashes.
_CODE_SHAPE_RE = re.compile(r"^[A-Z0-9][A-Z0-9\-]{3,14}$", re.IGNORECASE)

# Pipe-spec patterns to reject (e.g., "40S", "304L", "80SS" — digit-led + short).
# Heat codes like "1ZZWN" survive because they have multiple letters in a row.
_PIPE_SPEC_RE = re.compile(r"^\d{1,4}[A-Z]{1,2}$", re.IGNORECASE)

# Phone number pattern: e.g. "910-738-8904". Common false positive when
# OCR picks up the vendor's contact info from the document header/footer.
_PHONE_RE = re.compile(r"^\d{3}-\d{3}-\d{4}$")

# Sales order pattern: single letter followed by 6+ digits, often with
# multiple leading zeros (e.g., "W000009810" is a WOI sales order #).
# Real heat codes don't follow this shape — they have at most 4 digits
# after a single letter, or use dashes when longer.
_ORDER_NUMBER_RE = re.compile(r"^[A-Z]\d{6,}$", re.IGNORECASE)

# Tokens to reject regardless of context.
_SPEC_PREFIXES: tuple[str, ...] = (
    "ASTM", "ASME", "SA-", "EN", "ISO", "DIN", "API", "AISI", "JIS", "NACE",
    "MSS", "MR0",
)

# All-caps words that look code-like but are technical jargon.
_REJECT_WORDS: frozenset[str] = frozenset({
    "MAX", "MIN", "AVG", "STD", "NA", "N/A", "TBD", "TBC", "ACCEPT", "REJECT",
    "PASS", "FAIL", "PASSED", "FAILED", "OK", "ACCEPTED",
    "REMARK", "REMARKS", "NOTE", "NOTES", "TOTAL", "VALUE", "RESULT",
    "TYPE", "GRADE", "SIZE", "DESCRIPTION", "QTY", "QUANTITY",
    "TEMP", "SPEC", "SPECIFICATION", "REV", "DATE", "PAGE", "BLANK", "NONE",
    "LADLE", "PRODUCT", "MILL", "NORMALIZED", "ANNEALED", "TEMPERED",
    "CUSTOMER", "ITEM", "ORDER", "PO", "INVOICE", "QUANTITY",
    "HEAT", "TREATMENT", "COMPOSITION", "ANALYSIS",
    "BUNDLE", "DRAWING", "MFG", "MFGID", "POR", "TAG",
    "DIMENSION", "DIMENSIONAL", "VISUAL", "CHEMICAL", "MECHANICAL", "HARDNESS",
    "TENSILE", "YIELD", "ELONG", "ELONGATION", "REDUCTION",
    "RESULT", "RESULTS", "EXAMINATION", "INSPECTION", "MATERIAL",
    "TENSION", "LOCATION", "ABBREVIATION", "WELDED", "SEAMLESS",
    "MARKING", "CERTIFIED", "CERTIFICATE", "REPORT",
})

# Spec number patterns to reject (e.g., A312, SA-105, A105-21).
_SPEC_NUM_RE = re.compile(r"^[A-Z]{1,4}-?\d{2,5}(-\d{1,3})?$", re.IGNORECASE)

# Year-like standalone tokens (2010-2039). Expanded from 2020-2030 because
# MTRs reference older standards years like "ISO 9001:2015" and Tesseract
# may misread small text producing tokens in this range.
_YEAR_RE = re.compile(r"^(20[0-3]\d)$")

# Material-spec pattern: "A106B", "A105N", "A234B" — letter A + 3 digits +
# single trailing letter. ASTM material grade specs. Distinct from real
# heat codes like "A8633" (letter + 4 digits, no trailing letter).
_MATERIAL_SPEC_RE = re.compile(r"^A\d{3}[A-Z]$", re.IGNORECASE)

# Standards-prefix patterns: anywhere in the document text, if a token
# follows one of these prefixes, that token is a standards reference, NOT
# a heat code. Examples: "ISO 15156", "NACE MR 0103", "ASTM A312", "EN 10204".
# The capture group is the standards number we want to reject.
_STANDARDS_REFERENCE_RES: tuple[re.Pattern[str], ...] = (
    re.compile(r"\bISO\s+(\d{3,6})\b", re.IGNORECASE),
    re.compile(r"\bNACE\s*MR\s*(\d{3,6})\b", re.IGNORECASE),
    re.compile(r"\bMR\s*(\d{3,6})\b", re.IGNORECASE),
    # ASTM revisions can carry a single trailing letter (e.g. "A370-12a"),
    # so allow an optional lowercase letter at the end of the revision suffix.
    re.compile(r"\bASTM\s+([A-Z]?\d{3,5}[A-Z]?(?:-\d{1,3}[a-zA-Z]?)?)\b", re.IGNORECASE),
    # Capture the SA prefix when present so "SA105-21" makes it into the
    # reject set (not just "105-21"). Otherwise "ASME SA105-21" in the doc
    # text wouldn't filter "SA105-21" as an extracted candidate.
    re.compile(r"\bASME\s+((?:SA-?)?[A-Z]?\d{2,5}[A-Z]?(?:[.\-]\d{1,4}[a-zA-Z]?)?)\b", re.IGNORECASE),
    re.compile(r"\bEN\s+(\d{3,6}(?::\d{4})?)\b", re.IGNORECASE),
    re.compile(r"\bDIN\s+(\d{3,6})\b", re.IGNORECASE),
    re.compile(r"\bAPI\s+(\d+L?)\b", re.IGNORECASE),
    re.compile(r"\bAISI\s+(\d{3,4}[A-Z]?)\b", re.IGNORECASE),
    re.compile(r"\bJIS\s+(\w+)\b", re.IGNORECASE),
    re.compile(r"\bMSS\s*SP\s*(\d{1,4})\b", re.IGNORECASE),
    re.compile(r"\bPED\s+(\d{4}/\d{2}/[A-Z]+)\b", re.IGNORECASE),
    # "ISO 9001:2015" — both numbers are bad
    re.compile(r"(\d{4}):(\d{4})"),
    # "A106/A106M-18" or "SA106/SA106M" — slash-separated material specs
    re.compile(r"\b([A-Z]?\d{3,4}[A-Z]?(?:-\d{1,3})?)/([A-Z]?\d{3,4}[A-Z]?(?:-\d{1,3})?)\b"),
)


# ─── Public API ────────────────────────────────────────────────────────────────

def extract_heat_codes(file_bytes: bytes, filename: str | None = None) -> list[str]:
    """Extract heat codes from an MTR PDF.

    Returns a deduplicated, ordered list. Empty if no codes can be confidently
    extracted — caller files to a `_no_heat_code/` fallback path.
    """
    candidates: list[tuple[str, str]] = []   # (code, source) for logging

    # Strategy 1: pdfplumber tables (text-layer PDFs)
    try:
        for code in _from_pdfplumber_tables(file_bytes):
            candidates.append((code, "table"))
    except Exception as e:
        logger.debug("pdfplumber table extraction failed: %s", e)

    # Strategy 2: column-position aware extraction (text-layer)
    # Handles MTRs where the heat label is a column header and the value sits
    # directly below it in the same X-range (DSI, 24L0604, A8633, KB419-1).
    try:
        for code in _from_column_position(file_bytes):
            candidates.append((code, "column"))
    except Exception as e:
        logger.debug("column-position extraction failed: %s", e)

    # Strategy 3: label-anchored regex on text-layer extract
    text_layer = _text_layer(file_bytes)
    for code in _from_label_anchored(text_layer):
        candidates.append((code, "text_label"))

    # Strategy 4: filename hint
    if filename:
        for code in _from_filename(filename, text_layer):
            candidates.append((code, "filename"))

    # Strategy 5: OCR fallback if we found nothing or text layer is thin
    if not candidates or len(text_layer) < 500:
        try:
            ocr_text = _ocr_text(file_bytes)
            for code in _from_label_anchored(ocr_text):
                candidates.append((code, "ocr_label"))
            for code in _from_ocr_column_heuristic(ocr_text):
                candidates.append((code, "ocr_column"))
            if filename:
                for code in _from_filename(filename, ocr_text):
                    candidates.append((code, "ocr_filename"))
        except Exception as e:
            logger.debug("OCR fallback failed: %s", e)

    # Strategy 6: targeted cell-region OCR. Last-resort recovery for MTRs
    # where the heat code sits in a small bordered table cell that standard
    # OCR can't read at default DPI (e.g. WOI MTRs with 1ZZWN). We use
    # pytesseract's image_to_data at 600 DPI to find the bbox of the
    # "HEAT" / "CODE" / "LOT" label words, crop a generous region adjacent
    # to each label, apply aggressive preprocessing (autocontrast,
    # contrast x3, double sharpen), and OCR each crop with psm 6.
    # Skip if we already have plausible-looking candidates AND text_layer
    # was thick (cheap escape hatch — the strategy is slow).
    if not candidates or len(text_layer) < 500:
        try:
            for code in _from_targeted_cell_ocr(file_bytes):
                candidates.append((code, "ocr_cell"))
        except Exception as e:
            logger.debug("Targeted cell-region OCR failed: %s", e)

    # Strategy 7: EasyOCR (deep-learning OCR). Final fallback for documents
    # where Tesseract genuinely cannot read the heat code at any DPI / PSM
    # combination — typically small or thin-font codes inside multi-column
    # tables (e.g., Flotite's G18M). EasyOCR uses a CRNN-based recognizer
    # that handles such cases better than Tesseract's classical pipeline.
    # The dependency is heavy (~2GB with torch) so we lazy-import and only
    # invoke when (a) cheaper strategies haven't produced clean results
    # and (b) the document looks image-only. Falls back gracefully if
    # easyocr isn't installed in the runtime — code still works without it.
    if (not candidates or len(text_layer) < 500) and _easyocr_should_run(candidates):
        try:
            for code in _from_easyocr(file_bytes):
                candidates.append((code, "easyocr"))
        except Exception as e:
            logger.debug("EasyOCR fallback failed: %s", e)

    # OCR commonly confuses the digit '0' with the letter 'O' (e.g. "JN021-1"
    # becomes "JNO21-1"). None of our 14 confirmed real heat codes uses the
    # letter O — every one has digit 0 in positions where O would appear. So
    # for OCR-sourced candidates we replace O→0 in-place. This avoids creating
    # duplicate folders ("JNO21-1" AND "JN021-1") and produces the canonical
    # form. Text-layer candidates are left alone — they have the correct chars.
    for idx in range(len(candidates)):
        code, source = candidates[idx]
        if source.startswith("ocr") and ("O" in code or "o" in code):
            normalized = code.replace("O", "0").replace("o", "0")
            if normalized != code:
                candidates[idx] = (normalized, f"{source}_o2zero")

    # Build the set of tokens that appear as standards-references in the
    # document text. Combine text_layer + ocr_text if both were collected.
    combined_text = text_layer
    if not candidates or len(text_layer) < 500:
        # OCR fallback may have produced text we should also scan.
        try:
            combined_text = (text_layer or "") + "\n" + _ocr_text(file_bytes)
        except Exception:
            pass
    standards_tokens = _collect_standards_numbers(combined_text)

    # Filter + dedupe (case-insensitive on code; preserve first-seen casing)
    filtered = _filter_and_dedupe(candidates, standards_tokens)

    if logger.isEnabledFor(logging.DEBUG):
        logger.debug(
            "Heat code extraction: %d raw candidates → %d after filter. Final: %s",
            len(candidates), len(filtered), filtered,
        )

    return filtered


# ─── Strategy implementations ──────────────────────────────────────────────────

def _text_layer(file_bytes: bytes) -> str:
    """Extract embedded text layer via pdfplumber. '' if image-only."""
    try:
        import pdfplumber
        with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
            return "\n".join((page.extract_text() or "") for page in pdf.pages[:3])
    except Exception:
        return ""


def _ocr_text(file_bytes: bytes) -> str:
    """OCR first 3 pages at 300 DPI with multiple page-segmentation modes
    to maximize chances of catching small/isolated heat-code text.

    psm 6  — uniform block of text. Good for regular tables.
    psm 11 — sparse text. Good for small isolated values in form boxes
             (e.g., MTR W000009810's "1ZZWN" in the small HEAT CODE cell).
    psm 4  — single column. Good for column-aligned tables (Flotite).

    Outputs are concatenated; downstream code is order-insensitive and
    dedupes candidates.
    """
    from pdf2image import convert_from_bytes
    import pytesseract
    imgs = convert_from_bytes(file_bytes, dpi=300, first_page=1, last_page=3)
    parts: list[str] = []
    for img in imgs:
        for psm in ("6", "11", "4"):
            try:
                parts.append(pytesseract.image_to_string(img, config=f"--psm {psm}"))
            except Exception as e:
                logger.debug("OCR psm=%s failed: %s", psm, e)
    return "\n".join(parts)


def _from_pdfplumber_tables(file_bytes: bytes) -> Iterable[str]:
    """Find tables whose header column matches a heat label; yield column values."""
    import pdfplumber

    label_re = re.compile(r"(?i)" + "|".join(_LABEL_PATTERNS))

    with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
        for page in pdf.pages[:3]:
            for table in page.extract_tables() or []:
                if not table or len(table) < 2:
                    continue

                # Find heat-label column index by scanning header rows.
                heat_col_idx: int | None = None
                for row in table[: min(3, len(table))]:
                    for idx, cell in enumerate(row):
                        if cell and label_re.search(str(cell)):
                            heat_col_idx = idx
                            break
                    if heat_col_idx is not None:
                        break

                if heat_col_idx is None:
                    continue

                # Pull values from data rows in that column.
                for row in table:
                    if heat_col_idx >= len(row):
                        continue
                    cell = row[heat_col_idx]
                    if not cell:
                        continue
                    for token in _tokenize_cell(str(cell)):
                        if _CODE_SHAPE_RE.match(token):
                            yield token


def _from_column_position(file_bytes: bytes) -> Iterable[str]:
    """Find the X-position of a heat-label word, then yield tokens whose
    x-center sits within ±half the label's width on the lines directly below.

    Handles MTRs where the heat label is a column header (in a wide
    whitespace-aligned table) and the data row sits below, e.g. DSI:
        Part   Specif.   HEAT No   C   Si   Mn   ...
        BODY   ASTM A216 WCB-2021    VC9387   0.213 ...
    """
    import pdfplumber

    label_re = re.compile(r"(?i)" + "|".join(_LABEL_PATTERNS))

    with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
        for page in pdf.pages[:3]:
            words = page.extract_words(use_text_flow=True) or []
            if not words:
                continue

            # Group words by approximate line (round y0 to nearest 2pt).
            lines: dict[int, list[dict]] = {}
            for w in words:
                key = round(w["top"] / 2)
                lines.setdefault(key, []).append(w)

            sorted_keys = sorted(lines.keys())

            # Look for a label word; capture its x-range; scan lines below.
            for i, key in enumerate(sorted_keys):
                line_words = lines[key]
                # Reconstruct the line text to test label patterns against
                # multi-word labels like "HEAT No" or "Heat Code".
                # Match individual words first, then consider their pairings.
                for j, word in enumerate(line_words):
                    pair_text = word["text"]
                    if j + 1 < len(line_words):
                        pair_text = pair_text + " " + line_words[j + 1]["text"]

                    if not label_re.search(pair_text):
                        continue

                    # X-range of the label (use bbox of just this word, or
                    # span both words if the match was over the pair).
                    x0 = word["x0"]
                    x1 = line_words[j + 1]["x1"] if (
                        j + 1 < len(line_words) and label_re.search(pair_text)
                        and not label_re.search(word["text"])
                    ) else word["x1"]
                    label_center = (x0 + x1) / 2
                    half_width = max((x1 - x0) / 2, 20.0)   # min 20pt to be forgiving

                    # Scan up to 20 lines below for tokens whose center is within
                    # half-width of the label. Multi-row MTR tables (KB419-1
                    # case has 6 rows; some have more) need a generous lookahead
                    # or we miss data rows farther down the table.
                    for next_key in sorted_keys[i + 1 : i + 21]:
                        for nw in lines[next_key]:
                            nw_center = (nw["x0"] + nw["x1"]) / 2
                            if abs(nw_center - label_center) <= half_width:
                                tok = nw["text"].strip().strip(",;|")
                                if _CODE_SHAPE_RE.match(tok):
                                    yield tok


# Cached EasyOCR reader. Module-level so we pay the model-loading cost once
# per process, not per MTR. ~50-100MB RAM + ~200MB models downloaded on first
# call (cached in ~/.EasyOCR/ thereafter).
_EASYOCR_READER = None


def _easyocr_should_run(candidates: list) -> bool:
    """Decide whether the EasyOCR fallback is worth the time. Run if we
    have very few candidates so far (zero or one) — this is the case for
    image-only docs where Tesseract is struggling. Skip for docs that
    already look healthy."""
    return len(candidates) <= 2


def _from_easyocr(file_bytes: bytes) -> Iterable[str]:
    """Deep-learning OCR fallback via EasyOCR. Recovers heat codes that
    Tesseract cannot read at any configuration (typically small / thin
    fonts in multi-column tables).

    Lazy-imports easyocr to avoid the heavy torch dependency cost when
    the runtime doesn't have it installed. The reader instance is cached
    at module scope so subsequent MTRs in the same process don't pay the
    model-load cost (~3-5 seconds first call, instant after).
    """
    global _EASYOCR_READER

    try:
        import easyocr
        from pdf2image import convert_from_bytes
        import numpy as np
    except ImportError:
        logger.debug("EasyOCR not available — install easyocr to enable")
        return

    try:
        if _EASYOCR_READER is None:
            # CPU-only, English-only. Set verbose=False to suppress the
            # noisy progress bars during model loading.
            _EASYOCR_READER = easyocr.Reader(["en"], gpu=False, verbose=False)
            logger.info("EasyOCR reader initialized")

        imgs = convert_from_bytes(
            file_bytes, dpi=400, first_page=1, last_page=2
        )
        for img in imgs:
            arr = np.array(img)
            # detail=0 returns just the text strings (no boxes/scores).
            results = _EASYOCR_READER.readtext(arr, detail=0)
            for text in results:
                if not text:
                    continue
                for tok in _tokenize_cell(text):
                    if _CODE_SHAPE_RE.match(tok):
                        yield tok
    except Exception as e:
        logger.debug("EasyOCR run failed: %s", e)


def _from_targeted_cell_ocr(file_bytes: bytes) -> Iterable[str]:
    """Targeted cell-region OCR for MTRs where the heat code sits in a
    small bordered table cell that default-DPI Tesseract can't read.

    Approach:
        1. Render the PDF at 600 DPI (3-4x larger than _ocr_text).
        2. Use image_to_data to find the bbox of every "HEAT" / "CODE" /
           "LOT" / "NO." label word.
        3. For each label box, crop a generous adjacent region (~1200×400
           pixels at 600 DPI = ~2x2 inch chunk to the right/below the
           label).
        4. Apply aggressive preprocessing: grayscale → autocontrast →
           contrast x3 → double sharpen. Recovers small/faint text that
           default OCR misses.
        5. OCR each crop with psm 6 and yield heat-code-shaped tokens.

    Proven to recover '1ZZWN' from MTR W000009810 which was unreadable
    by every other strategy.
    """
    try:
        from pdf2image import convert_from_bytes
        from PIL import ImageEnhance, ImageFilter, ImageOps
        import pytesseract
    except ImportError:
        return

    # Render at 600 DPI. Cap at 2 pages — performance vs coverage tradeoff.
    try:
        imgs = convert_from_bytes(file_bytes, dpi=600, first_page=1, last_page=2)
    except Exception:
        return

    # Code-indicator words. We only crop when "HEAT" is followed by one
    # of these in the next few words, OR when the word itself is one of
    # these (e.g., a standalone "CODE" header in some MTRs). This avoids
    # spuriously cropping around "HEAT TREATMENT" which sits adjacent to
    # mechanical-properties values in many MTR layouts.
    code_indicators = {"CODE", "NO", "NO.", "NR", "NR.", "NUMBER", "#"}
    lot_indicators = {"LOT"}

    for img in imgs:
        try:
            # Find word-level bboxes.
            data = pytesseract.image_to_data(
                img, output_type=pytesseract.Output.DICT
            )
        except Exception:
            continue

        n = len(data.get("text", []))
        for i in range(n):
            word = (data["text"][i] or "").strip().upper()
            if not word:
                continue

            # Determine if this is a valid heat-code label position.
            is_heat_label = False
            if word == "HEAT":
                # Check next 1-2 words for a code indicator on the same line.
                for k in range(1, 3):
                    if i + k >= n:
                        break
                    nxt = (data["text"][i + k] or "").strip().upper().rstrip(".")
                    if not nxt:
                        continue
                    # Allow column-aligned lookahead — same approximate row.
                    if abs(int(data["top"][i + k]) - int(data["top"][i])) > 30:
                        break
                    if nxt in code_indicators or nxt in {"CODE", "NO", "NR", "NUMBER"}:
                        is_heat_label = True
                        break
                    # Stop searching after first non-matching word.
                    break
            elif word in lot_indicators:
                # "LOT NO." style label — check next word for NO/NUMBER
                if i + 1 < n:
                    nxt = (data["text"][i + 1] or "").strip().upper().rstrip(".")
                    if nxt in {"NO", "NUMBER", "#"}:
                        is_heat_label = True
            elif word == "CODE":
                # Standalone CODE only valid if preceded by HEAT.
                if i > 0:
                    prev = (data["text"][i - 1] or "").strip().upper()
                    if prev == "HEAT":
                        # Already handled by the "HEAT" branch above; skip
                        # to avoid duplicate crops around the same label.
                        continue

            if not is_heat_label:
                continue

            try:
                x0 = int(data["left"][i])
                y0 = int(data["top"][i])
                # Generous crop: ~1200×400 px at 600 DPI is ~2 inches.
                x1 = min(img.width, x0 + 1200)
                y1 = min(img.height, y0 + 400)
                crop = img.crop((x0, y0, x1, y1))
                # Aggressive preprocessing.
                g = crop.convert("L")
                g = ImageOps.autocontrast(g, cutoff=1)
                g = ImageEnhance.Contrast(g).enhance(3.0)
                g = g.filter(ImageFilter.SHARPEN)
                g = g.filter(ImageFilter.SHARPEN)
                # OCR with psm 6 (uniform block of text).
                txt = pytesseract.image_to_string(g, config="--psm 6")
            except Exception:
                continue
            # Yield heat-code-shaped tokens from the crop.
            for tok in _tokenize_cell(txt):
                if _CODE_SHAPE_RE.match(tok):
                    yield tok


def _from_ocr_column_heuristic(ocr_text: str) -> Iterable[str]:
    """For OCR text: when a line contains a heat label, yield code-shaped
    tokens from the next 3 lines. Cheap-and-pretty-good for the cases where
    the OCR preserves line breaks between header and data rows.
    """
    if not ocr_text:
        return
    label_re = re.compile(r"(?i)" + "|".join(_LABEL_PATTERNS))
    # Section headers that mark the end of a "heat code" data region.
    # When we hit a line containing one of these (and no heat-code-shaped
    # tokens), stop scanning — we've left the heat code table and entered
    # a different section (mechanical properties, chemical analysis, etc.)
    # that contains tensile-strength numbers, yield values, etc. that look
    # heat-code-shaped but are not.
    section_break_re = re.compile(
        r"\b(?:MECHANICAL|CHEMICAL|TENSILE|HARDNESS|IMPACT|HYDROSTATIC|"
        r"PROPERTIES|COMPOSITIONS?|TREATMENT|ANALYSIS|BEND|EDDY|MICROSTRUCTURE|"
        r"REMARKS?|NOTES?|"
        # Address-block markers: catches lines like "4815 West 5th Street
        # Lumberton NC 28358" that contain ZIP codes / street numbers OCR'd
        # as heat-code-shaped tokens.
        r"STREET|AVENUE|ROAD|BOULEVARD|DRIVE|HIGHWAY|"
        r"WEBSITE|EMAIL|PHONE|FAX|TEL)\b",
        re.IGNORECASE,
    )
    lines = ocr_text.splitlines()
    for i, line in enumerate(lines):
        if not label_re.search(line):
            continue
        # Look at the next 20 non-empty lines OR until we hit a section
        # break. Multi-row MTR tables can have 6+ data rows (KB419-1 case);
        # a tight lookahead drops rows past the first few. The section-
        # break stop prevents this wider lookahead from leaking into
        # downstream sections like mechanical/chemical properties.
        looked = 0
        for j in range(i + 1, min(i + 40, len(lines))):
            nxt = lines[j].strip()
            if not nxt:
                continue
            # If this line contains a section-break keyword, we've left the
            # heat-code data region — stop scanning unconditionally. Words
            # like "CHEMICAL", "MECHANICAL", "TENSILE" themselves match the
            # code-shape regex but are clearly section markers, not heat
            # codes.
            if section_break_re.search(nxt):
                break
            looked += 1
            if looked > 20:
                break
            for tok in _tokenize_cell(nxt):
                if _CODE_SHAPE_RE.match(tok):
                    yield tok


def _from_label_anchored(text: str) -> Iterable[str]:
    """Yield candidates following 'Heat No.', 'Lot No.', etc."""
    if not text:
        return
    for m in _LABEL_ANCHORED_RE.finditer(text):
        yield m.group("code")


def _from_filename(filename: str, doc_text: str) -> Iterable[str]:
    """If filename contains a heat-code-shaped token that appears (exactly or
    fuzzily) in document text, yield it.

    Fuzzy match handles OCR-mangled docs where e.g. "SD55822" in the filename
    shows up as "S055822" in the OCR text (D→0 confusion). Uses simple
    same-length / one-edit-away check (good enough for OCR drift).
    """
    # Split filename on whitespace and underscore ONLY — preserve dashes so
    # codes like "Y211213D40-1" and "KB419-1" survive intact.
    base = filename.rsplit(".", 1)[0]   # drop .pdf
    tokens = re.split(r"[\s_]+", base)
    text_upper = doc_text.upper() if doc_text else ""

    for tok in tokens:
        if not _CODE_SHAPE_RE.match(tok):
            continue
        tok_upper = tok.upper()
        if tok_upper in text_upper:
            yield tok
        elif _fuzzy_token_in_text(tok_upper, text_upper, max_edits=2):
            yield tok


def _fuzzy_token_in_text(token: str, text: str, max_edits: int = 2) -> bool:
    """True if `text` contains a same-length window within `max_edits` of `token`.

    Cheap-and-good-enough: scans every same-length substring and counts char
    differences. O(len(text) * len(token)) but n is small (n ≤ 15) and we
    only call it as a fallback.
    """
    n = len(token)
    if n == 0 or len(text) < n:
        return False
    for i in range(len(text) - n + 1):
        window = text[i : i + n]
        diff = sum(1 for a, b in zip(token, window) if a != b)
        if diff <= max_edits:
            return True
    return False


def _tokenize_cell(cell_text: str) -> list[str]:
    """Split a table cell into candidate tokens, preserving dashes."""
    return [t.strip() for t in re.split(r"[\s,;|]+", cell_text) if t.strip()]


def _collect_standards_numbers(doc_text: str) -> set[str]:
    """Scan document text for standards references (ISO 15156, NACE MR 0103,
    ASTM A312, EN 10204, etc.) and return the set of standard-number tokens
    that should never be treated as heat codes.

    Returned tokens are uppercased and stripped of leading non-alphanumerics
    so they can be compared directly against `code.upper()` in the filter.
    """
    bad: set[str] = set()
    if not doc_text:
        return bad
    for pat in _STANDARDS_REFERENCE_RES:
        for m in pat.finditer(doc_text):
            for grp in m.groups():
                if not grp:
                    continue
                tok = re.sub(r"^[^A-Za-z0-9]+", "", grp).upper()
                if tok:
                    bad.add(tok)
                    # Also reject the no-dash form (e.g., "A106-18" → also drop "A106")
                    no_dash = tok.split("-")[0]
                    if no_dash != tok:
                        bad.add(no_dash)
    return bad


def _filter_and_dedupe(
    candidates: list[tuple[str, str]],
    standards_tokens: set[str] | None = None,
) -> list[str]:
    """Apply rejection rules + dedupe case-insensitively, preserving order.
    Also drops substring duplicates: if 'KB419' and 'KB419-1' are both present,
    keep only the longer one (avoid over-counting when codes get split).
    `standards_tokens` is an optional set of uppercased tokens that were
    detected as standards references in the document text — these are rejected
    regardless of their shape."""
    standards_tokens = standards_tokens or set()
    seen: dict[str, str] = {}   # upper -> original-cased

    for code, _source in candidates:
        code = code.strip()
        upper = code.upper()

        if upper in seen:
            continue
        if upper in standards_tokens:
            continue
        if not _is_plausible_heat_code(code):
            continue
        seen[upper] = code

    # Substring dedup: drop any code that is a strict prefix OR strict suffix
    # of another. Prefix case handles e.g. KB419 / KB419-1 splits. Suffix case
    # handles OCR-truncation: "N021-1" is the suffix of "JN021-1" — happens
    # when Tesseract drops the leading letter of a heat code.
    uppers = list(seen.keys())
    keep = set(uppers)
    for u in uppers:
        for other in uppers:
            if u == other or len(other) <= len(u):
                continue
            if other.startswith(u) or other.endswith(u):
                keep.discard(u)
                break
    return [seen[u] for u in uppers if u in keep]


def _is_plausible_heat_code(code: str) -> bool:
    upper = code.upper()

    if not _CODE_SHAPE_RE.match(code):
        return False
    if upper in _REJECT_WORDS:
        return False
    if _YEAR_RE.match(upper):
        return False
    # Pure-numeric tokens with 8 or more digits look like dates (YYYYMMDD)
    # or order/quantity values, not heat codes. The largest pure-numeric
    # real heat code we've seen is 5 digits (10720); 8+ digits is always
    # something else.
    if upper.isdigit() and len(upper) >= 8:
        return False
    # Pure-numeric tokens starting with '0' are not heat codes — they're
    # OCR artifacts from chemical-analysis decimal values (".035" read as
    # "0035") or from index/sequence numbers. None of our 14 confirmed
    # heat codes starts with a digit-0; pure-numeric reals (10720, etc.)
    # all start with 1–9.
    if upper.isdigit() and upper.startswith("0"):
        return False
    # Pure-numeric 4-digit tokens are not heat codes in our confirmed
    # samples — the smallest pure-numeric real code is 10720 (5 digits).
    # 4-digit pure numerics are typically mechanical test values
    # (3400 PSI, 2500 PSI, etc.) or quantities.
    if upper.isdigit() and len(upper) == 4:
        return False
    # Phone numbers: 910-738-8904 style. Common false positive when OCR
    # picks up the vendor's contact info from a document header/footer.
    if _PHONE_RE.match(upper):
        return False
    # Sales order numbers: single letter + 6+ digits, often W000... style.
    # Real heat codes max out at letter + 4 digits before a dash.
    if _ORDER_NUMBER_RE.match(upper):
        return False
    # Material-spec pattern (A106B, A105N, A234B) — letter A + 3 digits +
    # single trailing letter. Real heat codes like A8633 have 4 digits and
    # no trailing letter, so they're not affected.
    if _MATERIAL_SPEC_RE.match(upper):
        return False
    # Pipe-spec rejection (40S, 304L, 80SS-style schedule/grade designators).
    if _PIPE_SPEC_RE.match(upper):
        return False
    # Spec numbers — only reject when the token explicitly starts with a known
    # spec prefix (ASTM, ASME, SA-, EN, ISO, etc.). Pure shape ("D460", "KB419-1")
    # overlaps with real heat codes; we don't reject those on shape alone.
    if any(upper.startswith(p) for p in _SPEC_PREFIXES):
        return False
    # Reject tokens with no digits at all. Real heat codes always contain at
    # least one digit (verified across all 14 confirmed samples). This catches
    # both pure-alpha tokens (DIMENSION) and alpha+dash tokens (CON-REDUCER,
    # WP304-S, etc.) that the per-char isalpha() check would miss.
    digits = [c for c in upper if c.isdigit()]
    if not digits:
        return False

    # Reject tokens whose ONLY digit(s) are zero(s). This is the OCR-mangled-
    # English-word signature: Tesseract reads letter O as digit 0, so words
    # like ACCORDING, BELOW, FROM, DONE, COMPOSITION, etc. become ACC0RDING,
    # BEL0W, FR0M, D0NE, C0MP0SITI0N. None of our 14 confirmed real heat codes
    # have all-zero digits; they always have at least one 1–9 digit.
    if all(d == "0" for d in digits):
        return False

    return True