from __future__ import annotations

from dataclasses import dataclass
import os
from typing import Any, Iterable, Optional, TYPE_CHECKING, cast


class OcrError(RuntimeError):
    """Raised when OCR processing fails."""


@dataclass(frozen=True)
class OcrLine:
    text: str
    score: float
    bbox: list[list[float]]


_OCR_CACHE: dict[tuple[str, str, Optional[str]], object] = {}


if TYPE_CHECKING:
    from typing import Any as PaddleOCRType
else:
    PaddleOCRType = object


def _load_paddleocr() -> type[PaddleOCRType]:
    from paddleocr import PaddleOCR  # type: ignore[import-not-found]

    return PaddleOCR


def _get_ocr_engine(
    *,
    lang: str,
    ocr_version: str,
    device: Optional[str],
) -> object:
    key = (lang, ocr_version, device)
    cached = _OCR_CACHE.get(key)
    if cached is not None:
        return cached
    try:
        PaddleOCR = _load_paddleocr()
    except Exception as exc:
        raise OcrError("PaddleOCR is not available") from exc
    kwargs: dict[str, object] = {"lang": lang, "ocr_version": ocr_version}
    if device:
        kwargs["device"] = device
    try:
        engine = PaddleOCR(**kwargs)
    except Exception as exc:
        raise OcrError(f"PaddleOCR init failed: {exc}") from exc
    _OCR_CACHE[key] = engine
    return engine


def _read_key(obj: object, key: str) -> object:
    if isinstance(obj, dict):
        return obj.get(key)
    getter = getattr(obj, "__getitem__", None)
    if callable(getter):
        try:
            return getter(key)
        except Exception:
            pass
    return getattr(obj, key, None)


def _is_bbox_like(value: object) -> bool:
    if not isinstance(value, (list, tuple)):
        return False
    if len(value) != 4:
        return False
    for point in value:
        if not isinstance(point, (list, tuple)):
            return False
        if len(point) != 2:
            return False
    return True


def _coerce_bbox(value: object) -> list[list[float]]:
    if not _is_bbox_like(value):
        return []
    bbox = cast(list[list[object]], value)
    out: list[list[float]] = []
    for point in bbox:
        if len(point) < 2:
            continue
        x = _coerce_float(point[0])
        y = _coerce_float(point[1])
        out.append([x, y])
    return out


def _coerce_float(value: object) -> float:
    if isinstance(value, (int, float)):
        return float(value)  # type: ignore[arg-type]
    try:
        return float(value)  # type: ignore[arg-type]
    except Exception:
        return 0.0


def _coerce_score(value: object) -> float:
    return _coerce_float(value)


def _parse_v3_page(res: object) -> list[OcrLine]:
    items = _read_key(res, "ocr_result")
    if not isinstance(items, list):
        items = None
    lines: list[OcrLine] = []
    if isinstance(items, list):
        for item in items:
            text = _read_key(item, "text")
            if text is None:
                continue
            bbox = _read_key(item, "bbox")
            score = _coerce_score(_read_key(item, "score"))
            lines.append(OcrLine(text=str(text), score=score, bbox=_coerce_bbox(bbox)))
        return lines

    texts = _read_key(res, "rec_texts")
    scores = _read_key(res, "rec_scores")
    polys = _read_key(res, "rec_polys")
    if isinstance(texts, list) and isinstance(scores, list) and isinstance(polys, list):
        for text, score, poly in zip(texts, scores, polys):
            lines.append(
                OcrLine(
                    text=str(text),
                    score=_coerce_score(score),
                    bbox=_coerce_bbox(poly),
                )
            )
    return lines


def _looks_like_v3_result(result: object) -> bool:
    if not isinstance(result, list) or not result:
        return False
    probe = result[0]
    items = _read_key(probe, "ocr_result")
    if isinstance(items, list):
        return True
    texts = _read_key(probe, "rec_texts")
    polys = _read_key(probe, "rec_polys")
    return isinstance(texts, list) and isinstance(polys, list)


def _looks_like_v2_line(value: object) -> bool:
    if not isinstance(value, (list, tuple)):
        return False
    if len(value) < 2:
        return False
    if not _is_bbox_like(value[0]):
        return False
    payload = value[1]
    if not isinstance(payload, (list, tuple)):
        return False
    return len(payload) >= 2


def _parse_v2_page(page: object) -> list[OcrLine]:
    if not isinstance(page, list):
        return []
    lines: list[OcrLine] = []
    for line in page:
        if not _looks_like_v2_line(line):
            continue
        bbox = _coerce_bbox(line[0])
        payload = line[1]
        text = payload[0]
        score = _coerce_score(payload[1])
        lines.append(OcrLine(text=str(text), score=score, bbox=bbox))
    return lines


def _parse_ocr_result(result: object) -> list[list[OcrLine]]:
    if result is None:
        return []
    if _looks_like_v3_result(result):
        result_list = cast(list[Any], result)
        return [_parse_v3_page(res) for res in result_list]
    if isinstance(result, list):
        if not result:
            return []
        if _looks_like_v2_line(result[0]):
            return [_parse_v2_page(result)]
        pages: list[list[OcrLine]] = []
        for page in result:
            if page is None:
                pages.append([])
                continue
            if _looks_like_v2_line(page[0]) if isinstance(page, list) and page else False:
                pages.append(_parse_v2_page(page))
                continue
            if _looks_like_v2_line(page):
                pages.append(_parse_v2_page([page]))
                continue
            pages.append([])
        return pages
    return []


def _bbox_stats(bbox: list[list[float]]) -> tuple[float, float, float]:
    if not bbox:
        return 0.0, 0.0, 0.0
    xs = [point[0] for point in bbox]
    ys = [point[1] for point in bbox]
    min_x, max_x = min(xs), max(xs)
    min_y, max_y = min(ys), max(ys)
    x_left = min_x
    y_center = (min_y + max_y) / 2.0
    height = max_y - min_y
    return x_left, y_center, height


def _median(values: list[float]) -> float:
    if not values:
        return 0.0
    sorted_vals = sorted(values)
    mid = len(sorted_vals) // 2
    if len(sorted_vals) % 2 == 1:
        return sorted_vals[mid]
    return (sorted_vals[mid - 1] + sorted_vals[mid]) / 2.0


def _order_lines(lines: Iterable[OcrLine]) -> list[OcrLine]:
    stats: list[tuple[OcrLine, float, float, float]] = []
    heights: list[float] = []
    for line in lines:
        x_left, y_center, height = _bbox_stats(line.bbox)
        stats.append((line, x_left, y_center, height))
        if height > 0:
            heights.append(height)
    if not stats:
        return []
    median_height = _median(heights) if heights else 0.0
    row_tol = max(6.0, median_height * 0.6)

    stats.sort(key=lambda item: item[2])
    rows: list[list[tuple[OcrLine, float]]] = []
    current: list[tuple[OcrLine, float]] = []
    row_center: Optional[float] = None

    for line, x_left, y_center, _height in stats:
        if row_center is None or abs(y_center - row_center) <= row_tol:
            current.append((line, x_left))
            if row_center is None:
                row_center = y_center
            else:
                row_center = (row_center * (len(current) - 1) + y_center) / len(current)
        else:
            rows.append(current)
            current = [(line, x_left)]
            row_center = y_center

    if current:
        rows.append(current)

    ordered: list[OcrLine] = []
    for row in rows:
        row.sort(key=lambda item: item[1])
        ordered.extend([line for line, _x in row])
    return ordered


def _pages_to_text(pages: Iterable[list[OcrLine]], min_score: Optional[float]) -> str:
    page_texts: list[str] = []
    for page in pages:
        ordered = _order_lines(page)
        parts = []
        for line in ordered:
            if not line.text:
                continue
            if min_score is not None and line.score < min_score:
                continue
            parts.append(str(line.text))
        page_texts.append("\n".join(parts))
    return "\n\n".join(page_texts)


def _run_ocr(engine: object, input_data: object) -> object:
    engine_obj = cast(Any, engine)
    predict = getattr(engine_obj, "predict", None)
    if callable(predict):
        return predict(input=input_data)
    ocr = getattr(engine_obj, "ocr", None)
    if callable(ocr):
        return ocr(input_data, cls=True)
    raise OcrError("PaddleOCR API does not expose predict or ocr")


def ocr_image(
    path: os.PathLike | str,
    *,
    lang: str,
    ocr_version: str,
    device: Optional[str],
    min_score: Optional[float],
) -> str:
    path_str = os.fspath(path)
    engine = _get_ocr_engine(lang=lang, ocr_version=ocr_version, device=device)
    try:
        result = _run_ocr(engine, path_str)
    except Exception as exc:
        raise OcrError(f"OCR failed: {exc}") from exc
    pages = _parse_ocr_result(result)
    return _pages_to_text(pages, min_score)


def _render_pdf_pages(path: str, *, scale: float) -> list[object]:
    if scale <= 0:
        raise OcrError("OCR render scale must be positive")
    try:
        import pypdfium2 as pdfium
    except Exception as exc:
        raise OcrError("pypdfium2 is required for PDF OCR") from exc
    try:
        pdf = pdfium.PdfDocument(path)
    except Exception as exc:
        raise OcrError(f"Failed to open PDF for OCR: {exc}") from exc
    images: list[object] = []
    for page in pdf:
        try:
            scale_value = int(round(scale))
            bitmap = page.render(scale=scale_value, rev_byteorder=True)
            image = bitmap.to_numpy()
        except Exception as exc:
            raise OcrError(f"Failed to render PDF page for OCR: {exc}") from exc
        finally:
            try:
                page.close()
            except Exception:
                pass
        if hasattr(image, "shape") and len(image.shape) == 3 and image.shape[2] == 4:
            image = image[:, :, :3]
        images.append(image)
    return images


def ocr_pdf(
    path: os.PathLike | str,
    *,
    lang: str,
    ocr_version: str,
    device: Optional[str],
    render_scale: float,
    min_score: Optional[float],
) -> str:
    path_str = os.fspath(path)
    engine = _get_ocr_engine(lang=lang, ocr_version=ocr_version, device=device)
    pages_text: list[list[OcrLine]] = []
    for image in _render_pdf_pages(path_str, scale=render_scale):
        try:
            result = _run_ocr(engine, image)
        except Exception as exc:
            raise OcrError(f"OCR failed: {exc}") from exc
        parsed = _parse_ocr_result(result)
        pages_text.append(parsed[0] if parsed else [])
    return _pages_to_text(pages_text, min_score)
