"""Screenshot -> layout JSON, for Page Designer's "AI Assist" (see
app/routers/ai_assistant.py, target_type "page"). The AI Assistant's own
model (Ollama or a configured cloud endpoint — see Settings.ai_provider) is a
*text* reasoning model (e.g. gpt-oss:120b), not a vision one: handing it raw
image bytes either fails outright or is silently ignored, and there is no
reliable way to tell which from the response alone. So this module does the
"vision" step itself, deterministically:

    screenshot bytes -> OpenCV preprocessing -> OCR (Tesseract or EasyOCR)
    -> a JSON list of {"text", "left", "top", "width", "height"} elements,
       one per detected line/phrase, with the box normalized to 0..1 of the
       image's own width/height (so the model never needs to know the
       screenshot's pixel size to reason about relative position).

That JSON is what actually reaches the model, as plain text alongside the
user's instructions — see ai_assistant.py's send_message. Nothing here talks
to an LLM; this is pure image -> text extraction.
"""

import threading
from dataclasses import dataclass
from typing import Any

import cv2
import numpy as np

from app.core.config import settings

# Tesseract's own per-word confidence is 0-100, or -1 for a non-text row
# (image_to_data returns one row per detected box at every level: page,
# block, paragraph, line, word — the -1 rows are the non-word levels).
_MIN_WORD_CONFIDENCE = 40
# Screenshots of software UI are often small (a cropped panel, a phone-sized
# capture); Tesseract's accuracy drops sharply below ~300 DPI-equivalent text
# height, so a small image is upscaled before OCR. Never shrinks a large one.
_MIN_OCR_DIMENSION_PX = 1400


class OcrUnavailableError(RuntimeError):
    """The configured OCR engine can't run right now — e.g. the `tesseract`
    binary isn't installed/on PATH, or (for EasyOCR) its Python package isn't
    installed. Message is written to be shown to the end user as-is (mirrors
    ai_assistant.py's existing `ollama_error` field's own convention for a
    provider-availability failure)."""


@dataclass
class LayoutElement:
    text: str
    left: float
    top: float
    width: float
    height: float

    def to_json(self) -> dict[str, Any]:
        return {
            "text": self.text,
            "left": round(self.left, 4),
            "top": round(self.top, 4),
            "width": round(self.width, 4),
            "height": round(self.height, 4),
        }


def _decode_image(image_bytes: bytes) -> np.ndarray:
    arr = np.frombuffer(image_bytes, dtype=np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise OcrUnavailableError("Couldn't read the attached file as an image — try a PNG or JPEG screenshot.")
    return img


def _preprocess_for_ocr(img: np.ndarray) -> tuple[np.ndarray, float]:
    """Grayscale + upscale-if-small + adaptive contrast (Otsu threshold) —
    OpenCV's own job here, ahead of OCR. Returns (processed_image, scale)
    where `scale` is how much the image was enlarged, so callers can map
    detected boxes back to the ORIGINAL image's coordinate space."""
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    h, w = gray.shape[:2]
    scale = max(1.0, _MIN_OCR_DIMENSION_PX / max(h, w))
    if scale > 1.0:
        gray = cv2.resize(gray, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_CUBIC)
    # Otsu's threshold picks its own cutoff from the image's own histogram —
    # robust to a screenshot's usual near-flat backgrounds without any
    # hand-tuned brightness constant.
    _, binarized = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return binarized, scale


def _tesseract_extract(img: np.ndarray) -> list[LayoutElement]:
    import pytesseract

    if settings.tesseract_cmd:
        pytesseract.pytesseract.tesseract_cmd = settings.tesseract_cmd

    orig_h, orig_w = img.shape[:2]
    processed, scale = _preprocess_for_ocr(img)

    # settings.ocr_languages is comma-separated (the one format that reads the
    # same for either engine — see _get_easyocr_reader below); Tesseract's own
    # CLI/library convention instead joins language codes with "+".
    tesseract_lang = "+".join(lang.strip() for lang in settings.ocr_languages.split(",") if lang.strip()) or "eng"
    try:
        data = pytesseract.image_to_data(processed, lang=tesseract_lang, output_type=pytesseract.Output.DICT)
    except pytesseract.TesseractNotFoundError as exc:
        raise OcrUnavailableError(
            "Tesseract OCR isn't installed on the server (or its path isn't configured) — "
            "install it and/or set TESSERACT_CMD in the backend's .env, then try again."
        ) from exc

    # Group words into lines by Tesseract's own (block, paragraph, line)
    # grouping — already in reading order (its layout analysis, not raw pixel
    # sort), so lines are emitted in the order Tesseract produced them rather
    # than re-sorted here.
    lines: dict[tuple[int, int, int, int], dict[str, Any]] = {}
    order: list[tuple[int, int, int, int]] = []
    n = len(data["text"])
    for i in range(n):
        text = data["text"][i].strip()
        try:
            conf = float(data["conf"][i])
        except (TypeError, ValueError):
            conf = -1.0
        if not text or conf < _MIN_WORD_CONFIDENCE:
            continue
        key = (data["block_num"][i], data["par_num"][i], data["line_num"][i], data.get("page_num", [0] * n)[i])
        left, top, width, height = data["left"][i], data["top"][i], data["width"][i], data["height"][i]
        if key not in lines:
            lines[key] = {"words": [text], "left": left, "top": top, "right": left + width, "bottom": top + height}
            order.append(key)
        else:
            box = lines[key]
            box["words"].append(text)
            box["left"] = min(box["left"], left)
            box["top"] = min(box["top"], top)
            box["right"] = max(box["right"], left + width)
            box["bottom"] = max(box["bottom"], top + height)

    elements: list[LayoutElement] = []
    for key in order:
        box = lines[key]
        # Map back from the (possibly upscaled) OCR image to the original.
        elements.append(
            LayoutElement(
                text=" ".join(box["words"]),
                left=(box["left"] / scale) / orig_w,
                top=(box["top"] / scale) / orig_h,
                width=((box["right"] - box["left"]) / scale) / orig_w,
                height=((box["bottom"] - box["top"]) / scale) / orig_h,
            )
        )
    return elements


# EasyOCR's Reader loads its detection/recognition models from disk (and
# downloads them on first use) — expensive enough that it must be built once
# and reused, never per-request. Built lazily (only if OCR_ENGINE=easyocr is
# actually selected) since the package itself is an optional, heavy
# dependency (see requirements.txt) most deployments won't install.
_easyocr_reader: Any = None
_easyocr_lock = threading.Lock()


def _get_easyocr_reader() -> Any:
    global _easyocr_reader
    if _easyocr_reader is not None:
        return _easyocr_reader
    with _easyocr_lock:
        if _easyocr_reader is None:
            try:
                import easyocr
            except ImportError as exc:
                raise OcrUnavailableError(
                    "EasyOCR isn't installed on the server — run `pip install easyocr` (a heavy, "
                    "PyTorch-based package), or set OCR_ENGINE=tesseract in the backend's .env instead."
                ) from exc
            langs = [lang.strip() for lang in settings.ocr_languages.split(",") if lang.strip()] or ["en"]
            _easyocr_reader = easyocr.Reader(langs, gpu=False)
    return _easyocr_reader


def _easyocr_extract(img: np.ndarray) -> list[LayoutElement]:
    orig_h, orig_w = img.shape[:2]
    reader = _get_easyocr_reader()
    # EasyOCR does its own internal preprocessing; it wants the (unscaled,
    # color) image directly rather than the Tesseract-oriented binarized one.
    results = reader.readtext(img)  # [(box_points, text, confidence), ...]

    elements = [
        LayoutElement(
            text=text.strip(),
            left=min(p[0] for p in box_points) / orig_w,
            top=min(p[1] for p in box_points) / orig_h,
            width=(max(p[0] for p in box_points) - min(p[0] for p in box_points)) / orig_w,
            height=(max(p[1] for p in box_points) - min(p[1] for p in box_points)) / orig_h,
        )
        for box_points, text, confidence in results
        if text.strip() and confidence >= 0.4
    ]
    # Unlike Tesseract's image_to_data, EasyOCR's results aren't in reading
    # order — approximate top-to-bottom/left-to-right by rounding each
    # element's top to a coarse "row" band before sorting, so words on the
    # same visual line don't get separated by sub-pixel top differences.
    return sorted(elements, key=lambda e: (round(e.top * 40), e.left))


def screenshot_to_layout_json(image_bytes: bytes) -> dict[str, Any]:
    """The module's one real entry point — see the module docstring for the
    full pipeline. Runs synchronously (OpenCV/Tesseract/EasyOCR are all
    blocking C-extension calls); callers on the request path must offload via
    asyncio.to_thread, same convention app/core/ai_client.py's
    _sync_chat_completion already uses for its own blocking HTTP call."""
    img = _decode_image(image_bytes)
    orig_h, orig_w = img.shape[:2]
    engine = (settings.ocr_engine or "tesseract").strip().lower()
    if engine == "easyocr":
        elements = _easyocr_extract(img)
    elif engine == "tesseract":
        elements = _tesseract_extract(img)
    else:
        raise OcrUnavailableError(f"Unknown OCR_ENGINE '{engine}' — set it to \"tesseract\" or \"easyocr\".")
    return {
        "engine": engine,
        "image_width": orig_w,
        "image_height": orig_h,
        "elements": [e.to_json() for e in elements],
    }