engine/serial_reader.py

"""Product serial number detection and extraction.

Pipeline: Center-cropped camera frame → VLM (Qwen3-VL-8B) → digit-density
scoring → validation. The client crops the camera feed to the scan guide area
before uploading, so the VLM receives a focused image.
"""

import base64
import logging
import os
import re

import cv2
import numpy as np

logger = logging.getLogger(__name__)

# ── VLM config ────────────────────────────────────────────────────────────
_SERIAL_VLM_MODEL = os.environ.get(
    "OCR_SERIAL_VLM_MODEL", "Qwen/Qwen3-VL-8B-Instruct"
)
_SERIAL_VLM_TIMEOUT = 60

_SERIAL_VLM_PROMPT = (
    "Read ALL text visible in this image. "
    "Include numbers printed below barcodes, text on labels, stickers, plates, and screens. "
    "List each text block you can see, one per line. "
    "Read characters exactly as they appear (letters, digits, dashes, dots, slashes)."
)

# ── Validation config ─────────────────────────────────────────────────────
_SERIAL_MIN_LENGTH = 4
_SERIAL_MAX_LENGTH = 30
_SERIAL_ALLOWED = re.compile(r"^[A-Z0-9\-\./ ]+$")

_KNOWN_PREFIXES = {
    "S/N": "Serial number",
    "SN": "Serial number",
    "P/N": "Part number",
    "PN": "Part number",
    "MN": "Model number",
    "M/N": "Model number",
    "REF": "Reference number",
    "LOT": "Lot number",
    "NO.": "Number",
    "NO": "Number",
}

_REJECT_PATTERNS = [
    re.compile(r"^[0\s]+$"),
    re.compile(r"^(.)\1+$"),
    re.compile(r"(.)\1{7,}"),  # 8+ consecutive identical chars (e.g. WTP00000000)
    re.compile(r"(NONE|NULL|N/A|UNKNOWN|TEST|SAMPLE|DEMO)", re.I),
    re.compile(r"(DESIGNED|WARRANTY|MANUFACTURED|ASSEMBLED|MADE IN|PRINTED)", re.I),
    re.compile(r"\b(THE|AND|FOR|WITH|FROM|THIS|THAT|VOID|REMOVED|SEALED)\b", re.I),
]

# Labels used in extraction scoring.
# Serial-number label — "the usual culprits": S/N, SN, S.N, S-N, SERIAL,
# SER., optionally trailed by NO / NO. / NUMBER / #.
_SN_LABEL = r"(?:S[\s/.\-]?N|SER(?:IAL|\.))(?:\s*(?:NUMBER|NO|#)\.?)?"
_SN_LABEL_RE = re.compile(r"^" + _SN_LABEL + r"\s*:?\s*", re.I)
# Labels appearing in the text immediately before a candidate. An S/N label
# is decisive (the candidate always wins); other labels only earn a boost.
_SN_CONTEXT_RE = re.compile(_SN_LABEL + r"\s*:?\s*$", re.I)
_OTHER_LABEL_CONTEXT_RE = re.compile(r"(PART|REF|LOT|NO\.?|#)\s*:?\s*$", re.I)
_NOT_SERIAL_LABELS = r"I?MEI\d?|MEID|MAC|MODEL|HW\s*VERSION|VERSION|INPUT|POWER|SCAN"
_NOT_SERIAL_RE = re.compile(r"^(" + _NOT_SERIAL_LABELS + r")\s*:?\s*", re.I)
_NOT_SERIAL_CONTEXT_RE = re.compile(r"(" + _NOT_SERIAL_LABELS + r")\s*:?\s*$", re.I)


def _is_sequential(s: str) -> bool:
    """Detect sequential digit runs (hallucination indicator like 1234567)."""
    digits = re.sub(r"[^0-9]", "", s)
    if len(digits) < 5:
        return False
    seq_up = 0
    seq_down = 0
    for i in range(1, len(digits)):
        if int(digits[i]) == int(digits[i - 1]) + 1:
            seq_up += 1
        elif int(digits[i]) == int(digits[i - 1]) - 1:
            seq_down += 1
    return max(seq_up, seq_down) >= 4


# ── VLM extraction ────────────────────────────────────────────────────────

def extract_serial_vlm(img_bgr: np.ndarray) -> dict | None:
    """Send image to VLM, ask it to read all text, then extract serial pattern."""
    from engine.vlm_extractor import _VLM_PROVIDER, _get_hf_token

    api_token = _get_hf_token()
    if not api_token:
        logger.info("HF_API_TOKEN not set, skipping serial VLM extraction")
        return None

    _, buf = cv2.imencode(".jpg", img_bgr, [cv2.IMWRITE_JPEG_QUALITY, 90])
    image_data = base64.b64encode(buf.tobytes()).decode("utf-8")

    try:
        from huggingface_hub import InferenceClient

        logger.info("Calling serial VLM (%s via %s)", _SERIAL_VLM_MODEL, _VLM_PROVIDER)
        client = InferenceClient(provider=_VLM_PROVIDER, token=api_token, timeout=_SERIAL_VLM_TIMEOUT)
        response = client.chat.completions.create(
            model=_SERIAL_VLM_MODEL,
            messages=[{
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_data}"}},
                    {"type": "text", "text": _SERIAL_VLM_PROMPT},
                ],
            }],
            max_tokens=1024,
        )

        content = response.choices[0].message.content
        # Strip Qwen3 thinking tags if present
        think_end = content.find("</think>")
        if think_end != -1:
            content = content[think_end + len("</think>"):].strip()

        logger.info("Serial VLM text response: %.500s", content)

        raw_text = content.upper()
        serial = _extract_serial_from_text(raw_text)
        if serial:
            return {"serial_number": serial, "confidence": "medium", "type": "unknown", "source": "vlm"}

        logger.info("No serial pattern found in VLM text")
        return None

    except Exception as e:
        logger.warning("Serial VLM extraction failed: %s", e)
        return None


# ── Text extraction & scoring ─────────────────────────────────────────────

def _extract_serial_from_text(text: str) -> str | None:
    """Extract the most likely serial number from OCR/VLM text.

    A candidate explicitly labelled S/N (also SN, S.N, SERIAL, SER.) always
    wins over non-labelled ones. Otherwise scores by digit density (serials
    are digit-heavy), rejects IMEI/MEID prefixes, and boosts matches near
    other labels.
    """
    pattern = r"[A-Z0-9][A-Z0-9\-\./ ]{2,28}[A-Z0-9]"
    matches = re.findall(pattern, text)

    if not matches:
        cleaned = text.strip()
        if _SERIAL_MIN_LENGTH <= len(cleaned) <= _SERIAL_MAX_LENGTH:
            return cleaned
        return None

    def _serial_score(candidate: str) -> tuple[bool, float, str]:
        """Return (is_sn_labelled, score, cleaned_value) for a candidate.

        is_sn_labelled is True when the candidate carries an explicit S/N
        label — those win selection outright over non-S/N candidates.
        """
        s = candidate.strip()
        if len(s) < _SERIAL_MIN_LENGTH:
            return (False, -1.0, s)

        if _NOT_SERIAL_RE.match(s):
            return (False, -1.0, s)

        # Strip an attached S/N label; flag it and boost (if value is 8+ chars)
        sn_match = _SN_LABEL_RE.match(s)
        is_sn = False
        near_label = False
        if sn_match:
            s = s[sn_match.end():].strip()
            if len(s) < _SERIAL_MIN_LENGTH:
                return (False, -1.0, s)
            is_sn = True
            near_label = len(s) >= 8

        total = len(s.replace(" ", ""))
        if total == 0:
            return (False, -1.0, s)
        digits = sum(1 for c in s if c.isdigit())
        digit_ratio = digits / total

        if digits == 0:
            return (False, -1.0, s)
        if digit_ratio < 0.2:
            return (False, -0.5, s)
        if _is_sequential(s):
            return (False, -0.5, s)

        has_separators = bool(re.search(r"[\-\./]", s))

        # Check context before candidate in the original text
        pos = text.find(candidate.strip())
        if pos >= 0:
            prefix = text[max(0, pos - 20):pos]
            # Reject if preceded by a non-serial label (e.g. "IMEI1:862933...", "HW version: P052...")
            if _NOT_SERIAL_CONTEXT_RE.search(prefix):
                return (False, -1.0, s)
            # An S/N label in the surrounding text is decisive; other
            # labels (PART/REF/LOT/...) only earn the generic boost.
            if _SN_CONTEXT_RE.search(prefix):
                is_sn = True
                near_label = True
            elif _OTHER_LABEL_CONTEXT_RE.search(prefix):
                near_label = True

        score = digit_ratio * 2.0
        score += 0.3 if has_separators else 0.0
        score += 0.5 if near_label else 0.0

        slen = len(s.replace(" ", ""))
        if slen < 6:
            score -= 0.5
        elif slen < 8:
            score -= 0.2
        score += min(slen, 20) * 0.03

        return (is_sn, score, s)

    scored = [r for m in matches for r in [_serial_score(m)] if r[1] > 0]

    if not scored:
        return None

    # A candidate labelled S/N always wins over non-S/N ones; among
    # equally-labelled candidates the higher score breaks the tie.
    best = max(scored, key=lambda r: (r[0], r[1]))
    return best[2]


# ── Validation ────────────────────────────────────────────────────────────

def validate_serial(serial: str) -> dict:
    """Validate and normalize a candidate serial number."""
    if not serial:
        return {"valid": False, "normalized": None, "issues": ["empty"], "prefix_match": None}

    normalized = serial.upper().strip()

    # Strip known prefixes
    prefix_match = None
    for prefix, desc in _KNOWN_PREFIXES.items():
        pat = re.compile(r"^" + re.escape(prefix) + r"[\s:\.]*", re.I)
        if pat.match(normalized):
            prefix_match = desc
            normalized = pat.sub("", normalized).strip()
            break

    issues = []

    if len(normalized) < _SERIAL_MIN_LENGTH:
        issues.append(f"too_short (min {_SERIAL_MIN_LENGTH})")
    if len(normalized) > _SERIAL_MAX_LENGTH:
        issues.append(f"too_long (max {_SERIAL_MAX_LENGTH})")

    if not _SERIAL_ALLOWED.match(normalized):
        issues.append("invalid_characters")

    alphanum = normalized.replace(" ", "").replace("-", "").replace(".", "").replace("/", "")
    if alphanum:
        digit_count = sum(1 for c in alphanum if c.isdigit())
        digit_ratio = digit_count / len(alphanum)
        if digit_count == 0:
            issues.append("no_digits")
        elif digit_ratio < 0.2 and len(alphanum) > 6:
            issues.append("too_few_digits")

    word_groups = re.findall(r"[A-Z]{3,}", normalized)
    if len(word_groups) >= 3 and all(len(w) >= 3 for w in word_groups[:3]):
        issues.append("looks_like_text")

    if _is_sequential(normalized):
        issues.append("sequential_digits")

    for rp in _REJECT_PATTERNS:
        if rp.search(normalized):
            issues.append("rejected_pattern")
            break

    return {
        "valid": len(issues) == 0,
        "normalized": normalized,
        "issues": issues,
        "prefix_match": prefix_match,
    }


# ── Main pipeline ─────────────────────────────────────────────────────────

def process_serial_image(image_path: str) -> dict:
    """Full serial number extraction pipeline.

    1. Load image, basic quality check
    2. Try VLM on full image (fast, one call)
    3. Validate extracted serial, return result or failure
    """
    img_bgr = cv2.imread(image_path)
    if img_bgr is None:
        return _serial_fail("Image unreadable")

    h, w = img_bgr.shape[:2]
    quality = {"width": w, "height": h, "usable": True, "issues": []}

    if max(h, w) < 200:
        quality["usable"] = False
        quality["issues"].append("Image too small")
        return _serial_fail("Image too small", quality=quality)

    logger.info("Phase 1: Trying VLM on full image (%dx%d)", w, h)
    vlm_result = extract_serial_vlm(img_bgr)
    if vlm_result and vlm_result.get("serial_number"):
        validation = validate_serial(vlm_result["serial_number"])
        if validation["valid"]:
            logger.info("Valid serial from full-image VLM: %s", validation["normalized"])
            return {
                "serial_number": validation["normalized"],
                "bounding_box": None,
                "cropped_image": None,
                "extraction_source": "vlm",
                "confidence": vlm_result.get("confidence", "medium"),
                "validation": validation,
                "failure_reason": None,
                "regions_detected": 0,
                "quality": quality,
            }
        else:
            logger.info("Full-image VLM serial rejected: %s → %s",
                        vlm_result["serial_number"], validation["issues"])

    return _serial_fail("No valid serial number extracted", quality=quality)


def _serial_fail(reason: str, quality: dict | None = None) -> dict:
    return {
        "serial_number": None,
        "bounding_box": None,
        "cropped_image": None,
        "extraction_source": None,
        "confidence": "low",
        "validation": None,
        "failure_reason": reason,
        "regions_detected": 0,
        "quality": quality or {},
    }