from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from ocr_pipeline.config import settings
from ocr_pipeline.utils.logging import get_logger

logger = get_logger(__name__)


@dataclass
class DetectedFigure:
    bbox: list[float]
    confidence: float
    label: str
    caption: str | None = None
    image_crop: np.ndarray | None = None


class FigureDetector:
    def __init__(self):
        self._predictor = None
        self._initialized = False

    def _init(self):
        if self._initialized:
            return
        if not settings.detectors.figures.enabled:
            self._initialized = True
            return
        try:
            import layoutparser as lp

            self._predictor = lp.Detectron2LayoutModel(
                config_path=settings.detectors.figures.model,
                label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
                extra_config=[
                    "MODEL.ROI_HEADS.SCORE_THRESH_TEST",
                    settings.detectors.figures.confidence_threshold,
                ],
            )
            self._initialized = True
            logger.info("figure_detector_initialized", model=settings.detectors.figures.model)
        except Exception as e:
            logger.warning("figure_detector_init_failed", error=str(e))
            self._initialized = True

    def detect(self, image: np.ndarray) -> list[DetectedFigure]:
        if not settings.detectors.figures.enabled:
            return []

        self._init()
        if self._predictor is None:
            return []

        try:
            layout = self._predictor.detect(image)
            figures = []

            for block in layout:
                if block.type == "Figure":
                    x1, y1, x2, y2 = (
                        block.block.x_1,
                        block.block.y_1,
                        block.block.x_2,
                        block.block.y_2,
                    )
                    crop = image[int(y1) : int(y2), int(x1) : int(x2)]
                    figures.append(
                        DetectedFigure(
                            bbox=[float(x1), float(y1), float(x2), float(y2)],
                            confidence=float(block.score),
                            label="Figure",
                            image_crop=crop if crop.size > 0 else None,
                        )
                    )

            logger.debug("figures_detected", count=len(figures))
            return figures
        except Exception as e:
            logger.warning("figure_detection_failed", error=str(e))
            return []


class CaptionExtractor:
    def __init__(self):
        self._predictor = None
        self._initialized = False

    def _init(self):
        if self._initialized:
            return
        if not settings.detectors.captions.enabled:
            self._initialized = True
            return
        try:
            import layoutparser as lp

            self._predictor = lp.Detectron2LayoutModel(
                config_path=settings.detectors.captions.model,
                label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
                extra_config=[
                    "MODEL.ROI_HEADS.SCORE_THRESH_TEST",
                    settings.detectors.captions.confidence_threshold,
                ],
            )
            self._initialized = True
        except Exception as e:
            logger.warning("caption_detector_init_failed", error=str(e))
            self._initialized = True

    def extract_captions(
        self, image: np.ndarray, figures: list[DetectedFigure]
    ) -> list[DetectedFigure]:
        if not settings.detectors.captions.enabled or not figures:
            return figures

        self._init()
        if self._predictor is None:
            return figures

        try:
            layout = self._predictor.detect(image)
            captions = [
                b
                for b in layout
                if b.type in ("Title", "Text")
                and b.score > settings.detectors.captions.confidence_threshold
            ]

            for fig in figures:
                fig_x1, fig_y1, fig_x2, fig_y2 = fig.bbox
                fig_center_y = (fig_y1 + fig_y2) / 2

                best_caption = None
                min_dist = float("inf")

                for cap in captions:
                    cap_y1, cap_y2 = cap.block.y_1, cap.block.y_2
                    cap_center_y = (cap_y1 + cap_y2) / 2
                    dist = abs(cap_center_y - fig_center_y)

                    if cap_y2 < fig_y1 or cap_y1 > fig_y2:
                        if dist < min_dist:
                            min_dist = dist
                            best_caption = cap

                if best_caption and min_dist < 200:
                    fig.caption = best_caption.block.text

            return figures
        except Exception as e:
            logger.warning("caption_extraction_failed", error=str(e))
            return figures


def detect_figures(image: np.ndarray) -> list[DetectedFigure]:
    detector = FigureDetector()
    figures = detector.detect(image)

    captioner = CaptionExtractor()
    figures = captioner.extract_captions(image, figures)

    return figures
