from __future__ import annotations

import os
import tempfile
import logging
from collections.abc import Iterable
from dataclasses import dataclass, replace
from numbers import Real
from pathlib import Path
from typing import Any

from core.models import Bounds, SceneElement

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class OCRBox:
    text: str
    bounds: Bounds
    confidence: float | None = None
    foreground_color: str | None = None
    background_color: str | None = None

    def to_scene_element(self, element_id: str) -> SceneElement:
        return SceneElement(
            id=element_id,
            type="text",
            text=self.text,
            bounds=self.bounds,
            confidence=self.confidence,
            source="ocr",
            foreground_color=self.foreground_color,
            background_color=self.background_color,
        )


class PaddleOCREngine:
    def __init__(self, **kwargs: Any) -> None:
        # Device screenshots are already upright, flat digital captures, not
        # photographed paper documents, so PaddleOCR's document-preprocessing
        # models (orientation classification + UVDoc unwarping) have nothing
        # real to correct. Left at their library defaults (True), they still
        # run, geometrically warp the image, and detect/recognize text against
        # that warped image, returning box coordinates that no longer line up
        # with the original screenshot. Callers can still opt back in via an
        # explicit kwarg.
        kwargs.setdefault("use_doc_orientation_classify", False)
        kwargs.setdefault("use_doc_unwarping", False)
        self.kwargs = kwargs
        self._engine: Any | None = None

    def extract(self, image: bytes | str | Path) -> list[OCRBox]:
        engine = self._load()
        image_input, temp_path = _image_input(image)
        try:
            if hasattr(engine, "predict"):
                raw = engine.predict(input=image_input)
            else:
                raw = engine.ocr(image_input, cls=True)
            boxes = parse_paddle_result(raw)
            return _with_sampled_colors(boxes, image_input)
        finally:
            if temp_path:
                temp_path.unlink(missing_ok=True)

    def _load(self) -> Any:
        if self._engine is None:
            from paddleocr import PaddleOCR

            self._engine = PaddleOCR(**self.kwargs)
        return self._engine


def run_ocr(
    image: bytes | str | Path,
    *,
    engine: PaddleOCREngine | None = None,
    strict: bool = False,
) -> list[SceneElement]:
    try:
        boxes = (engine or PaddleOCREngine()).extract(image)
    except Exception:
        if strict:
            raise
        logger.warning("OCR extraction failed; continuing without OCR", exc_info=True)
        return []
    return [box.to_scene_element(f"ocr-{index:03d}") for index, box in enumerate(boxes)]


def _with_sampled_colors(boxes: list[OCRBox], image_path: str) -> list[OCRBox]:
    if not boxes:
        return boxes
    try:
        from PIL import Image

        with Image.open(image_path) as source:
            picture = source.convert("RGB")
            sampled = []
            for box in boxes:
                foreground, background = _estimate_colors(picture, box.bounds)
                sampled.append(
                    replace(
                        box,
                        foreground_color=foreground,
                        background_color=background,
                    )
                )
            return sampled
    except Exception:
        logger.warning(
            "OCR color sampling failed; continuing without foreground/background colors",
            exc_info=True,
        )
        return boxes


def _estimate_colors(picture: Any, bounds: Bounds) -> tuple[str | None, str | None]:
    left = max(0, int(bounds.x))
    top = max(0, int(bounds.y))
    right = min(picture.width, int(round(bounds.right)))
    bottom = min(picture.height, int(round(bounds.bottom)))
    if right - left < 2 or bottom - top < 2:
        return (None, None)

    pixels = list(picture.crop((left, top, right, bottom)).get_flattened_data())
    threshold = _otsu_threshold(pixels)
    if threshold is None:
        return (None, None)

    dark = [pixel for pixel in pixels if _luminance(pixel) <= threshold]
    light = [pixel for pixel in pixels if _luminance(pixel) > threshold]
    if not dark or not light:
        return (None, None)

    # Text strokes normally cover a minority of the box's pixels regardless of
    # whether the text is dark-on-light or light-on-dark, so the smaller of
    # the two luminance clusters is treated as the foreground (text) color.
    foreground, background = (dark, light) if len(dark) <= len(light) else (light, dark)
    return _average_hex(foreground), _average_hex(background)


def _luminance(pixel: tuple[int, int, int]) -> float:
    r, g, b = pixel
    return 0.299 * r + 0.587 * g + 0.114 * b


def _otsu_threshold(pixels: list[tuple[int, int, int]]) -> float | None:
    total = len(pixels)
    if total == 0:
        return None

    histogram = [0] * 256
    for pixel in pixels:
        histogram[int(_luminance(pixel))] += 1

    sum_total = sum(level * count for level, count in enumerate(histogram))
    sum_background = 0.0
    weight_background = 0
    best_variance = -1.0
    threshold = 0
    for level, count in enumerate(histogram):
        weight_background += count
        if weight_background == 0:
            continue
        weight_foreground = total - weight_background
        if weight_foreground == 0:
            break
        sum_background += level * count
        mean_background = sum_background / weight_background
        mean_foreground = (sum_total - sum_background) / weight_foreground
        variance = (
            weight_background
            * weight_foreground
            * (mean_background - mean_foreground) ** 2
        )
        if variance > best_variance:
            best_variance = variance
            threshold = level
    return float(threshold)


def _average_hex(pixels: list[tuple[int, int, int]]) -> str:
    count = len(pixels)
    r = sum(pixel[0] for pixel in pixels) // count
    g = sum(pixel[1] for pixel in pixels) // count
    b = sum(pixel[2] for pixel in pixels) // count
    return f"#{r:02x}{g:02x}{b:02x}"


def parse_paddle_result(raw: Any) -> list[OCRBox]:
    boxes: list[OCRBox] = []
    for item in _flatten_pages(raw):
        parsed = _parse_line(item)
        if parsed:
            boxes.append(parsed)
    return boxes


def _image_input(image: bytes | str | Path) -> tuple[str, Path | None]:
    if isinstance(image, bytes):
        handle = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        try:
            handle.write(image)
        finally:
            handle.close()
        return handle.name, Path(handle.name)
    return os.fspath(image), None


def _flatten_pages(raw: Any) -> Iterable[Any]:
    if raw is None:
        return []
    if isinstance(raw, dict):
        return _dict_lines(raw)
    if isinstance(raw, list):
        flattened: list[Any] = []
        for page in raw:
            if isinstance(page, dict):
                flattened.extend(_dict_lines(page))
            elif _looks_like_ocr_line(page):
                flattened.append(page)
            elif isinstance(page, list):
                flattened.extend(page)
        return flattened
    json_attr = getattr(raw, "json", None)
    if callable(json_attr):
        return _flatten_pages(json_attr())
    if isinstance(json_attr, (dict, list)):
        return _flatten_pages(json_attr)
    return []


def _dict_lines(data: dict[str, Any]) -> list[Any]:
    result = data.get("res")
    payload = result if isinstance(result, dict) else data
    texts = _first_nonempty(payload, "rec_texts", "texts")
    scores = _first_nonempty(payload, "rec_scores", "scores")
    boxes = _first_nonempty(payload, "rec_boxes", "rec_polys", "dt_polys")
    return [
        {
            "text": text,
            "confidence": scores[index] if index < len(scores) else None,
            "points": boxes[index] if index < len(boxes) else None,
        }
        for index, text in enumerate(texts)
    ]


def _looks_like_ocr_line(value: Any) -> bool:
    return isinstance(value, (list, tuple)) and len(value) >= 2


def _parse_line(line: Any) -> OCRBox | None:
    if isinstance(line, dict):
        text = line.get("text")
        points = _first_nonempty(line, "points", "box", "bounds")
        confidence = line.get("confidence")
        if not _has_items(text) or not _has_items(points):
            return None
        return OCRBox(
            str(text), _bounds_from_points(points), _float_or_none(confidence)
        )

    if not _looks_like_ocr_line(line):
        return None

    points = line[0]
    text_payload = line[1]
    if isinstance(text_payload, (list, tuple)) and text_payload:
        text = text_payload[0]
        confidence = text_payload[1] if len(text_payload) > 1 else None
    else:
        text = text_payload
        confidence = None
    if not text:
        return None
    return OCRBox(str(text), _bounds_from_points(points), _float_or_none(confidence))


def _bounds_from_points(points: Any) -> Bounds:
    if isinstance(points, dict):
        return Bounds.from_dict(points)
    if _has_length(points, 4) and all(isinstance(value, Real) for value in points):
        x1, y1, x2, y2 = [float(value) for value in points]
        return Bounds(x1, y1, x2 - x1, y2 - y1)

    xs: list[float] = []
    ys: list[float] = []
    for point in points:
        if isinstance(point, dict):
            xs.append(float(point["x"]))
            ys.append(float(point["y"]))
        else:
            xs.append(float(point[0]))
            ys.append(float(point[1]))
    return Bounds(min(xs), min(ys), max(xs) - min(xs), max(ys) - min(ys))


def _float_or_none(value: Any) -> float | None:
    if value is None:
        return None
    return float(value)


def _first_nonempty(data: dict[str, Any], *keys: str) -> Any:
    for key in keys:
        value = data.get(key)
        if _has_items(value):
            return value
    return []


def _has_items(value: Any) -> bool:
    if value is None:
        return False
    try:
        return len(value) > 0
    except TypeError:
        return bool(value)


def _has_length(value: Any, length: int) -> bool:
    try:
        return len(value) == length
    except TypeError:
        return False