io_utils.py

from __future__ import annotations

import base64
import difflib
import io
import json
import re
from typing import Iterable

from PIL import Image


def load_image_bytes_to_jpeg_b64(raw: bytes, *, max_side: int = 2200, quality: int = 92) -> tuple[str, int, int]:
    img = Image.open(io.BytesIO(raw))
    if img.mode not in ("RGB", "L"):
        img = img.convert("RGB")
    w, h = img.size
    longest = max(w, h)
    if longest > max_side:
        scale = max_side / longest
        new_size = (int(w * scale), int(h * scale))
        img = img.resize(new_size, Image.LANCZOS)
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=quality, optimize=True)
    data = buf.getvalue()
    b64 = base64.b64encode(data).decode("ascii")
    return b64, img.size[0], img.size[1]


def data_url(b64: str, mime: str = "image/jpeg") -> str:
    return f"data:{mime};base64,{b64}"


_WORD_RE = re.compile(r"\S+|\s+")


def _split_keep_ws(text: str) -> list[str]:
    return _WORD_RE.findall(text or "")


def diff_words(reference: str, corrected: str) -> list[dict]:
    """Return a per-word diff aligned to `corrected`.

    Each item is {"text": str, "edited": bool}. Whitespace tokens are kept
    as their own entries with edited=False so the frontend can render them
    verbatim.
    """
    ref_tokens = _split_keep_ws(reference)
    cor_tokens = _split_keep_ws(corrected)
    matcher = difflib.SequenceMatcher(a=ref_tokens, b=cor_tokens, autojunk=False)
    out: list[dict] = []
    for tag, _i1, _i2, j1, j2 in matcher.get_opcodes():
        for tok in cor_tokens[j1:j2]:
            if tok.strip() == "":
                out.append({"text": tok, "edited": False})
            else:
                out.append({"text": tok, "edited": tag != "equal"})
    return out


def _normalize_for_cer(s: str) -> str:
    return s.replace("\r\n", "\n").strip()


def compute_cer_wer(reference: str, hypothesis: str) -> dict:
    """Compute CER and WER from reference (= corrected) to hypothesis (= OCR)."""
    ref = _normalize_for_cer(reference)
    hyp = _normalize_for_cer(hypothesis)
    if not ref:
        return {"cer": 0.0, "wer": 0.0, "n_chars": 0, "n_words": 0}
    try:
        import jiwer
        cer = jiwer.cer(ref, hyp)
        wer = jiwer.wer(ref, hyp)
    except Exception:
        cer = _levenshtein(list(ref), list(hyp)) / max(1, len(ref))
        ref_words = ref.split()
        hyp_words = hyp.split()
        wer = _levenshtein(ref_words, hyp_words) / max(1, len(ref_words))
    return {
        "cer": round(float(cer), 4),
        "wer": round(float(wer), 4),
        "n_chars": len(ref),
        "n_words": len(ref.split()),
    }


def _levenshtein(a: list, b: list) -> int:
    if a == b:
        return 0
    if not a:
        return len(b)
    if not b:
        return len(a)
    prev = list(range(len(b) + 1))
    for i, ca in enumerate(a, 1):
        curr = [i] + [0] * len(b)
        for j, cb in enumerate(b, 1):
            cost = 0 if ca == cb else 1
            curr[j] = min(
                curr[j - 1] + 1,
                prev[j] + 1,
                prev[j - 1] + cost,
            )
        prev = curr
    return prev[-1]


def export_icl_jsonl(items: Iterable[dict]) -> str:
    return "\n".join(json.dumps(it, ensure_ascii=False) for it in items)


def _strip_fence(raw: str) -> str:
    text = (raw or "").strip()
    if text.startswith("```"):
        text = text.strip("`")
        if text.lower().startswith("json"):
            text = text[4:]
        text = text.strip()
    return text


def _line_to_text(item, prefer: str = "raw") -> str:
    """Coerce one line item to a plain string.

    Supports str, or dict with one of: 'raw', 'expanded', 'text'.
    `prefer` decides which field wins when several are present.
    """
    if isinstance(item, str):
        return item
    if isinstance(item, dict):
        for key in (prefer, "raw", "text", "expanded", "line", "content"):
            v = item.get(key)
            if isinstance(v, str):
                return v
        return json.dumps(item, ensure_ascii=False)
    return str(item)


def parse_lines_from_model_response(raw: str, mode: str = "lines") -> tuple[list[str], dict | None]:
    """Return (lines_for_display, structured_payload).

    - mode="lines": expect {"lines": [str]}. Structured payload is None.
    - mode="custom_json": parse as a free-form JSON object (whatever shape the user
      requested via their template) and extract a best-effort list of lines for the
      editor; the full dict is returned as structured payload.

    Both modes fall back to plain splitlines for non-JSON replies.
    """
    text = _strip_fence(raw)
    try:
        data = json.loads(text)
    except json.JSONDecodeError:
        data = None

    if mode == "custom_json":
        if isinstance(data, dict):
            return _extract_display_lines(data), data
        if isinstance(data, list):
            return [_line_to_text(it) for it in data], {"lines": data}
        return [ln for ln in text.splitlines() if ln.strip()], None

    # mode == "lines" (default)
    structured: dict | None = None
    if isinstance(data, dict):
        if isinstance(data.get("lines"), list):
            lines = [_line_to_text(it) for it in data["lines"]]
            has_extra = any(k for k in data.keys() if k != "lines")
            non_str_items = any(not isinstance(it, str) for it in data["lines"])
            if has_extra or non_str_items:
                structured = data
            return lines, structured
        if isinstance(data.get("text"), str):
            return data["text"].splitlines(), None
    if isinstance(data, list):
        return [_line_to_text(it) for it in data], None
    return [ln for ln in text.splitlines() if ln.strip()], None


def _extract_display_lines(obj: dict) -> list[str]:
    """Best-effort: walk the JSON, prefer a 'lines' or 'text'/'body' array, else
    collect string leaves so the editor has something to show.
    """
    for key in ("lines", "body", "text", "content"):
        v = obj.get(key)
        if isinstance(v, list):
            return [_line_to_text(it) for it in v]
        if isinstance(v, str):
            return v.splitlines()
    # Fallback: walk all values, collect strings
    out: list[str] = []
    _walk_strings(obj, out)
    return out


def _walk_strings(node, out: list[str]) -> None:
    if isinstance(node, str):
        if node.strip():
            out.append(node)
    elif isinstance(node, list):
        for it in node:
            _walk_strings(it, out)
    elif isinstance(node, dict):
        for v in node.values():
            _walk_strings(v, out)