"""
Text Extraction Module
Handles PDF, DOCX, and Image file types with layout-aware extraction.
"""
import io
import logging
from pathlib import Path
from typing import Optional

import pdfplumber
import fitz  # PyMuPDF
from docx import Document as DocxDocument
from PIL import Image

from app.modules.ocr_processor import OCRProcessor

logger = logging.getLogger(__name__)


class TextExtractor:
    """
    Routes documents to the appropriate extraction pipeline based on file type.
    Supports PDF (text-based and scanned), DOCX, PNG, and JPG.
    """

    SUPPORTED_TYPES = {
        "application/pdf": "_extract_pdf",
        "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "_extract_docx",
        "image/png": "_extract_image",
        "image/jpeg": "_extract_image",
        "image/jpg": "_extract_image",
    }

    EXTENSION_MAP = {
        ".pdf": "application/pdf",
        ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
    }

    def __init__(self):
        self.ocr = OCRProcessor()

    def detect_mime_type(self, filename: str, content: bytes) -> str:
        """Detect MIME type from filename extension with fallback to content sniffing."""
        ext = Path(filename).suffix.lower()
        if ext in self.EXTENSION_MAP:
            return self.EXTENSION_MAP[ext]

        # Content-based sniffing
        if content[:4] == b"%PDF":
            return "application/pdf"
        if content[:2] in (b"PK",):
            return "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        if content[:8] in (b"\x89PNG\r\n\x1a\n",):
            return "image/png"
        if content[:3] == b"\xff\xd8\xff":
            return "image/jpeg"

        raise ValueError(f"Unsupported or undetectable file type for: {filename}")

    def extract(self, content: bytes, filename: str) -> str:
        """
        Main entry point. Detects file type and routes to correct extractor.
        Returns cleaned, normalized text string.
        """
        mime = self.detect_mime_type(filename, content)
        method_name = self.SUPPORTED_TYPES.get(mime)
        if not method_name:
            raise ValueError(f"Unsupported MIME type: {mime}")

        method = getattr(self, method_name)
        logger.info(f"Extracting text from {filename} (MIME: {mime})")
        return method(content, filename)

    # ─── PDF Extraction ────────────────────────────────────────────────────────

    def _extract_pdf(self, content: bytes, filename: str) -> str:
        """
        Extract text from PDF. Tries selectable text first; falls back to OCR.
        Uses pdfplumber for layout-aware extraction.
        """
        text = self._extract_pdf_selectable(content)

        if not text or len(text.strip()) < 50:
            logger.info(f"PDF '{filename}' appears scanned. Falling back to OCR.")
            text = self._extract_pdf_ocr(content)

        return text

    def _extract_pdf_selectable(self, content: bytes) -> str:
        """Extract selectable text from PDF using pdfplumber (layout-aware)."""
        pages_text = []
        try:
            with pdfplumber.open(io.BytesIO(content)) as pdf:
                for page in pdf.pages:
                    page_text = page.extract_text(x_tolerance=3, y_tolerance=3)
                    if page_text:
                        pages_text.append(page_text.strip())
        except Exception as e:
            logger.warning(f"pdfplumber extraction failed: {e}")
        return "\n\n".join(pages_text)

    def _extract_pdf_ocr(self, content: bytes) -> str:
        """
        Render each PDF page as an image and apply OCR.
        Uses PyMuPDF (fitz) for reliable PDF→image rendering.
        """
        pages_text = []
        try:
            doc = fitz.open(stream=content, filetype="pdf")
            for page_num in range(len(doc)):
                page = doc[page_num]
                # Render at 300 DPI for good OCR quality
                mat = fitz.Matrix(300 / 72, 300 / 72)
                pix = page.get_pixmap(matrix=mat, alpha=False)
                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                page_text = self.ocr.extract(img)
                if page_text.strip():
                    pages_text.append(page_text.strip())
            doc.close()
        except Exception as e:
            logger.error(f"PDF OCR failed: {e}")
        return "\n\n".join(pages_text)

    # ─── DOCX Extraction ───────────────────────────────────────────────────────

    def _extract_docx(self, content: bytes, filename: str) -> str:
        """
        Extract text from DOCX preserving paragraph integrity.
        Maintains logical structure: headings are separated by blank lines.
        """
        try:
            doc = DocxDocument(io.BytesIO(content))
            paragraphs = []
            for para in doc.paragraphs:
                text = para.text.strip()
                if text:
                    # Preserve heading structure
                    if para.style.name.startswith("Heading"):
                        paragraphs.append(f"\n{text}\n")
                    else:
                        paragraphs.append(text)

            # Also extract from tables
            for table in doc.tables:
                for row in table.rows:
                    row_text = " | ".join(
                        cell.text.strip() for cell in row.cells if cell.text.strip()
                    )
                    if row_text:
                        paragraphs.append(row_text)

            return "\n".join(paragraphs)
        except Exception as e:
            logger.error(f"DOCX extraction failed: {e}")
            return ""

    # ─── Image Extraction ──────────────────────────────────────────────────────

    def _extract_image(self, content: bytes, filename: str) -> str:
        """Extract text from images using OCR with preprocessing."""
        try:
            img = Image.open(io.BytesIO(content))
            return self.ocr.extract(img)
        except Exception as e:
            logger.error(f"Image extraction failed: {e}")
            return ""