"""Offline worker for the full official PaddleOCR-VL layout pipeline.

The caller sends JPEG bytes on stdin. This process uses explicit local model
directories; it writes a mode-0600 temporary file and deletes it afterward.
No private images or OCR results are saved to the model cache or logs.
"""

from __future__ import annotations

import json
import os
from html.parser import HTMLParser
from pathlib import Path
import re
import sys
import tempfile

from PIL import Image


ROOT = Path(os.environ.get("OCR_MODEL_HOME", Path(__file__).resolve().parents[1] / ".cache")).expanduser().resolve()
os.environ["PADDLE_PDX_CACHE_HOME"] = str(ROOT)
os.environ["HF_HOME"] = str(ROOT / "hf")
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "True"
os.environ["OMP_NUM_THREADS"] = "1"


class _TableRows(HTMLParser):
    def __init__(self):
        super().__init__()
        self.rows = []
        self.row = None
        self.cell = None

    def handle_starttag(self, tag, attrs):
        if tag == "tr":
            self.row = []
        elif tag in ("td", "th") and self.row is not None:
            self.cell = []
        elif tag == "br" and self.cell is not None:
            self.cell.append(" ")

    def handle_data(self, data):
        if self.cell is not None:
            self.cell.append(data)

    def handle_endtag(self, tag):
        if tag in ("td", "th") and self.cell is not None and self.row is not None:
            self.row.append(" ".join("".join(self.cell).split())[:500])
            self.cell = None
        elif tag == "tr" and self.row is not None:
            if self.row:
                self.rows.append(self.row)
            self.row = None


def _table_rows(content: str) -> list[list[str]]:
    parser = _TableRows()
    parser.feed(content)
    parser.close()
    return parser.rows


def _box(value, width: int, height: int, offset=(0, 0)):
    if value is None or len(value) != 4:
        return None
    try:
        x1, y1, x2, y2 = (float(item) for item in value)
    except (TypeError, ValueError):
        return None
    if not (0 <= x1 < x2 <= width and 0 <= y1 < y2 <= height):
        return None
    ox, oy = offset
    return [round(x1 + ox, 1), round(y1 + oy, 1),
            round(x2 + ox, 1), round(y2 + oy, 1)]


def extract(payload: dict, offset=(0, 0)) -> dict:
    """Keep original block text and its actual layout region, never fake line boxes."""
    width, height = int(payload["width"]), int(payload["height"])
    lines, regions = [], []
    for block in payload.get("parsing_res_list", []):
        box = _box(block.get("block_bbox"), width, height, offset)
        label = str(block.get("block_label", "region"))[:80]
        content = block.get("block_content", "")
        if not isinstance(content, str):
            content = str(content)
        if box:
            regions.append({"label": label, "box": box})
        if label == "table" and "<tr" in content.lower():
            for cells in _table_rows(content):
                text = " | ".join(cells)
                if text.strip():
                    lines.append({"text": text[:1000], "cells": cells,
                                  "box": box, "source": "PaddleOCR-VL-Tabellenzeile; nur Tabellenregion als Bildbeleg"})
        else:
            for part in content.splitlines():
                text = part.strip()
                if text:
                    lines.append({"text": text[:1000], "box": box,
                                  "source": f"PaddleOCR-VL-Layoutblock {label}; keine eigene Zeilenbox"})
    return {
        "model": "PaddleOCR-VL 1.6 · vollständige lokale Layoutpipeline",
        "lines": lines[:500], "regions": regions[:200],
        "layout_note": "Vision-Bildbelege markieren echte Layoutregionen. Mehrere Textzeilen können dieselbe Region teilen; die Region ist keine präzise Feldbox.",
    }


def _retry_table(block: dict, width: int, height: int) -> bool:
    if block.get("block_label") != "table":
        return False
    box = _box(block.get("block_bbox"), width, height)
    if not box:
        return False
    x1, y1, x2, y2 = box
    area = (x2 - x1) * (y2 - y1) / (width * height)
    content = block.get("block_content", "")
    if not isinstance(content, str):
        return False
    visible_text = re.sub(r"<[^>]*>", "", content).strip()
    return area > .45 and len(visible_text) < 80


def _split_table(box: list[float], width: int, height: int) -> list[tuple[int, int, int, int]]:
    x1, y1, x2, y2 = box
    left = max(0, int(x1) - 20)
    top = max(0, int(y1) - 20)
    right = min(width, int(x2) + 20)
    bottom = min(height, int(y2) + 20)
    middle = (left + right) // 2
    overlap = min(120, max(40, round((right - left) * .045)))
    return [(left, top, min(right, middle + overlap), bottom),
            (max(left, middle - overlap), top, right, bottom)]


def _missing_neighbor_box(block: dict, blocks: list[dict], width: int, height: int):
    """Find an uncovered page column directly left of a detected right table."""
    if block.get("block_label") != "table":
        return None
    box = _box(block.get("block_bbox"), width, height)
    if not box:
        return None
    x1, y1, x2, y2 = box
    table_width, table_height = x2 - x1, y2 - y1
    if x1 < .52 * width or table_width > .4 * width or table_height < .2 * height:
        return None
    left = max(0, round(x1 - 1.2 * table_width))
    right = min(width, round(x1 + .07 * table_width))
    top = max(0, round(y1 - .1 * table_height))
    bottom = min(height, round(y2 + .3 * table_height))
    if right - left < 300 or bottom - top < 300:
        return None
    for other in blocks:
        if other is block or other.get("block_label") != "table":
            continue
        other_box = _box(other.get("block_bbox"), width, height)
        if not other_box:
            continue
        a, b, c, d = other_box
        overlap = max(0, min(x1, c) - max(left, a)) * max(0, min(bottom, d) - max(top, b))
        if overlap > .15 * (x1 - left) * (bottom - top):
            return None
    return (left, top, right, bottom)


def _run_tile(pipeline, image: Image.Image, crop_box, private: Path):
    descriptor, name = tempfile.mkstemp(prefix="table-", suffix=".jpg", dir=private)
    tile_path = Path(name)
    try:
        with os.fdopen(descriptor, "wb") as stream:
            image.crop(crop_box).save(stream, format="JPEG", quality=92)
        results = list(pipeline.predict(str(tile_path)))
        if len(results) != 1:
            return None
        payload = results[0].json
        return extract(payload.get("res", payload), crop_box[:2])
    except Exception:
        # Keep the original page result if a diagnostic retry fails.
        return None
    finally:
        tile_path.unlink(missing_ok=True)


def predict_tiled(pipeline, path: Path, private: Path) -> dict:
    """Generic full-page coverage for a wide page whose full parsing timed out."""
    with Image.open(path) as image:
        width, height = image.size
        boxes = [(0, 0, round(width * .34), height),
                 (round(width * .25), 0, round(width * .66), height),
                 (round(width * .61), 0, width, height)]
        tiles = [_run_tile(pipeline, image, box, private) for box in boxes]
    valid = [tile for tile in tiles if tile and tile["lines"]]
    if not valid:
        raise RuntimeError("tile_inference_failed")
    result = extract({"width": width, "height": height, "parsing_res_list": []})
    for tile in valid:
        result["lines"].extend(tile["lines"])
        result["regions"].extend(tile["regions"])
    result["lines"] = result["lines"][:500]
    result["regions"] = result["regions"][:200]
    result["retry_used"] = True
    result["layout_note"] = (
        f"Die Gesamtansicht brauchte zu lange. {len(valid)} von 3 überlappenden "
        "Ansichten derselben Seite wurden lokal gelesen; widersprüchliche "
        "Feldwerte bleiben ohne Zuordnung. Tabellenbelege sind grobe Regionen.")
    return result


def predict_adaptive(pipeline, path: Path, private: Path) -> dict:
    results = list(pipeline.predict(str(path)))
    if len(results) != 1:
        raise RuntimeError("invalid_page_count")
    original = results[0].json
    payload = original.get("res", original)
    width, height = int(payload["width"]), int(payload["height"])
    blocks = payload.get("parsing_res_list", [])
    large_table = next((block for block in blocks if _retry_table(block, width, height)), None)
    neighbor_box = next((box for block in blocks
                         if (box := _missing_neighbor_box(block, blocks, width, height))), None)
    if large_table is None and neighbor_box is None:
        return extract(payload)

    # A large, almost empty table is a measurable layout/scale failure. Retry
    # only that detected region in two smaller views; all pixels still come
    # from the same upload and all coordinates map back to the common image.
    tiles = []
    large_recovered = False
    neighbor_recovered = False
    with Image.open(path) as image:
        if large_table is not None:
            box = _box(large_table["block_bbox"], width, height)
            for crop_box in _split_table(box, width, height):
                parsed = _run_tile(pipeline, image, crop_box, private)
                if parsed and parsed["lines"]:
                    tiles.append(parsed)
                    large_recovered = True
        if neighbor_box is not None:
            parsed = _run_tile(pipeline, image, neighbor_box, private)
            if parsed and parsed["lines"]:
                tiles.append(parsed)
                neighbor_recovered = True
    if not tiles:
        return extract(payload)
    kept = [block for block in blocks if block is not large_table] if large_recovered else blocks
    result = extract({**payload, "parsing_res_list": kept})
    for tile in tiles:
        result["lines"].extend(tile["lines"])
        result["regions"].extend(tile["regions"])
    result["lines"] = result["lines"][:500]
    result["regions"] = result["regions"][:200]
    result["retry_used"] = True
    notes = []
    if large_recovered:
        notes.append("Ein übergroßer, fast leerer Tabellenblock wurde in zwei kleineren Ansichten erneut gelesen.")
    if neighbor_recovered:
        notes.append("Eine von der Layoutstufe ausgelassene Nachbarspalte wurde zusätzlich gelesen.")
    notes.append("Tabellenzeilen haben weiterhin nur einen groben Regionsbeleg.")
    result["layout_note"] = " ".join(notes)
    return result


def main() -> None:
    data = sys.stdin.buffer.read(16 * 1024 * 1024 + 1)
    if not data or len(data) > 16 * 1024 * 1024:
        print(json.dumps({"error": "invalid_input"}))
        return
    models = ROOT / "official_models"
    layout = models / "PP-DocLayoutV3"
    vlm = models / "PaddleOCR-VL-1.6"
    if not (layout / "inference.pdiparams").is_file() or not (vlm / "model.safetensors").is_file():
        print(json.dumps({"error": "models_missing"}))
        return
    private = Path(os.environ.get("OCR_PRIVATE_DIR", ROOT / "private"))
    private.mkdir(mode=0o700, exist_ok=True)
    descriptor, name = tempfile.mkstemp(prefix="page-", suffix=".jpg", dir=private)
    path = Path(name)
    try:
        with os.fdopen(descriptor, "wb") as stream:
            stream.write(data)
        # Paddle logs should never expose text or private file paths through
        # the SSH result channel. The app suppresses stderr as well.
        stdout_copy = os.dup(1)
        null_fd = os.open(os.devnull, os.O_WRONLY)
        os.dup2(null_fd, 1)
        os.close(null_fd)
        try:
            from paddleocr import PaddleOCRVL
            pipeline = PaddleOCRVL(pipeline_version="v1.6", use_layout_detection=True,
                                   layout_detection_model_dir=str(layout), vl_rec_model_dir=str(vlm),
                                   use_doc_orientation_classify=False,
                                   use_doc_unwarping=False,
                                   device=os.environ.get("OCR_VISION_DEVICE", "cpu"))
            result = (predict_tiled(pipeline, path, private)
                      if "--tiles-only" in sys.argv[1:] else
                      predict_adaptive(pipeline, path, private))
        finally:
            os.dup2(stdout_copy, 1)
            os.close(stdout_copy)
        print(json.dumps(result, ensure_ascii=False, separators=(",", ":")))
    except Exception:
        # No exception text: model libraries may include recognized content.
        print(json.dumps({"error": "inference_failed"}))
    finally:
        path.unlink(missing_ok=True)


if __name__ == "__main__":
    main()