src/services/label_ocr.py

"""Label OCR — DeepSeek-OCR-2 specialist + Gemma 4 reasoning 兩階段 pipeline。

【這個檔案在做什麼？】
  用戶拍一張衣服洗標（care label）的照片，
  這個模組負責從照片中辨識出：品牌、商品名、貨號、尺碼、材質、洗滌說明、產地。

【做法演進】
  v1（已棄用）：Claude Vision API → 每張 ~$0.0012，需要 API key
  v2（已棄用）：Florence-2 OCR + regex 解析 → regex 不夠泛用
  v3（已棄用）：Florence-2 OCR + Qwen2.5 LLM 解析 → 需要兩個模型
  v4（已棄用）：Gemma 4 E4B 視覺一步完成 → 通用 VLM 跑 OCR token 成本高
  v5（現在）：  DeepSeek-OCR-2 specialist 出 markdown → Gemma 4 reasoning 抽 7 欄位
                 OCR specialist + reasoning specialist 各司其職，user-facing schema 不變

【兩階段流程】
  1. Stage 1 (DeepSeek-OCR-2)：圖片 → markdown raw text（純 OCR，91% OmniDocBench）
  2. Stage 2 (Gemma 4)：markdown → 結構化 JSON（reasoning，brand vs product_name 等區分）
  3. Regex 後處理：補抓 LLM 漏掉的 product_code + 清理 material 格式

【為什麼分兩階段？】
  - DeepSeek-OCR-2（3B BF16）是 OCR specialist，比通用 VLM 精度高且 throughput 好
  - Gemma 4 E4B 已 resident（chat / stylist / qa 等 capability 共用），reasoning 零增量 cost
  - 對外 wire schema (LabelOcrResult) 不變，內部分工改變屬 implementation detail

【Rollback path】
  OCR_MODEL=gemma-4-e4b 環境變數切回單階段純 Gemma 視覺路徑（v4 行為），
  用於 git-bisect / Apple Silicon dev / DeepSeek prod issue 即時切換。
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import re
import tempfile
from string import Template
from typing import TYPE_CHECKING

from PIL import Image
from pydantic import BaseModel

from src.lib.prompt_safety import sanitize_for_prompt
from src.services.gemma_inference import gemma_text, gemma_vision

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


# ============================================================================
# 回傳資料結構
# ============================================================================


class LabelOcrResult(BaseModel):
    """
    OCR 辨識結果。每個欄位都是 optional（可能洗標上沒印）。

    - brand_raw:          洗標上的原始品牌名
                          → 上游會再經過 brand normalizer 對照 DB 品牌表
    - product_name:       商品名稱（例如 "CONDUIT DOWN JACKET"）
    - product_code:       商品貨號/型號（例如 "X000009960"）
    - size:               尺碼（各品牌格式不同：S/M/L、36/38/42、170/96A 等）
    - material:           材質組成（例如 "100% Nylon"）
    - care_instructions:  洗滌說明（文字描述）
    - country_of_origin:  產地（例如 "Italy"）
    """

    brand_raw: str | None = None
    product_name: str | None = None
    product_code: str | None = None
    size: str | None = None
    material: str | None = None
    care_instructions: str | None = None
    country_of_origin: str | None = None


# ============================================================================
# Prompt 定義
# ============================================================================

# Stage 1 — DeepSeek-OCR-2 純 OCR prompt（model card 標準格式）
# `<|grounding|>` 啟用 layout-aware 抽取（含 bounding box token），對洗標多區塊
# 排版（品牌印標 + 數字貨號 + 圖示）有幫助；輸出仍是 markdown 含 grounding tokens。
_DEEPSEEK_OCR_PROMPT = "<image>\n<|grounding|>Convert the document to markdown. "

# Stage 2 — Gemma 4 把 markdown 抽 7 個欄位
#
# 用 string.Template（$-placeholder）而非 str.format：OCR 文字（弱信任，洗標內容
# 攻擊者可控）含字面 `{` / `}` 會讓 str.format 拋 KeyError/IndexError/ValueError，
# 被 extract_label 外層 except 吞成空 LabelOcrResult — 一個會被惡意/雜亂洗標觸發的
# 靜默失效。Template.safe_substitute 對任意 `{}`/`$`（OCR 不太可能含 `$name`，含了
# 也原樣保留）都安全、永不對缺鍵或字面括號炸。JSON 範例的大括號在 Template 下是
# 普通字元（不需 `{{` 轉義），唯一特殊字元是 `$ocr_text` 佔位符。
_GEMMA_REASONING_PROMPT_TEMPLATE = Template("""\
You are a fashion label parser. Below is OCR text extracted from a clothing care label.
Parse it into structured JSON.

OCR text:
<<<
$ocr_text
>>>

Rules:
- brand_raw: the BRAND name (e.g. Nike, Gucci), NOT the product name. null if not visible.
- product_name: garment style name (e.g. "Conduit Down Jacket"). null if not visible.
- product_code: alphanumeric style code (e.g. X000009960, CW2288-111). Look for long alphanumeric sequences.
- size: the FULL size string. Prefer "170/96A(S)" over just "S". Include Asian sizing if present.
- material: main material with percentage, English only (e.g. "100% Nylon"). Ignore translations.
- care_instructions: washing/care text in English. null if only symbols.
- country_of_origin: country from "Made in ..." text. null if not found.

Return ONLY valid JSON with these exact keys, no explanation:
{"brand_raw": null, "product_name": null, "product_code": null, "size": null, "material": null, "care_instructions": null, "country_of_origin": null}""")

# v4 rollback path — Gemma 4 single-pass 直接看圖出 JSON（OCR_MODEL=gemma-4-e4b）
_GEMMA_VISION_PROMPT = """\
You are a fashion label OCR expert. Look at this care label photo and extract all visible information.

Rules:
- brand_raw: the BRAND name (e.g. Nike, Gucci), NOT the product name. null if not visible.
- product_name: garment style name (e.g. "Conduit Down Jacket"). null if not visible.
- product_code: alphanumeric style code (e.g. X000009960, CW2288-111). Look for long alphanumeric sequences.
- size: the FULL size string. Prefer "170/96A(S)" over just "S". Include Asian sizing if present.
- material: main material with percentage, English only (e.g. "100% Nylon"). Ignore translations.
- care_instructions: washing/care text in English. null if only symbols.
- country_of_origin: country from "Made in ..." text. null if not found.

Return ONLY valid JSON with these exact keys, no explanation:
{"brand_raw": null, "product_name": null, "product_code": null, "size": null, "material": null, "care_instructions": null, "country_of_origin": null}"""


# ============================================================================
# 核心函式 — capability-aware dispatcher
# ============================================================================


async def extract_label(image: Image.Image) -> LabelOcrResult:
    """
    從洗標照片中提取結構化資訊。

    Dispatch 邏輯：
    - OCR_MODEL=gemma-4-e4b（rollback）→ 走 v4 單階段 Gemma 視覺路徑
    - 預設 OCR_MODEL=deepseek-ocr-2 → 走 v5 兩階段（DeepSeek OCR + Gemma reasoning）

    對外 wire schema 不變：回傳 LabelOcrResult 7 欄位。
    """
    raw_text = ""
    try:
        # 走 OCR capability 看綁的是哪個 model — 從 registry 拿 model_id 判斷路徑
        # （沿 trend_predictor.py 的 TREND_MODEL env 模式）
        ocr_model_id = await _get_ocr_model_id()

        if "deepseek-ocr" in ocr_model_id.lower():
            # v5 兩階段 path
            raw_text = await _stage1_deepseek_ocr(image)
            logger.info("DeepSeek OCR raw output: %s", raw_text[:300])
            json_response = await _stage2_gemma_reasoning(raw_text)
        else:
            # v4 rollback path（OCR_MODEL=gemma-4-e4b）
            json_response = await gemma_vision(
                "ocr",
                image,
                _GEMMA_VISION_PROMPT,
                max_tokens=512,
            )
            raw_text = json_response  # 給 _postprocess regex 補漏用
            logger.info("Gemma OCR raw output: %s", raw_text[:300])

        result = _parse_llm_response(json_response)

    except Exception as e:
        logger.warning("OCR extraction failed: %s", e)
        result = LabelOcrResult()

    # Stage 3: Regex 後處理（補漏 + 清理）
    result = _postprocess(result, raw_text)

    return result


async def _get_ocr_model_id() -> str:
    """從 ModelRegistry 拿 ocr capability 綁的 model_id（決定 v4/v5 path）。

    fail-soft：registry 未初始化（unit test）/ get_*_config raise → 走 fallback。
    Fallback 不能無條件 return DeepSeek — 那會 invert rollback safety direction：
    當 operator 已設 OCR_MODEL=gemma-4-e4b 想跳開 DeepSeek，registry 又恰好壞掉時，
    我們仍應尊重 rollback env，否則 _stage1_deepseek_ocr 拿到 Gemma 的 model object
    跑 model.infer() AttributeError → outer except 吞成空 LabelOcrResult（reviewer P1-A）。
    """
    try:
        from src.registry import get_registry

        reg = get_registry()
        cap_config = reg.get_capability_config("ocr")
        model_key = cap_config.get("model", "deepseek-ocr-2")
        model_config = reg.get_model_config(model_key)
        return model_config.get("model_id", "deepseek-ai/DeepSeek-OCR-2")
    except (RuntimeError, KeyError, ValueError) as e:
        # Registry 壞掉時，仍尊重 OCR_MODEL rollback env。
        # widen 到 ValueError（review-task P1-A）：未來 reg.get_*_config 若改以
        # ValueError 表達 invalid model key，原本 (RuntimeError, KeyError) 會
        # silent escape 到 outer extract_label except，吞成空 LabelOcrResult，
        # rollback intent 也被跳過。log exception class 讓 prod incident 可診斷。
        logger.warning(
            "_get_ocr_model_id registry lookup failed (%s): %s — falling back to env",
            type(e).__name__,
            e,
        )
        rollback_env = os.environ.get("OCR_MODEL", "").strip().lower()
        if rollback_env == "gemma-4-e4b":
            return "google/gemma-4-E4B-it"
        return "deepseek-ai/DeepSeek-OCR-2"


# ============================================================================
# Stage 1 — DeepSeek-OCR-2 specialist OCR
# ============================================================================


async def _stage1_deepseek_ocr(image: Image.Image) -> str:
    """DeepSeek-OCR-2 純 OCR：PIL Image → markdown raw text。

    【為什麼用 tempfile？】
    DeepSeek-OCR-2 的 model.infer(image_file=...) 接受**檔案路徑字串**，
    不接受 PIL.Image 物件。我們的 caller 已經把 user upload 解碼成 PIL Image
    （走過 PIL bomb cap + content type 檢查），所以這裡再寫到 temp file 給 infer 用。
    delete=False + try/finally 確保異常時也清掉。

    【為什麼 to_thread？】
    model.infer() 跑 ~3B params forward + visual cropping，CUDA 上 ~1-3s，
    放 async def body 會 pin event loop（.claude/rules/ml-service.md async discipline）。
    """
    from src.registry import get_registry

    reg = get_registry()
    loaded = await reg.get("ocr")
    model = loaded.model
    tokenizer = loaded.tokenizer
    base_size = loaded.config.get("base_size", 1024)
    image_size = loaded.config.get("image_size", 768)
    crop_mode = loaded.config.get("crop_mode", True)

    # Per-request 獨立 temp dir — 避免並發 model.infer 共用 tempfile.gettempdir()
    # 競爭 grounding 中間檔（review-task P1-B）。即使 save_results=False，
    # trust_remote_code 載入的 deepseek_vl_v2 custom code 仍可能用 output_path
    # 為工作目錄寫 intermediate；L40S 多 concurrent 真風險，dev mock test 蓋不到。
    work_dir = tempfile.mkdtemp(prefix="ocr_deepseek_")
    tmp_path = os.path.join(work_dir, "input.jpg")

    try:
        # JPEG quality=95 在 OCR 任務中對精度影響可忽略，比 PNG 小一個數量級
        await asyncio.to_thread(image.save, tmp_path, "JPEG", quality=95)

        def _run_infer() -> str:
            return model.infer(
                tokenizer,
                prompt=_DEEPSEEK_OCR_PROMPT,
                image_file=tmp_path,
                output_path=work_dir,  # per-request isolated working dir
                base_size=base_size,
                image_size=image_size,
                crop_mode=crop_mode,
                save_results=False,
            )

        raw_output = await asyncio.to_thread(_run_infer)
    finally:
        # rmtree 比 unlink 安全：infer 可能在 work_dir 留 grounding 中間檔，
        # 不 rmtree 會 leak 永遠不清的 dir。ignore_errors=True：OS 自有 temp
        # cleanup，這裡優先不打擾 caller exception flow。
        import shutil

        try:
            shutil.rmtree(work_dir, ignore_errors=True)
        except OSError as cleanup_err:
            logger.debug("OCR temp dir cleanup failed: %s — %s", work_dir, cleanup_err)

    # 清掉 grounding tokens（<|ref|>...<|/ref|><|det|>...<|/det|>），保留純文字
    # re.DOTALL 確保 grounding span 跨行也能 strip（DeepSeek 多區塊輸出常含 \n）
    cleaned = re.sub(
        r"<\|ref\|>.*?<\|/ref\|><\|det\|>.*?<\|/det\|>",
        "",
        raw_output or "",
        flags=re.DOTALL,
    )
    return cleaned.strip()


# ============================================================================
# Stage 2 — Gemma 4 reasoning 把 markdown 抽 JSON
# ============================================================================


async def _stage2_gemma_reasoning(ocr_text: str) -> str:
    """Gemma 4 把 OCR markdown 解析成 7 欄位 JSON 字串。

    使用 ocr_reasoning capability（綁 gemma-4-e4b）— 與 chat / stylist / qa 共用同一
    Gemma 4 instance，reasoning 零增量 cost。
    """
    if not ocr_text or not ocr_text.strip():
        return "{}"

    # OCR 文字（DeepSeek-OCR 輸出）是弱信任的多行自由文字（品牌/尺寸/材質分行），
    # 雖非直接 user input 但洗標內容可被攻擊者控制 → 進 Gemma prompt 前必過單一來源
    # sanitize。preserve_newlines=True 保留分行結構（對齊 translator 那類弱信任長文字），
    # 只 strip 不可見/控制/bidi carrier；max_len=4000 對齊原本的 [:4000] slice。
    # boundary token 由 template 的 <<< >>> 提供（保留）。sanitize 回 None（清空後為空）
    # 時 fallback 空字串，不讓 prompt 變 "None"（記憶 no_fabricated_fallbacks）。
    safe_ocr_text = (
        sanitize_for_prompt(ocr_text, preserve_newlines=True, max_len=4000) or ""
    )
    # safe_substitute (not substitute/format): OCR text with literal `{`/`}`/`$`
    # never raises — it would otherwise be swallowed into an empty result by the
    # outer except. The text is already carrier-stripped + clamped above.
    prompt = _GEMMA_REASONING_PROMPT_TEMPLATE.safe_substitute(ocr_text=safe_ocr_text)
    return await gemma_text(
        "ocr_reasoning",
        prompt,
        max_tokens=512,
    )


# ============================================================================
# JSON 解析
# ============================================================================


def _parse_llm_response(response: str) -> LabelOcrResult:
    """
    從 LLM 回應中提取 JSON。

    Gemma / DeepSeek 有時會在 JSON 外面包 markdown code fence，
    需要把外殼剝掉再 json.loads()。
    解析失敗時回傳空的 LabelOcrResult。
    """
    text = response.strip()

    # 剝掉 markdown code fence
    if "```" in text:
        parts = text.split("```")
        for part in parts:
            part = part.strip()
            if part.startswith("json"):
                part = part[4:].strip()
            if part.startswith("{"):
                text = part
                break

    # 找到 JSON 物件
    if not text.startswith("{"):
        idx = text.find("{")
        if idx != -1:
            text = text[idx:]

    idx = text.rfind("}")
    if idx != -1:
        text = text[: idx + 1]

    try:
        parsed = json.loads(text)
        # 把 "unknown"/"N/A"/空字串 轉成 None
        for key in parsed:
            if isinstance(parsed[key], str) and parsed[key].lower() in (
                "unknown",
                "n/a",
                "none",
                "",
            ):
                parsed[key] = None
        return LabelOcrResult(**parsed)
    except (json.JSONDecodeError, TypeError) as e:
        logger.warning("Failed to parse OCR response as JSON: %s\nRaw: %s", e, response)
        return LabelOcrResult()


# ============================================================================
# Stage 3: Regex 後處理（補漏 + 清理）
# ============================================================================


def _postprocess(result: LabelOcrResult, raw_text: str) -> LabelOcrResult:
    """
    用 regex 補強 LLM 的結果。

    【分工原則】
    - LLM 負責「理解」（品牌、商品名、尺碼）→ 需要語意
    - Regex 負責「格式」（貨號、材質清理）→ 格式固定，regex 更可靠
    """
    data = result.model_dump()

    # --- 補 size（LLM 可能只抓到 S，漏掉完整的 170/96A(S)）---
    asian_size = re.search(r"(\d{3}/\d{2,3}[A-Z]?\([A-Z]+\))", raw_text)
    if asian_size:
        data["size"] = asian_size.group(1)

    # --- 驗證 brand_raw（不該是亂碼或非英文）---
    if data["brand_raw"] and not re.match(r"^[A-Za-z\s&\'.\-]+$", data["brand_raw"]):
        data["brand_raw"] = None

    # --- 驗證 product_name（不該包含非 ASCII 字元）---
    if data["product_name"] and not re.match(
        r"^[A-Za-z0-9\s\-\.\']+$", data["product_name"]
    ):
        data["product_name"] = None

    # --- 修正 product_code（取最精確的匹配）---
    regex_code = _regex_product_code(raw_text)
    if regex_code:
        data["product_code"] = regex_code

    # --- 驗證 country_of_origin（必須來自 "Made in" 文字）---
    if data["country_of_origin"] and "made in" not in raw_text.lower():
        data["country_of_origin"] = None

    # --- 清理 material（去除多語言翻譯）---
    if data["material"]:
        data["material"] = _clean_material(data["material"])

    return LabelOcrResult(**data)


def _regex_product_code(text: str) -> str | None:
    """
    用 regex 從文字中抓商品貨號。

    貨號格式：字母開頭+6位以上數字、字母數字混合+破折號、純數字6位以上。
    """
    patterns = [
        r"\b([A-Z]\d{6,})\b",
        r"\b([A-Z]{1,3}\d{3,}-[A-Z0-9]{2,})\b",
        r"(?<!/)\b(\d{6,})\b(?!/)",
    ]
    for pattern in patterns:
        match = re.search(pattern, text)
        if match:
            return match.group(1)
    return None


def _clean_material(material: str) -> str:
    """
    清理材質字串，只保留英文部分。
    """
    pct_match = re.search(r"(\d{1,3}%\s*[A-Za-z]+)", material)
    if pct_match:
        return pct_match.group(1).strip()

    if "/" in material:
        return material.split("/")[0].strip()

    return material.strip()