import os
import re
import tempfile
import threading
import logging
import fitz
from rapidocr_onnxruntime import RapidOCR

logger = logging.getLogger(__name__)

OCR_PAGE_TIMEOUT = 30  # detik maksimum per halaman OCR

class HybridParser:
    def __init__(self):
        self.ocr_engine = RapidOCR()
        self._ocr_lock = threading.Lock()

    def _ocr_single_page(self, img_bytes: bytes) -> str:
        """
        OCR satu halaman dengan timeout protection.
        Menggunakan threading.Thread.join(timeout) agar kompatibel lintas OS.
        """
        result_container = []
        error_container = []

        def _run():
            try:
                with self._ocr_lock:
                    result, _ = self.ocr_engine(img_bytes)
                result_container.append(result)
            except Exception as e:
                error_container.append(e)

        t = threading.Thread(target=_run, daemon=True)
        t.start()
        t.join(timeout=OCR_PAGE_TIMEOUT)

        if t.is_alive():
            logger.warning(f"[HybridParser] OCR timeout setelah {OCR_PAGE_TIMEOUT}s, halaman dilewati.")
            return ""

        if error_container:
            raise error_container[0]

        result = result_container[0] if result_container else None
        if result:
            return "\n".join([item[1] for item in result])
        return ""

    def is_gibberish(self, pdf_path: str, threshold: float = 0.6) -> bool:
        """
        Menghitung rasio karakter alfabetis vs non-alfabetis.
        Mengembalikan True jika rasio valid < threshold (default 60%).
        """
        try:
            doc = fitz.open(pdf_path)
            total_chars = 0
            valid_chars = 0

            # Cek 3 halaman pertama untuk kecepatan
            for i in range(min(3, len(doc))):
                text = doc[i].get_text()
                total_chars += len(text)
                for char in text:
                    if char.isalpha() or char.isspace():
                        valid_chars += 1

            if total_chars == 0:
                return True

            ratio = valid_chars / total_chars
            return ratio < threshold
        except Exception as e:
            logger.error(f"[HybridParser] Error pada is_gibberish: {e}", exc_info=True)
            return True

    def _full_ocr(self, pdf_path: str) -> str:
        """Fallback ke OCR penuh jika PDF terdeteksi sebagai scan/gibberish."""
        try:
            doc = fitz.open(pdf_path)
            full_text = []
            for page_num in range(len(doc)):
                page = doc[page_num]
                pix = page.get_pixmap(dpi=150)
                img_bytes = pix.tobytes("png")
                page_text = self._ocr_single_page(img_bytes)
                if page_text:
                    full_text.append(page_text)

            if not full_text:
                raise RuntimeError(f"_full_ocr menghasilkan teks kosong untuk {pdf_path}")

            return "\n\n".join(full_text)
        except Exception as e:
            logger.error(f"[HybridParser] Error pada _full_ocr: {e}", exc_info=True)
            raise

    def _extract_with_pymupdf4llm(self, pdf_path: str) -> str:
        """
        Ekstraksi teks digital ke Markdown menggunakan pymupdf4llm.
        Menggunakan tempfile untuk menyimpan gambar, melakukan OCR pada gambar,
        dan menginjeksikan teks hasil OCR kembali ke dalam Markdown.

        FIX: Menambahkan `use_ocr=False` untuk menghindari TypeError: 'NoneType' object is not iterable
        dari engine OCR bawaan pymupdf4llm.
        """
        import pymupdf4llm

        with tempfile.TemporaryDirectory() as temp_dir:
            try:
                md_text = pymupdf4llm.to_markdown(
                    pdf_path,
                    write_images=True,
                    image_path=temp_dir,
                    image_format="png",
                    header=False,
                    footer=False,
                    margins=(0, 72, 0, 72),
                    use_ocr=False  # FIX UTAMA
                )

                if md_text is None:
                    logger.warning(f"[HybridParser] pymupdf4llm returned None, attempting full OCR fallback...")
                    return self._full_ocr(pdf_path)

                md_text = (md_text or "").strip()
                if not md_text:
                    logger.warning(f"[HybridParser] Extracted markdown is empty, attempting full OCR fallback...")
                    return self._full_ocr(pdf_path)

                # Proses gambar dengan RapidOCR untuk injeksi
                for img_file in os.listdir(temp_dir):
                    if img_file.endswith(".png"):
                        img_path = os.path.join(temp_dir, img_file)

                        try:
                            with self._ocr_lock:
                                result, _ = self.ocr_engine(img_path)
                            ocr_text = ""
                            if result:
                                ocr_text = "\n".join([item[1] for item in result])

                            if ocr_text.strip():
                                injection = f"\n\n> **[Teks Gambar Hasil OCR - {img_file}]**:\n{ocr_text}\n"

                                md_img_pattern = re.compile(r'!\[.*?\]\([^)]*' + re.escape(img_file) + r'\)')
                                if md_img_pattern.search(md_text):
                                    md_text = md_img_pattern.sub(injection, md_text)
                                else:
                                    md_text += injection
                        except Exception as ocr_err:
                            logger.warning(f"[HybridParser] OCR failed for {img_file}: {ocr_err}")
                            continue

                return md_text
            except Exception as e:
                logger.error(
                    f"[HybridParser] Error pada _extract_with_pymupdf4llm: {type(e).__name__}: {e}",
                    exc_info=True,
                )
                logger.info(f"[HybridParser] Attempting full OCR fallback for {pdf_path}...")
                return self._full_ocr(pdf_path)

    def parse_pdf(self, pdf_path: str) -> str:
        """
        Fungsi utama Smart Dual-Pass.
        """
        if self.is_gibberish(pdf_path):
            logger.info(f"[HybridParser] PDF terdeteksi sebagai scan/gibberish. Memulai Full OCR...")
            return self._full_ocr(pdf_path)
        else:
            logger.info(f"[HybridParser] PDF terdeteksi sebagai digital. Memulai ekstraksi Markdown (pymupdf4llm)...")
            return self._extract_with_pymupdf4llm(pdf_path)
