"""
PP-StructureV3 Table Detection Wrapper

Provides automatic table detection in invoice images using PaddleOCR's
PP-StructureV3 pipeline. Supports both wired (bordered) and wireless
(borderless) tables commonly found in Swedish invoices.
"""

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Protocol
import logging

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class TableDetectorConfig:
    """Configuration for TableDetector."""

    device: str = "gpu:0"
    use_doc_orientation_classify: bool = False
    use_doc_unwarping: bool = False
    use_textline_orientation: bool = False
    # Use SLANeXt models for better table recognition accuracy
    # SLANeXt_wireless has ~6% higher accuracy than SLANet for borderless tables
    wired_table_model: str = "SLANeXt_wired"
    wireless_table_model: str = "SLANeXt_wireless"
    layout_model: str = "PP-DocLayout_plus-L"
    min_confidence: float = 0.5


@dataclass
class TableDetectionResult:
    """Result of table detection."""

    bbox: tuple[float, float, float, float]  # x1, y1, x2, y2 in pixels
    html: str  # Table structure as HTML
    confidence: float
    table_type: str  # 'wired' or 'wireless'
    cells: list[dict[str, Any]] = field(default_factory=list)  # Cell-level data


class PPStructureProtocol(Protocol):
    """Protocol for PP-StructureV3 pipeline interface."""

    def predict(self, image: str | np.ndarray, **kwargs: Any) -> Any:
        """Run prediction on image."""
        ...


class TableDetector:
    """
    Table detector using PP-StructureV3.

    Detects tables in invoice images and returns their bounding boxes,
    HTML structure, and cell-level data.
    """

    def __init__(
        self,
        config: TableDetectorConfig | None = None,
        pipeline: PPStructureProtocol | None = None,
    ):
        """
        Initialize table detector.

        Args:
            config: Configuration options. Uses defaults if None.
            pipeline: Optional pre-initialized PP-StructureV3 pipeline.
                     If None, will be lazily initialized on first use.
        """
        self.config = config or TableDetectorConfig()
        self._pipeline = pipeline
        self._initialized = pipeline is not None

    def _ensure_initialized(self) -> None:
        """Lazily initialize PP-Structure pipeline."""
        if self._initialized:
            return

        # Try PPStructureV3 first (paddleocr >= 3.0.0), fall back to PPStructure (2.x)
        try:
            from paddleocr import PPStructureV3

            self._pipeline = PPStructureV3(
                layout_detection_model_name=self.config.layout_model,
                wired_table_structure_recognition_model_name=self.config.wired_table_model,
                wireless_table_structure_recognition_model_name=self.config.wireless_table_model,
                use_doc_orientation_classify=self.config.use_doc_orientation_classify,
                use_doc_unwarping=self.config.use_doc_unwarping,
                use_textline_orientation=self.config.use_textline_orientation,
                device=self.config.device,
            )
            self._initialized = True
            logger.info("PP-StructureV3 pipeline initialized successfully")
        except ImportError:
            # Fall back to PPStructure (paddleocr 2.x)
            try:
                from paddleocr import PPStructure

                # Map device config to use_gpu for PPStructure 2.x
                use_gpu = "gpu" in self.config.device.lower()
                self._pipeline = PPStructure(
                    table=True,
                    ocr=True,
                    use_gpu=use_gpu,
                    show_log=False,
                )
                self._initialized = True
                logger.info("PPStructure (2.x) pipeline initialized successfully")
            except ImportError as e:
                raise ImportError(
                    "PPStructure requires paddleocr. "
                    "Install with: pip install paddleocr"
                ) from e

    def detect(
        self,
        image: np.ndarray | str | Path,
    ) -> list[TableDetectionResult]:
        """
        Detect tables in an image.

        Args:
            image: Input image as numpy array, file path, or Path object.

        Returns:
            List of TableDetectionResult for each detected table.
        """
        self._ensure_initialized()

        if self._pipeline is None:
            raise RuntimeError("Pipeline not initialized")

        # Convert Path to string
        if isinstance(image, Path):
            image = str(image)

        # Run detection
        results = self._pipeline.predict(image)

        return self._parse_results(results)

    def _parse_results(self, results: Any) -> list[TableDetectionResult]:
        """Parse PP-StructureV3 output into TableDetectionResult list.

        Supports both:
        - PaddleX 3.x API: dict-like LayoutParsingResultV2 with table_res_list
        - Legacy API: objects with layout_elements attribute
        """
        tables: list[TableDetectionResult] = []

        if results is None:
            logger.warning("PP-StructureV3 returned None results")
            return tables

        # Log raw result type for debugging
        logger.debug(f"PP-StructureV3 raw results type: {type(results).__name__}")

        # Handle case where results is a single dict-like object (PaddleX 3.x)
        # rather than a list of results
        if hasattr(results, "get") and not isinstance(results, list):
            # Single result object - wrap in list for uniform processing
            logger.debug("Results is dict-like, wrapping in list")
            results = [results]
        elif hasattr(results, "__iter__") and not isinstance(results, (list, tuple)):
            # Iterator or generator - convert to list
            try:
                results = list(results)
                logger.debug(f"Converted iterator to list with {len(results)} items")
            except Exception as e:
                logger.warning(f"Failed to convert results to list: {e}")
                return tables

        logger.debug(f"Processing {len(results)} result(s)")

        for i, result in enumerate(results):
            try:
                result_type = type(result).__name__
                has_get = hasattr(result, "get")
                has_layout = hasattr(result, "layout_elements")
                logger.debug(f"Result[{i}]: type={result_type}, has_get={has_get}, has_layout_elements={has_layout}")

                # Try PaddleX 3.x API first (dict-like with table_res_list)
                if has_get:
                    parsed = self._parse_paddlex_result(result)
                    logger.debug(f"Result[{i}]: parsed {len(parsed)} tables via PaddleX path")
                    tables.extend(parsed)
                    continue

                # Fall back to legacy API (layout_elements)
                if has_layout:
                    legacy_count = 0
                    for element in result.layout_elements:
                        if not self._is_table_element(element):
                            continue
                        table_result = self._extract_table_data(element)
                        if table_result and table_result.confidence >= self.config.min_confidence:
                            tables.append(table_result)
                            legacy_count += 1
                    logger.debug(f"Result[{i}]: parsed {legacy_count} tables via legacy path")
                else:
                    logger.warning(f"Result[{i}]: no recognized API (not dict-like and no layout_elements)")
            except Exception as e:
                logger.warning(f"Failed to parse result: {type(result).__name__}, error: {e}")
                continue

        logger.debug(f"Total tables detected: {len(tables)}")
        return tables

    def _parse_paddlex_result(self, result: Any) -> list[TableDetectionResult]:
        """Parse PaddleX 3.x LayoutParsingResultV2."""
        tables: list[TableDetectionResult] = []

        try:
            # Log result structure for debugging
            result_type = type(result).__name__
            result_keys = []
            if hasattr(result, "keys"):
                result_keys = list(result.keys())
            elif hasattr(result, "__dict__"):
                result_keys = list(result.__dict__.keys())
            logger.debug(f"Parsing PaddleX result: type={result_type}, keys={result_keys}")

            # Get table results from PaddleX 3.x API
            # Handle both dict.get() and attribute access
            if hasattr(result, "get"):
                table_res_list = result.get("table_res_list")
                parsing_res_list = result.get("parsing_res_list", [])
            else:
                table_res_list = getattr(result, "table_res_list", None)
                parsing_res_list = getattr(result, "parsing_res_list", [])

            logger.debug(f"table_res_list: {type(table_res_list).__name__}, count={len(table_res_list) if table_res_list else 0}")
            logger.debug(f"parsing_res_list: {type(parsing_res_list).__name__}, count={len(parsing_res_list) if parsing_res_list else 0}")

            if not table_res_list:
                # Log available keys/attributes for debugging
                logger.warning(f"No table_res_list found in result: {result_type}, available: {result_keys}")
                return tables

            # Get parsing_res_list to find table bounding boxes
            table_bboxes = {}
            for elem in parsing_res_list or []:
                try:
                    if isinstance(elem, dict):
                        label = elem.get("label", "")
                        bbox = elem.get("bbox", [])
                    else:
                        label = getattr(elem, "label", "")
                        bbox = getattr(elem, "bbox", [])
                    # Check bbox has items (handles numpy arrays safely)
                    has_bbox = False
                    try:
                        has_bbox = len(bbox) >= 4 if hasattr(bbox, "__len__") else False
                    except (TypeError, ValueError):
                        pass
                    if label == "table" and has_bbox:
                        # Map by index (parsing_res_list tables appear in order)
                        idx = len(table_bboxes)
                        table_bboxes[idx] = bbox
                except Exception as e:
                    logger.debug(f"Failed to parse parsing_res element: {e}")
                    continue

            for i, table_res in enumerate(table_res_list):
                try:
                    # Extract from PaddleX 3.x table result format
                    # Handle both dict and object access (SingleTableRecognitionResult)
                    if isinstance(table_res, dict):
                        cell_boxes = table_res.get("cell_box_list", [])
                        html = table_res.get("pred_html", "")
                        ocr_data = table_res.get("table_ocr_pred", {})
                    else:
                        cell_boxes = getattr(table_res, "cell_box_list", [])
                        html = getattr(table_res, "pred_html", "")
                        ocr_data = getattr(table_res, "table_ocr_pred", {})

                    # table_ocr_pred can be dict (PaddleOCR 3.x) or list (older versions)
                    # For dict format: {"rec_texts": [...], "rec_scores": [...], ...}
                    ocr_texts = []
                    if isinstance(ocr_data, dict):
                        ocr_texts = ocr_data.get("rec_texts", [])
                    elif isinstance(ocr_data, list):
                        ocr_texts = ocr_data

                    # Try to get bbox from parsing_res_list
                    bbox = table_bboxes.get(i, [0.0, 0.0, 0.0, 0.0])
                    # Handle numpy arrays - check length explicitly to avoid boolean ambiguity
                    try:
                        bbox_len = len(bbox) if hasattr(bbox, "__len__") else 0
                        if bbox_len < 4:
                            bbox = [0.0, 0.0, 0.0, 0.0]
                    except (TypeError, ValueError):
                        bbox = [0.0, 0.0, 0.0, 0.0]

                    # Build cells from cell_box_list and OCR text
                    cells = []
                    # Check cell_boxes length explicitly to avoid numpy array boolean issues
                    has_cell_boxes = False
                    try:
                        has_cell_boxes = len(cell_boxes) > 0 if hasattr(cell_boxes, "__len__") else bool(cell_boxes)
                    except (TypeError, ValueError):
                        pass
                    if has_cell_boxes:
                        # Check ocr_texts length safely for numpy arrays
                        ocr_texts_len = 0
                        try:
                            ocr_texts_len = len(ocr_texts) if hasattr(ocr_texts, "__len__") else 0
                        except (TypeError, ValueError):
                            pass
                        for j, cell_bbox in enumerate(cell_boxes):
                            cell_text = ocr_texts[j] if ocr_texts_len > j else ""
                            # Convert cell_bbox to list safely (may be numpy array)
                            cell_bbox_list = []
                            try:
                                cell_bbox_list = list(cell_bbox) if hasattr(cell_bbox, "__iter__") else []
                            except (TypeError, ValueError):
                                pass
                            cells.append({
                                "text": cell_text,
                                "bbox": cell_bbox_list,
                                "row": 0,  # Row/col info not directly available
                                "col": j,
                            })

                    # Default confidence for PaddleX 3.x results
                    confidence = 0.9

                    logger.debug(f"Table {i}: html_len={len(html)}, cells={len(cells)}")
                    tables.append(TableDetectionResult(
                        bbox=(float(bbox[0]), float(bbox[1]), float(bbox[2]), float(bbox[3])),
                        html=html,
                        confidence=confidence,
                        table_type="wired",  # PaddleX 3.x handles both types
                        cells=cells,
                    ))
                except Exception as e:
                    import traceback
                    logger.warning(f"Failed to parse table_res {i}: {e}\n{traceback.format_exc()}")
                    continue

        except Exception as e:
            logger.warning(f"Failed to parse PaddleX result: {type(e).__name__}: {e}")

        return tables

    def _is_table_element(self, element: Any) -> bool:
        """Check if element is a table."""
        if hasattr(element, "label"):
            return element.label.lower() in ("table", "wired_table", "wireless_table")
        if hasattr(element, "type"):
            return element.type.lower() in ("table", "wired_table", "wireless_table")
        return False

    def _extract_table_data(self, element: Any) -> TableDetectionResult | None:
        """Extract table data from PP-StructureV3 element."""
        try:
            # Get bounding box
            bbox = self._get_bbox(element)
            if bbox is None:
                return None

            # Get HTML content
            html = self._get_html(element)

            # Get confidence
            confidence = getattr(element, "score", 0.9)
            if isinstance(confidence, (list, tuple)):
                confidence = float(confidence[0]) if confidence else 0.9

            # Determine table type
            table_type = self._get_table_type(element)

            # Get cells if available
            cells = self._get_cells(element)

            return TableDetectionResult(
                bbox=bbox,
                html=html,
                confidence=float(confidence),
                table_type=table_type,
                cells=cells,
            )
        except Exception as e:
            logger.warning(f"Failed to extract table data: {e}")
            return None

    def _get_bbox(self, element: Any) -> tuple[float, float, float, float] | None:
        """Extract bounding box from element."""
        if hasattr(element, "bbox"):
            bbox = element.bbox
            if len(bbox) >= 4:
                return (float(bbox[0]), float(bbox[1]), float(bbox[2]), float(bbox[3]))
        if hasattr(element, "box"):
            box = element.box
            if len(box) >= 4:
                return (float(box[0]), float(box[1]), float(box[2]), float(box[3]))
        return None

    def _get_html(self, element: Any) -> str:
        """Extract HTML content from element."""
        if hasattr(element, "html"):
            return str(element.html)
        if hasattr(element, "table_html"):
            return str(element.table_html)
        if hasattr(element, "res") and isinstance(element.res, dict):
            return element.res.get("html", "")
        return ""

    def _get_table_type(self, element: Any) -> str:
        """Determine table type (wired or wireless)."""
        label = ""
        if hasattr(element, "label"):
            label = str(element.label).lower()
        elif hasattr(element, "type"):
            label = str(element.type).lower()

        if "wireless" in label or "borderless" in label:
            return "wireless"
        return "wired"

    def _get_cells(self, element: Any) -> list[dict[str, Any]]:
        """Extract cell-level data from element."""
        cells: list[dict[str, Any]] = []

        if hasattr(element, "cells"):
            for cell in element.cells:
                cell_data = {
                    "text": getattr(cell, "text", ""),
                    "row": getattr(cell, "row", 0),
                    "col": getattr(cell, "col", 0),
                    "row_span": getattr(cell, "row_span", 1),
                    "col_span": getattr(cell, "col_span", 1),
                }
                if hasattr(cell, "bbox"):
                    cell_data["bbox"] = cell.bbox
                cells.append(cell_data)

        return cells

    def detect_from_pdf(
        self,
        pdf_path: str | Path,
        page_number: int = 0,
        dpi: int = 300,
    ) -> list[TableDetectionResult]:
        """
        Detect tables from a PDF page.

        Args:
            pdf_path: Path to PDF file.
            page_number: Page number (0-indexed).
            dpi: Resolution for rendering.

        Returns:
            List of TableDetectionResult for the specified page.
        """
        from shared.pdf.renderer import render_pdf_to_images
        from PIL import Image
        import io

        pdf_path = Path(pdf_path)
        if not pdf_path.exists():
            raise FileNotFoundError(f"PDF not found: {pdf_path}")

        logger.debug(f"detect_from_pdf: {pdf_path}, page={page_number}, dpi={dpi}")

        # Render specific page
        for page_no, image_bytes in render_pdf_to_images(str(pdf_path), dpi=dpi):
            if page_no == page_number:
                image = Image.open(io.BytesIO(image_bytes))
                image_array = np.array(image)
                logger.debug(f"detect_from_pdf: rendered page {page_no}, image shape={image_array.shape}")
                return self.detect(image_array)

        raise ValueError(f"Page {page_number} not found in PDF")
