"""Core OCR: PaddleOCR + PPStructureV3."""

from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Callable

# Desabilita check de fonte de modelo antes de qualquer import do paddle
os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "True"

from .config import OcrConfig
from .device import resolve_device
from .renderer import get_page_count, render_all_parallel

if TYPE_CHECKING:
    import numpy as np


@dataclass
class PageResult:
    """Resultado OCR de uma página."""

    page_num: int
    texts: list[str]
    scores: list[float]


@dataclass
class PdfResult:
    """Resultado OCR de um PDF completo."""

    pdf_path: Path
    pages: list[PageResult]
    num_pages: int

    @property
    def total_texts(self) -> int:
        return sum(len(p.texts) for p in self.pages)


def create_ocr(config: OcrConfig):
    """Cria instância PaddleOCR configurada."""
    from paddleocr import PaddleOCR

    device = resolve_device(config.device)
    suffix = "server" if config.model == "server" else "mobile"

    return PaddleOCR(
        text_detection_model_name=f"PP-OCRv5_{suffix}_det",
        text_recognition_model_name=f"PP-OCRv5_{suffix}_rec",
        use_doc_orientation_classify=False,
        use_doc_unwarping=False,
        use_textline_orientation=False,
        device=device,
    )


def create_structure(config: OcrConfig):
    """Cria instância PPStructureV3 configurada."""
    from paddleocr import PPStructureV3

    device = resolve_device(config.device)

    return PPStructureV3(
        use_doc_orientation_classify=False,
        use_doc_unwarping=False,
        device=device,
    )


def process_pdf_text(
    pdf_path: str | Path,
    config: OcrConfig,
    on_render_progress: Callable[[int, int], None] | None = None,
    on_ocr_progress: Callable[[int, int], None] | None = None,
) -> PdfResult:
    """Processa um PDF com PP-OCRv5 (modo text).

    Fase 1: render paralelo das páginas.
    Fase 2: OCR sequencial (mais rápido que batch conforme benchmarks).
    """
    pdf_path = Path(pdf_path)
    num_pages = get_page_count(pdf_path)

    # Fase 1: render paralelo
    images = render_all_parallel(
        pdf_path,
        num_pages,
        scale=config.scale,
        workers=config.workers,
        on_progress=on_render_progress,
    )

    # Fase 2: OCR sequencial
    ocr = create_ocr(config)
    pages: list[PageResult] = []

    for i, img in enumerate(images):
        result = ocr.predict(img)
        texts: list[str] = []
        scores: list[float] = []

        for res in result:
            try:
                for text, score in zip(res["rec_texts"], res["rec_scores"]):
                    if score >= config.min_score:
                        texts.append(text)
                        scores.append(float(score))
            except (KeyError, TypeError):
                pass

        pages.append(PageResult(page_num=i + 1, texts=texts, scores=scores))

        if on_ocr_progress:
            on_ocr_progress(i + 1, num_pages)

    return PdfResult(pdf_path=pdf_path, pages=pages, num_pages=num_pages)


def process_pdf_structure(
    pdf_path: str | Path,
    config: OcrConfig,
    on_progress: Callable[[int, int], None] | None = None,
) -> PdfResult:
    """Processa um PDF com PPStructureV3 (modo structure).

    Usa o pipeline nativo que faz render + OCR internamente,
    e salva markdown via save_to_markdown().
    """
    import tempfile

    pdf_path = Path(pdf_path)
    pipeline = create_structure(config)

    result = pipeline.predict(str(pdf_path))

    # Salva markdown temporário para extrair conteúdo
    tmp_dir = tempfile.mkdtemp()
    pages: list[PageResult] = []

    for i, res in enumerate(result):
        res.save_to_markdown(tmp_dir)
        # Extrair textos dos resultados estruturais
        texts: list[str] = []
        try:
            if hasattr(res, "get") or isinstance(res, dict):
                for text in res.get("rec_texts", []):
                    texts.append(text)
        except (KeyError, TypeError):
            pass

        pages.append(PageResult(page_num=i + 1, texts=texts, scores=[]))

        if on_progress:
            on_progress(i + 1, len(result))

    # Ler os markdowns gerados
    md_files = sorted(Path(tmp_dir).glob("*.md"))
    md_content = "\n\n".join(f.read_text() for f in md_files)

    pdf_result = PdfResult(pdf_path=pdf_path, pages=pages, num_pages=len(result))
    pdf_result._structure_markdown = md_content  # type: ignore[attr-defined]
    pdf_result._structure_tmp_dir = tmp_dir  # type: ignore[attr-defined]

    return pdf_result


def process_pdf(
    pdf_path: str | Path,
    config: OcrConfig,
    on_render_progress: Callable[[int, int], None] | None = None,
    on_ocr_progress: Callable[[int, int], None] | None = None,
) -> PdfResult:
    """Processa um PDF de acordo com o modo configurado."""
    if config.mode == "structure":
        return process_pdf_structure(pdf_path, config, on_progress=on_ocr_progress)
    return process_pdf_text(
        pdf_path, config, on_render_progress=on_render_progress, on_ocr_progress=on_ocr_progress
    )
