from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import cv2
import numpy as np

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

logger = get_logger(__name__)


@dataclass
class DetectedTable:
    bbox: list[float]
    confidence: float
    label: str
    structure: list[dict[str, Any]] | None = None
    markdown: str | None = None
    image_crop: np.ndarray | None = None


class TableDetector:
    def __init__(self):
        self._model = None
        self._processor = None
        self._initialized = False

    def _init(self):
        if self._initialized:
            return
        if not settings.detectors.tables.enabled:
            self._initialized = True
            return
        try:
            from transformers import AutoImageProcessor, AutoModelForObjectDetection

            self._processor = AutoImageProcessor.from_pretrained(settings.detectors.tables.model)
            self._model = AutoModelForObjectDetection.from_pretrained(
                settings.detectors.tables.model
            )
            self._model.eval()
            self._initialized = True
            logger.info("table_detector_initialized", model=settings.detectors.tables.model)
        except Exception as e:
            logger.warning("table_detector_init_failed", error=str(e))
            self._initialized = True

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

        self._init()
        if self._model is None or self._processor is None:
            return []

        try:
            import torch

            pil_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            pil_image = cv2.resize(pil_image, (800, 800))

            inputs = self._processor(images=pil_image, return_tensors="pt")
            with torch.no_grad():
                outputs = self._model(**inputs)

            target_sizes = torch.tensor([pil_image.shape[:2]])
            results = self._processor.post_process_object_detection(
                outputs,
                threshold=settings.detectors.tables.confidence_threshold,
                target_sizes=target_sizes,
            )[0]

            tables = []
            for score, _label, box in zip(
                results["scores"], results["labels"], results["boxes"], strict=True
            ):
                box = [float(x) for x in box]
                x1, y1, x2, y2 = map(int, box)
                crop = image[y1:y2, x1:x2]
                tables.append(
                    DetectedTable(
                        bbox=box,
                        confidence=float(score),
                        label="Table",
                        image_crop=crop if crop.size > 0 else None,
                    )
                )

            logger.debug("tables_detected", count=len(tables))
            return tables
        except Exception as e:
            logger.warning("table_detection_failed", error=str(e))
            return []


class TableStructureRecognizer:
    def __init__(self):
        self._model = None
        self._processor = None
        self._initialized = False

    def _init(self):
        if self._initialized:
            return
        try:
            from transformers import AutoImageProcessor, AutoModelForObjectDetection

            self._processor = AutoImageProcessor.from_pretrained(
                "microsoft/table-transformer-structure-recognition"
            )
            self._model = AutoModelForObjectDetection.from_pretrained(
                "microsoft/table-transformer-structure-recognition"
            )
            self._model.eval()
            self._initialized = True
            logger.info("table_structure_recognizer_initialized")
        except Exception as e:
            logger.warning("table_structure_recognizer_init_failed", error=str(e))
            self._initialized = True

    def recognize(self, table_crop: np.ndarray) -> tuple[list[dict[str, Any]] | None, str | None]:
        if self._model is None or self._processor is None:
            return None, None

        try:
            import torch

            pil_image = cv2.cvtColor(table_crop, cv2.COLOR_BGR2RGB)
            inputs = self._processor(images=pil_image, return_tensors="pt")
            with torch.no_grad():
                outputs = self._model(**inputs)

            target_sizes = torch.tensor([pil_image.shape[:2]])
            results = self._processor.post_process_object_detection(
                outputs, threshold=0.5, target_sizes=target_sizes
            )[0]

            structure = []
            for score, label, box in zip(
                results["scores"], results["labels"], results["boxes"], strict=True
            ):
                structure.append(
                    {
                        "label": self._model.config.id2label[int(label)],
                        "confidence": float(score),
                        "bbox": [float(x) for x in box],
                    }
                )

            markdown = self._structure_to_markdown(structure, table_crop.shape)
            return structure, markdown
        except Exception as e:
            logger.warning("table_structure_recognition_failed", error=str(e))
            return None, None

    def _structure_to_markdown(self, structure: list[dict[str, Any]], shape: tuple) -> str | None:
        if not structure:
            return None

        rows = {}
        for cell in structure:
            if cell["label"] in ("table row", "table cell"):
                y_center = (cell["bbox"][1] + cell["bbox"][3]) / 2
                x_center = (cell["bbox"][0] + cell["bbox"][2]) / 2
                row_idx = int(y_center / (shape[0] / 20))
                if row_idx not in rows:
                    rows[row_idx] = []
                rows[row_idx].append((x_center, cell["label"]))

        markdown_lines = []
        for row_idx in sorted(rows.keys()):
            cells = sorted(rows[row_idx], key=lambda x: x[0])
            row_cells = [c[1] for c in cells]
            markdown_lines.append("| " + " | ".join(row_cells) + " |")

        if markdown_lines:
            header = markdown_lines[0]
            separator = "| " + " | ".join(["---"] * len(header.split("|")[1:-1])) + " |"
            markdown_lines.insert(1, separator)
            return "\n".join(markdown_lines)

        return None


def detect_tables(image: np.ndarray) -> list[DetectedTable]:
    detector = TableDetector()
    tables = detector.detect(image)

    if not tables:
        return tables

    recognizer = TableStructureRecognizer()
    for table in tables:
        if table.image_crop is not None:
            structure, markdown = recognizer.recognize(table.image_crop)
            table.structure = structure
            table.markdown = markdown

    return tables
