from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import fitz
import pytesseract
from pdf2image import convert_from_path
from PIL import Image


@dataclass
class LayoutOCRResult:
    engine_name: str
    engine_version: str
    pages: list[dict[str, Any]]

    def to_analysis_json(self) -> dict[str, Any]:
        return {
            "schema_version": 1,
            "analysis_type": "canonical",
            "engine": {
                "name": self.engine_name,
                "version": self.engine_version,
            },
            "pages": self.pages,
        }


def _group_words_into_lines(words: list[dict[str, Any]], y_tol: float = 12.0) -> list[dict[str, Any]]:
    if not words:
        return []

    words = sorted(words, key=lambda w: (w["bbox"][1], w["bbox"][0]))
    groups: list[list[dict[str, Any]]] = []

    for word in words:
        placed = False
        wy = word["bbox"][1]
        for group in groups:
            gy = sum(item["bbox"][1] for item in group) / len(group)
            if abs(wy - gy) <= y_tol:
                group.append(word)
                placed = True
                break
        if not placed:
            groups.append([word])

    lines: list[dict[str, Any]] = []
    for group in groups:
        group = sorted(group, key=lambda w: w["bbox"][0])
        text = " ".join((w.get("text") or "").strip() for w in group).strip()
        if not text:
            continue
        left = min(w["bbox"][0] for w in group)
        top = min(w["bbox"][1] for w in group)
        right = max(w["bbox"][2] for w in group)
        bottom = max(w["bbox"][3] for w in group)
        avg_height = max(1.0, sum((w["bbox"][3] - w["bbox"][1]) for w in group) / len(group))
        lines.append(
            {
                "text": text,
                "bbox": [left, top, right, bottom],
                "confidence": None,
                "font_family_guess": "Helvetica",
                "font_size_guess": max(6.0, avg_height * 0.75),
                "text_color_guess": "#000000",
                "words": group,
            }
        )
    return lines


def run_layout_ocr(pdf_path: str | Path, dpi: int = 300) -> LayoutOCRResult:
    pdf_path = Path(pdf_path)
    if not pdf_path.exists():
        raise FileNotFoundError(f"PDF not found: {pdf_path}")

    doc = fitz.open(pdf_path)
    pil_pages = convert_from_path(str(pdf_path), dpi=dpi)

    pages: list[dict[str, Any]] = []

    for idx, (pdf_page, pil_img) in enumerate(zip(doc, pil_pages), start=1):
        page_w = float(pdf_page.rect.width)
        page_h = float(pdf_page.rect.height)

        if not isinstance(pil_img, Image.Image):
            raise ValueError(f"Rendered page {idx} is not a PIL image")

        img_w, img_h = pil_img.size
        scale_x = page_w / float(img_w)
        scale_y = page_h / float(img_h)

        data = pytesseract.image_to_data(
            pil_img,
            output_type=pytesseract.Output.DICT,
            config="--oem 3 --psm 6",
        )

        words: list[dict[str, Any]] = []
        n = len(data.get("text", []))
        for i in range(n):
            text = (data["text"][i] or "").strip()
            if not text:
                continue

            try:
                conf = float(data["conf"][i])
            except Exception:
                conf = None

            left_px = float(data["left"][i])
            top_px = float(data["top"][i])
            width_px = float(data["width"][i])
            height_px = float(data["height"][i])

            if width_px <= 0 or height_px <= 0:
                continue

            left = left_px * scale_x
            top = top_px * scale_y
            right = (left_px + width_px) * scale_x
            bottom = (top_px + height_px) * scale_y

            words.append(
                {
                    "text": text,
                    "bbox": [left, top, right, bottom],
                    "confidence": conf,
                }
            )

        lines = _group_words_into_lines(words)

        pages.append(
            {
                "page": idx,
                "page_width": page_w,
                "page_height": page_h,
                "image_width": page_w,
                "image_height": page_h,
                "lines": lines,
                "words": words,
            }
        )

    return LayoutOCRResult(
        engine_name="tesseract_layout",
        engine_version=str(pytesseract.get_tesseract_version()),
        pages=pages,
    )