"""PaddleOCR wrapper.

PaddleOCR has a heavy initialization cost (~2-5s on CPU as model files load),
so we keep a process-global instance behind a lazy accessor.

The wrapper exposes a small, stable surface so the rest of the pipeline does
not depend directly on paddleocr's evolving API.
"""

from __future__ import annotations

from dataclasses import dataclass
from threading import Lock
from typing import TYPE_CHECKING

import numpy as np

from ocr_sprint.config import get_settings
from ocr_sprint.pipeline.ingest import NDArrayU8
from ocr_sprint.utils.logging import get_logger

if TYPE_CHECKING:
    from paddleocr import PaddleOCR

_logger = get_logger(__name__)
_lock = Lock()
_instance: PaddleOCR | None = None


@dataclass(frozen=True)
class OCRLine:
    """One recognized line with its bounding polygon and confidence."""

    text: str
    confidence: float
    box: tuple[tuple[float, float], ...]  # 4 (x, y) corner points


def _line_y_center(line: OCRLine) -> float:
    return sum(p[1] for p in line.box) / len(line.box)


def _line_x_left(line: OCRLine) -> float:
    return min(p[0] for p in line.box)


def _line_height(line: OCRLine) -> float:
    ys = [p[1] for p in line.box]
    return max(ys) - min(ys)


def sort_lines_by_layout(lines: list[OCRLine]) -> list[OCRLine]:
    """Reorder lines into top-to-bottom, left-to-right reading order.

    PaddleOCR's natural output order reflects detection order, not visual
    layout. On dense tables (e.g. Polda Kalbar Akpol-panitia sprint) this
    interleaves rows and columns — Paddle may emit a row's KET column
    before its NAMA column, breaking every downstream extractor that
    assumes top-to-bottom row order.

    We rebuild reading order by:

    1. Sorting by ``y_center``.
    2. Grouping consecutive lines into row-bands when their ``y_center``
       differs by less than half the median line height (so visually
       same-row cells stay together even when their boxes don't perfectly
       align).
    3. Sorting each band left-to-right by ``x_left``.
    """
    if not lines:
        return []

    heights = [_line_height(ln) for ln in lines if _line_height(ln) > 0]
    if not heights:
        return list(lines)
    median_height = sorted(heights)[len(heights) // 2]
    band_threshold = max(1.0, median_height * 0.5)

    by_y = sorted(lines, key=_line_y_center)
    bands: list[list[OCRLine]] = []
    current_band: list[OCRLine] = []
    current_y: float | None = None
    for ln in by_y:
        y = _line_y_center(ln)
        if current_y is None or abs(y - current_y) <= band_threshold:
            current_band.append(ln)
            # Track the band's running y-center as the mean of its
            # members so a slowly-drifting set of cells doesn't split
            # mid-row.
            current_y = (
                sum(_line_y_center(b) for b in current_band) / len(current_band)
            )
        else:
            bands.append(current_band)
            current_band = [ln]
            current_y = y
    if current_band:
        bands.append(current_band)

    ordered: list[OCRLine] = []
    for band in bands:
        ordered.extend(sorted(band, key=_line_x_left))
    return ordered


@dataclass(frozen=True)
class OCRPage:
    """OCR output for a single page."""

    lines: list[OCRLine]

    @property
    def text(self) -> str:
        """Reconstruct page text in visual reading order (top-to-bottom, left-to-right)."""
        return "\n".join(line.text for line in sort_lines_by_layout(self.lines))

    @property
    def mean_confidence(self) -> float:
        if not self.lines:
            return 0.0
        return float(np.mean([line.confidence for line in self.lines]))


def _build_paddleocr() -> PaddleOCR:
    s = get_settings()
    if s.ocr_use_gpu:
        from ocr_sprint.utils.gpu import configure_nvidia_dll_path

        configure_nvidia_dll_path()

    from paddleocr import PaddleOCR

    kwargs: dict[str, object] = {
        "lang": s.ocr_lang,
        "use_angle_cls": True,
        "use_gpu": s.ocr_use_gpu,
        "show_log": False,
    }
    if s.ocr_det_model_dir:
        kwargs["det_model_dir"] = s.ocr_det_model_dir
    if s.ocr_rec_model_dir:
        kwargs["rec_model_dir"] = s.ocr_rec_model_dir
    if s.ocr_cls_model_dir:
        kwargs["cls_model_dir"] = s.ocr_cls_model_dir
    _logger.info("paddleocr.init", lang=s.ocr_lang, use_gpu=s.ocr_use_gpu)
    return PaddleOCR(**kwargs)


def get_ocr() -> PaddleOCR:
    """Lazy, thread-safe singleton accessor for the PaddleOCR engine."""
    global _instance
    if _instance is None:
        with _lock:
            if _instance is None:
                _instance = _build_paddleocr()
    return _instance


def warmup() -> None:
    """Eagerly initialize the PaddleOCR engine.

    Call this during application startup so the first real request does not
    pay the model-loading cost (~2-5s on CPU). Also prevents the process from
    entering Disk-Sleep state (state D) mid-request when memory is tight,
    because the OS has already paged in all model weights during startup.
    """
    _logger.info("paddleocr.warmup.start")
    get_ocr()
    _logger.info("paddleocr.warmup.done")


def run_ocr(image: NDArrayU8) -> OCRPage:
    """Run OCR on a single BGR image and return a structured page result."""
    engine = get_ocr()
    raw = engine.ocr(image, cls=True)
    # PaddleOCR returns [[ [box, (text, conf)], ... ]] — one outer list per image.
    if not raw or raw[0] is None:
        return OCRPage(lines=[])
    page_raw = raw[0]
    lines: list[OCRLine] = []
    for item in page_raw:
        if not item or len(item) < 2:
            continue
        box_raw, text_conf = item[0], item[1]
        text, conf = text_conf[0], float(text_conf[1])
        try:
            box = tuple((float(p[0]), float(p[1])) for p in box_raw)
        except (TypeError, ValueError, IndexError):
            continue
        lines.append(OCRLine(text=text, confidence=conf, box=box))
    return OCRPage(lines=lines)
