from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
from transformers import AutoTokenizer
import pdfplumber
import logging
import re

logger = logging.getLogger(__name__)

class DoclingParser:
    """Uses IBM Docling's advanced structural parsing and HybridChunker."""
    
    def __init__(self, tokenizer_model: str, max_tokens: int, ocr_mode: str):
        self.ocr_mode = (ocr_mode or "").lower()
        if self.ocr_mode not in {"never", "auto", "always"}:
            raise ValueError("rag.ocr_mode must be one of: never, auto, always")

        # Non-OCR parser for fast PDFs with embedded text.
        no_ocr_options = PdfPipelineOptions(
            do_ocr=False, 
            generate_picture_images=False,
            # Reduce memory bloat
            layout_batch_size=2,
            table_batch_size=2
        )
        self.no_ocr_converter = DocumentConverter(
            format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=no_ocr_options)}
        )
        # OCR parser for scanned/badly encoded PDFs.
        ocr_options = PdfPipelineOptions(
            do_ocr=True, 
            generate_picture_images=False,
            # Aggressively reduce batch sizes to prevent OOM
            ocr_batch_size=1,
            layout_batch_size=1,
            table_batch_size=1,
            queue_max_size=2
        )
        self.ocr_converter = DocumentConverter(
            format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=ocr_options)}
        )
        
        # 2. Setup IBM's HybridChunker (Chunks structurally by tables/headings, not blindly by character count)
        try:
            tokenizer = AutoTokenizer.from_pretrained(tokenizer_model)
            self.chunker = HybridChunker(tokenizer=tokenizer, max_tokens=max_tokens)
        except Exception as e:
            logger.warning(
                f"Could not load tokenizer {tokenizer_model}. "
                f"Falling back to raw document text chunking. Error: {e}"
            )
            self.chunker = None
            
        logger.info("Initialized Advanced Docling Parser & HybridChunker.")

    def parse_and_chunk(self, file_path: str) -> list[str]:
        """Parses a PDF and returns a list of structurally intact text/table chunks."""
        logger.info(f"Parsing structure for: {file_path}")
        
        # Test standard fast extraction first to preserve natural layout 
        # (e.g., forms where Docling's layout AI incorrectly splits columns)
        fast_text = self._extract_with_pdfplumber(file_path)
        fast_chunks = [fast_text] if fast_text else []

        if self.ocr_mode == "never":
            return fast_chunks if fast_chunks and not self._is_likely_garbled(fast_chunks) else self._parse_with_converter(self.no_ocr_converter, file_path)

        if self.ocr_mode == "always":
            return self._parse_with_converter(self.ocr_converter, file_path)

        # Auto mode: Use pdfplumber if it yields clean layout-preserved text.
        if fast_chunks and not self._is_likely_garbled(fast_chunks):
            logger.info("Auto mode: pdfplumber successfully extracted clean text format.")
            return fast_chunks

        # Fallback to Docling non-OCR, then OCR if still garbled.
        logger.warning(f"Auto mode: fast extraction garbled or empty. Trying Docling for: {file_path}")
        chunks = self._parse_with_converter(self.no_ocr_converter, file_path)
        if chunks and not self._is_likely_garbled(chunks):
            return chunks
            
        logger.warning(f"Auto OCR fallback triggered (running layout and OCR models) for: {file_path}")
        ocr_chunks = self._parse_with_converter(self.ocr_converter, file_path)
        return ocr_chunks if ocr_chunks else chunks

    def _parse_with_converter(self, converter, file_path: str) -> list[str]:
        try:
            result = converter.convert(file_path)
            return self._doc_to_chunks(result)
        except Exception as e:
            logger.warning(f"Docling conversion failed for {file_path}; using pdfplumber fallback. Error: {e}")
            text = self._extract_with_pdfplumber(file_path)
            return [text] if text else []

    def _doc_to_chunks(self, result) -> list[str]:
        # To guarantee no text is dropped or limited by the chunking process,
        # we extract the complete document as Markdown from A-Z.
        # This will then be passed to the text splitters to split exhaustively.
        if hasattr(result.document, "export_to_markdown"):
            return [result.document.export_to_markdown()]
        elif hasattr(result.document, "export_to_text"):
            return [result.document.export_to_text()]
        return [str(result.document)]

    def _is_likely_garbled(self, chunks: list[str], threshold: float = 0.08) -> bool:
        text = "\n".join(chunks or [])
        if not text:
            return True
        # Very short text can be sparse/noisy naturally; avoid aggressive OCR fallback.
        if len(text) < 400:
            return False
        total = len(text)
        control_like = 0
        for ch in text:
            code = ord(ch)
            if code < 32 and ch not in ("\n", "\r", "\t"):
                control_like += 1
            elif 127 <= code <= 159:
                control_like += 1
        control_ratio = control_like / total

        # If there are enough readable words, keep non-OCR output even with symbols.
        word_count = len(re.findall(r"[A-Za-z]{2,}", text))
        has_readable_text = word_count >= 40
        if has_readable_text:
            return False

        return control_ratio >= threshold

    def _extract_with_pdfplumber(self, file_path: str) -> str:
        parts = []
        with pdfplumber.open(file_path) as pdf:
            for page in pdf.pages:
                page_text = page.extract_text() or ""
                if page_text.strip():
                    parts.append(page_text)
        return "\n\n".join(parts).strip()