#!/usr/bin/env python3
"""PaddleOCR-VL client: parallel layout crops + merge.

Proven on RTX 5090 + vLLM next to Qwen3-4B. Prompt is ``OCR:``, temperature 0.
"""
from __future__ import annotations

import base64
import io
import re
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path

from openai import OpenAI
from PIL import Image

DEFAULT_BASE_URL = "http://127.0.0.1:18000/v1"
DEFAULT_MODEL = "PaddlePaddle/PaddleOCR-VL"
PROMPT = "OCR:"
_MIN_SIDE = 420

_REPEAT_TOKEN = re.compile(r"(\S+)(?:\s+\1){6,}")
_REPEAT_CHAR = re.compile(r"(.{6,}?)\1{5,}", re.DOTALL)
_REFUSAL = re.compile(
    r"(not fully clear|cannot extract|cannot准确|无法|没有文字)",
    re.I,
)


@dataclass
class RegionResult:
    name: str
    text: str
    latency_s: float


def _to_data_url(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.convert("RGB").save(buf, format="PNG", optimize=True)
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii")


def _pad_min(img: Image.Image, min_side: int = _MIN_SIDE) -> Image.Image:
    w, h = img.size
    if w >= min_side and h >= min_side:
        return img
    nw, nh = max(w, min_side), max(h, min_side)
    canvas = Image.new("RGB", (nw, nh), (255, 255, 255))
    canvas.paste(img, ((nw - w) // 2, (nh - h) // 2))
    return canvas


def _band(img: Image.Image, top: float, bot: float, left: float = 0.0, right: float = 1.0) -> Image.Image:
    w, h = img.size
    box = (
        int(w * left),
        int(h * top),
        max(int(w * left) + 8, int(w * right)),
        max(int(h * top) + 8, int(h * bot)),
    )
    return img.crop(box)


def regions_for(img: Image.Image) -> list[tuple[str, Image.Image]]:
    w, h = img.size
    out: list[tuple[str, Image.Image]] = [
        ("full", img),
        ("header", _pad_min(_band(img, 0.0, 0.16))),
        ("top_right", _pad_min(_band(img, 0.0, 0.22, 0.48, 1.0))),
        ("right", _pad_min(_band(img, 0.0, 0.62, 0.52, 1.0))),
        ("totals", _pad_min(_band(img, 0.34, 0.72))),
    ]
    if w < 900:
        out.append(("masthead", img.crop((0, 0, w, min(h, 130)))))
    return out


def strip_collapse(text: str) -> str:
    if not text:
        return text
    m = _REPEAT_TOKEN.search(text)
    if m:
        text = text[: m.start()] + m.group(1)
    m = _REPEAT_CHAR.search(text)
    if m:
        text = text[: m.start()] + m.group(1)
    lines = text.splitlines()
    cleaned: list[str] = []
    run = 0
    prev = None
    short_run = 0
    for ln in lines:
        key = ln.strip()
        if len(key) <= 2:
            short_run += 1
            if short_run >= 8:
                break
            continue
        short_run = 0
        if key and key == prev:
            run += 1
            if run >= 6:
                break
            continue
        prev = key
        run = 0
        cleaned.append(ln)
    return "\n".join(cleaned).strip()


def is_junk_region(text: str) -> bool:
    if not text or not text.strip():
        return True
    if _REFUSAL.search(text):
        return True
    cjk = sum(1 for ch in text if "\u4e00" <= ch <= "\u9fff")
    return cjk > 12


class OcrEngine:
    def __init__(
        self,
        base_url: str = DEFAULT_BASE_URL,
        model: str = DEFAULT_MODEL,
        timeout: float = 180.0,
    ) -> None:
        self.model = model
        self.client = OpenAI(api_key="EMPTY", base_url=base_url, timeout=timeout)

    def _complete(self, img: Image.Image, *, max_tokens: int, penalty: bool) -> str:
        extra = {"repetition_penalty": 1.05}
        if penalty:
            extra = {"repetition_penalty": 1.12, "frequency_penalty": 0.22}
        resp = self.client.chat.completions.create(
            model=self.model,
            temperature=0.0,
            max_tokens=max_tokens,
            extra_body=extra,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": _to_data_url(img)}},
                        {"type": "text", "text": PROMPT},
                    ],
                }
            ],
        )
        return strip_collapse(resp.choices[0].message.content or "")

    def ocr_image(self, image: Image.Image | Path | str) -> tuple[str, float, list[RegionResult]]:
        if not isinstance(image, Image.Image):
            image = Image.open(image).convert("RGB")
        else:
            image = image.convert("RGB")
        jobs = regions_for(image)
        t0 = time.perf_counter()
        results: list[RegionResult] = []

        def run(name: str, crop: Image.Image) -> RegionResult:
            t1 = time.perf_counter()
            light = name in {"full", "masthead"}
            text = self._complete(
                crop,
                max_tokens=1024 if name == "full" else 512,
                penalty=not light,
            )
            return RegionResult(name, text, time.perf_counter() - t1)

        with ThreadPoolExecutor(max_workers=len(jobs)) as pool:
            futs = [pool.submit(run, name, crop) for name, crop in jobs]
            for fut in as_completed(futs):
                results.append(fut.result())

        order = {"full": 0, "header": 1, "masthead": 2, "top_right": 3, "right": 4, "totals": 5}
        results.sort(key=lambda r: order.get(r.name, 9))
        parts = []
        for r in results:
            chunk = strip_collapse(r.text)
            if chunk.strip() and not is_junk_region(chunk):
                parts.append(chunk)
        return "\n\n".join(parts), time.perf_counter() - t0, results