"""PDFApps – TabOCR: OCR text recognition tool."""

import os

from PySide6.QtCore import Qt
from PySide6.QtWidgets import (
    QHBoxLayout, QLabel, QComboBox, QFileDialog, QMessageBox,
    QProgressDialog,
)
from pypdf import PdfReader, PdfWriter

from app.base import BasePage
from app.i18n import t
from app.utils import section, info_lbl
from app.worker import TaskRunner, run_task
from app.constants import TEXT_SEC, DESKTOP
from app.widgets import DropFileEdit


def _find_tesseract() -> str | None:
    """Returns the tesseract executable path or None if not found."""
    import shutil, sys
    found = shutil.which("tesseract")
    if found:
        return found
    if sys.platform == "win32":
        candidates = [
            r"C:\Program Files\Tesseract-OCR\tesseract.exe",
            r"C:\Program Files (x86)\Tesseract-OCR\tesseract.exe",
        ]
    elif sys.platform == "darwin":
        candidates = [
            "/opt/homebrew/bin/tesseract",
            "/usr/local/bin/tesseract",
        ]
    else:
        candidates = [
            "/usr/bin/tesseract",
            "/usr/local/bin/tesseract",
            "/snap/bin/tesseract",
        ]
    for p in candidates:
        if os.path.isfile(p):
            return p
    return None


def _find_tessdata(tess_exe: str | None) -> str | None:
    """Locate the tessdata directory across platforms.

    Windows / macOS Homebrew put tessdata next to the binary. Debian /
    Ubuntu put it under /usr/share/tesseract-ocr/<version>/tessdata,
    and the binary itself often has a default path baked in that
    points at the wrong version (Ubuntu 24.04 ships v5 but the
    default still points at .../4.00/tessdata, see issue #27).
    Without TESSDATA_PREFIX set explicitly, OCR fails with
    'Error opening data file .../4.00/tessdata/eng.traineddata'."""
    import sys
    if tess_exe:
        adjacent = os.path.join(os.path.dirname(tess_exe), "tessdata")
        if os.path.isdir(adjacent):
            return adjacent
    if sys.platform.startswith(("linux", "darwin")):
        import glob
        # Sort descending so the latest version wins (5 over 4.00 on
        # Ubuntu 24.04 where both folders may coexist).
        for p in sorted(glob.glob("/usr/share/tesseract-ocr/*/tessdata"),
                        reverse=True):
            if os.path.isdir(p):
                return p
        # Older Debian, manual installs, snap fallbacks.
        for p in ("/usr/share/tessdata",
                  "/usr/local/share/tessdata",
                  "/snap/tesseract/current/usr/share/tesseract-ocr/tessdata"):
            if os.path.isdir(p):
                return p
    return None


class TabOCR(BasePage):
    _LANG_KEYS = [
        ("tool.ocr.lang.pt",    "por"),
        ("tool.ocr.lang.en",    "eng"),
        ("tool.ocr.lang.pt_en", "por+eng"),
        ("tool.ocr.lang.es",    "spa"),
        ("tool.ocr.lang.fr",    "fra"),
        ("tool.ocr.lang.de",    "deu"),
    ]

    @property
    def _LANGS(self):
        return [(t(k), c) for k, c in self._LANG_KEYS]

    def __init__(self, status_fn):
        super().__init__("fa5s.search", t("tool.ocr.name"),
                         t("tool.ocr.desc"),
                         t("tool.ocr.btn"), status_fn)
        f = self._form

        sec_src = section(t("tool.ocr.source"))
        f.addWidget(sec_src)
        self.drop_in = DropFileEdit()
        try: self.drop_in.btn.clicked.disconnect()
        except RuntimeError: pass
        self.drop_in.btn.clicked.connect(self._pick_input)
        self.drop_in.path_changed.connect(self._load_input)
        self.lbl_info = info_lbl()
        f.addWidget(self.drop_in); f.addWidget(self.lbl_info)

        f.addWidget(section(t("tool.ocr.options")))
        row_lang = QHBoxLayout()
        lbl_lang = QLabel(t("tool.ocr.lang_label"))
        lbl_lang.setStyleSheet(f"color:{TEXT_SEC};")
        self.cmb_lang = QComboBox()
        self._lang_codes = [c for _, c in self._LANGS]
        for name, _ in self._LANGS:
            self.cmb_lang.addItem(name)
        row_lang.addWidget(lbl_lang); row_lang.addWidget(self.cmb_lang); row_lang.addStretch()
        f.addLayout(row_lang)

        row_fmt = QHBoxLayout()
        lbl_fmt = QLabel(t("tool.ocr.format_label"))
        lbl_fmt.setStyleSheet(f"color:{TEXT_SEC};")
        self.cmb_fmt = QComboBox()
        self.cmb_fmt.addItems([t("tool.ocr.format.pdf"), t("tool.ocr.format.txt")])
        self.cmb_fmt.currentIndexChanged.connect(self._on_fmt_change)
        row_fmt.addWidget(lbl_fmt); row_fmt.addWidget(self.cmb_fmt); row_fmt.addStretch()
        f.addLayout(row_fmt)

        sec_out = section(t("tool.ocr.output"))
        f.addWidget(sec_out)
        self.drop_out = DropFileEdit("ocr_output.pdf", save=True, default_name="ocr_output.pdf")
        f.addWidget(self.drop_out)
        f.addStretch()
        self._compact_hidden = [sec_src, self.drop_in, self.lbl_info]
        sec_out.setVisible(False)
        self.drop_out.setVisible(False)

    def _on_fmt_change(self, idx):
        p = self.drop_out.path()
        if p:
            base = os.path.splitext(p)[0]
            self.drop_out.set_path(base + (".pdf" if idx == 0 else ".txt"))

    def _pick_input(self):
        p, _ = QFileDialog.getOpenFileName(self, t("btn.open_pdf"), DESKTOP, t("file_filter.pdf"))
        if p: self._load_input(p)

    def _load_input(self, p: str):
        self.drop_in.blockSignals(True)
        self.drop_in.set_path(p)
        self.drop_in.blockSignals(False)
        if not self._maybe_prompt_password(p):
            self.drop_in.blockSignals(True); self.drop_in.set_path("")
            self.drop_in.blockSignals(False); return
        if not self.drop_out.path():
            base = os.path.splitext(p)[0]
            ext = ".pdf" if self.cmb_fmt.currentIndex() == 0 else ".txt"
            self.drop_out.set_path(base + "_ocr" + ext)
        try:
            doc = self._open_fitz(p)
            self.lbl_info.setText(t("edit.status.pages", n=doc.page_count))
            doc.close()
        except Exception as e:
            self.lbl_info.setText(t("tool.split.error_info", e=e))

    def auto_load(self, path: str):
        if path and not self.drop_in.path(): self._load_input(path)

    def _ensure_tesseract(self):
        """Locate Tesseract, set TESSDATA_PREFIX and update available languages."""
        import pytesseract, sys
        tess_exe = _find_tesseract()

        if tess_exe:
            pytesseract.pytesseract.tesseract_cmd = tess_exe
        tessdata = _find_tessdata(tess_exe)
        if tessdata:
            os.environ["TESSDATA_PREFIX"] = tessdata

        try:
            pytesseract.get_tesseract_version()
        except Exception:
            if sys.platform == "win32":
                install_hint = "https://github.com/UB-Mannheim/tesseract/wiki"
            elif sys.platform == "darwin":
                install_hint = "brew install tesseract tesseract-lang"
            else:
                install_hint = "sudo apt install tesseract-ocr tesseract-ocr-por tesseract-ocr-eng"
            QMessageBox.critical(self, t("tool.ocr.tess_not_found"),
                t("tool.ocr.tess_msg", hint=install_hint))
            return None

        try:
            installed = pytesseract.get_languages(config="")
            self._update_lang_combo(installed)
        except Exception:
            pass

        return pytesseract

    def _update_lang_combo(self, installed: list):
        """Show only installed languages + possible combinations."""
        entries = []
        label_map = {"por": t("tool.ocr.lang.pt"), "eng": t("tool.ocr.lang.en"), "spa": t("tool.ocr.lang.es"),
                     "fra": t("tool.ocr.lang.fr"),     "deu": t("tool.ocr.lang.de"),  "ita": "Italian"}
        for code, label in [(c, label_map.get(c, c)) for c in installed if c != "osd"]:
            entries.append((label, code))
        if "por" in installed and "eng" in installed:
            entries.append((t("tool.ocr.lang.pt_en"), "por+eng"))
        if not entries:
            entries = [(t("tool.ocr.lang.en"), "eng")]
        current_codes = [e[1] for e in entries]
        self.cmb_lang.clear()
        for name, _ in entries:
            self.cmb_lang.addItem(name)
        self._lang_codes = current_codes

    def _run(self):
        pdf_path = self.drop_in.path()
        if not pdf_path or not os.path.isfile(pdf_path):
            QMessageBox.warning(self, t("msg.warning"), t("msg.select_valid_pdf")); return
        out_path = self._resolve_output_file(self.drop_out, pdf_path)
        if not out_path: return
        try:
            import pytesseract
        except ImportError:
            QMessageBox.critical(self, t("msg.missing_dep"), t("tool.ocr.dep_pytesseract"))
            return
        tess = self._ensure_tesseract()
        if tess is None: return
        try:
            import fitz  # noqa: F401 — surface ImportError before launching thread
        except ImportError:
            QMessageBox.critical(self, t("msg.missing_dep"), t("tool.ocr.dep_pymupdf"))
            return
        try:
            from PIL import Image  # noqa: F401
        except ImportError:
            QMessageBox.critical(self, t("msg.missing_dep"), t("tool.ocr.dep_pillow"))
            return

        codes = getattr(self, "_lang_codes", [c for _, c in self._LANGS])
        if not codes:
            QMessageBox.warning(self, t("msg.warning"), t("tool.ocr.no_langs"))
            return
        idx = max(0, min(self.cmb_lang.currentIndex(), len(codes) - 1))
        lang = codes[idx]
        fmt = self.cmb_fmt.currentIndex()
        pwd = self._pdf_password

        # Quickly count pages so the progress dialog has the correct max.
        try:
            with self._open_fitz(pdf_path) as _doc:
                n_pages = _doc.page_count
        except Exception as e:
            QMessageBox.critical(self, t("tool.ocr.error"), str(e)); return

        progress = QProgressDialog(t("progress.ocr.page", current=1, total=n_pages),
                                   t("progress.cancel"), 0, n_pages, self)
        progress.setWindowTitle(t("progress.ocr.title"))
        progress.setWindowModality(Qt.WindowModality.WindowModal)
        progress.setMinimumDuration(0)
        progress.setValue(0)

        class _OcrRunner(TaskRunner):
            def do_work(_self):
                import io as _io
                import fitz
                from PIL import Image
                import pytesseract
                doc = fitz.open(pdf_path)
                if doc.needs_pass and pwd:
                    doc.authenticate(pwd)
                try:
                    if fmt == 1:
                        texts = []
                        for i, page in enumerate(doc):
                            if _self.is_cancelled():
                                return None
                            _self.progress.emit(
                                i, t("progress.ocr.page",
                                     current=i + 1, total=n_pages))
                            pix = page.get_pixmap(dpi=300)
                            img = Image.frombytes(
                                "RGB", (pix.width, pix.height), pix.samples)
                            texts.append(
                                pytesseract.image_to_string(img, lang=lang))
                        with open(out_path, "w", encoding="utf-8") as fh:
                            fh.write("\f".join(texts))
                    else:
                        pdf_pages = []
                        for i, page in enumerate(doc):
                            if _self.is_cancelled():
                                return None
                            _self.progress.emit(
                                i, t("progress.ocr.page",
                                     current=i + 1, total=n_pages))
                            pix = page.get_pixmap(dpi=300)
                            img = Image.frombytes(
                                "RGB", (pix.width, pix.height), pix.samples)
                            pdf_pages.append(
                                pytesseract.image_to_pdf_or_hocr(
                                    img, lang=lang, extension="pdf"))
                        writer = PdfWriter()
                        for page_bytes in pdf_pages:
                            writer.append(PdfReader(_io.BytesIO(page_bytes)))
                        with open(out_path, "wb") as fh:
                            writer.write(fh)
                finally:
                    doc.close()
                return out_path

        self.action_btn.setEnabled(False)

        def _on_done(result):
            self.action_btn.setEnabled(True)
            if result is None:
                self._status(t("progress.cancelled"))
                return
            self._status(f"✔  OCR → {result}")
            QMessageBox.information(self, t("msg.done"),
                                    t("tool.ocr.done", path=result))

        def _on_err(msg):
            self.action_btn.setEnabled(True)
            QMessageBox.critical(self, t("tool.ocr.error"), msg)

        self._runner = _OcrRunner()
        self._runner_thread = run_task(self, self._runner, progress, _on_done, _on_err)
