import base64
import re
from collections import OrderedDict
from collections.abc import Generator
from typing import Any

import fitz
from dify_plugin import Tool
from dify_plugin.entities.tool import ToolInvokeMessage


class PdfToMarkdownTool(Tool):
    """Convert PDF to a single Markdown file. No LLM needed.

    - Auto-detect TOC and organize content by chapters.
    - Extract text and tables as Markdown.
    - Embed raster images as base64.
    - Render vector drawings as base64 PNG.
    - Output one .md file via create_blob_message.
    """

    _TOC_PATTERNS = [
        r"目录", r"目　录", r"目\u3000录",
        r"Table of Contents", r"Contents", r"目次",
    ]

    # ── entry point ──────────────────────────────────────────

    def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage]:
        file = tool_parameters.get("file")
        if not file:
            yield self.create_text_message("Error: file is required")
            return

        include_images = self._to_bool(tool_parameters.get("include_images"), True)
        image_dpi = self._to_int(tool_parameters.get("image_dpi"), 150)
        image_dpi = max(72, min(image_dpi, 300))
        max_image_bytes = 2 * 1024 * 1024  # skip images > 2 MB raw

        doc = fitz.open(stream=file.blob, filetype="pdf")
        try:
            num_pages = len(doc)

            # 1) Build chapter map (metadata TOC → printed TOC → none)
            chapters, content_offset = self._build_chapter_map(doc, num_pages)

            # 2) Convert every page
            page_mds: list[str] = []
            for idx in range(num_pages):
                md = self._page_to_markdown(
                    doc, doc[idx], idx,
                    include_images, image_dpi, max_image_bytes,
                )
                page_mds.append(md)

            # 3) Assemble
            if chapters:
                final_md = self._assemble_by_chapters(
                    chapters, page_mds, content_offset, num_pages,
                )
            else:
                final_md = "\n\n---\n\n".join(m for m in page_mds if m.strip())

            # 4) Output: text (for variable aggregation) + blob (.md file)
            yield self.create_text_message(final_md)
            md_bytes = final_md.encode("utf-8")
            yield self.create_blob_message(
                blob=md_bytes,
                meta={"mime_type": "text/markdown"},
            )
        finally:
            doc.close()

    # ── chapter detection ────────────────────────────────────

    def _build_chapter_map(
        self, doc: fitz.Document, num_pages: int,
    ) -> tuple[dict, int]:
        """Return (chapters_dict, content_offset).

        Try embedded PDF TOC metadata first (reliable page mapping).
        Fall back to scanning printed TOC pages.
        """
        toc = doc.get_toc()
        if toc:
            chapters = self._chapters_from_metadata(toc, num_pages)
            if chapters:
                return chapters, 0

        toc_start, toc_end = self._find_toc_pages(doc, num_pages)
        if toc_start is not None and toc_end is not None:
            toc_text = "\n".join(
                doc[i].get_text() or "" for i in range(toc_start, toc_end + 1)
            )
            chapters = self._parse_toc_lines(toc_text)
            if chapters:
                offset = self._guess_offset(chapters, toc_end)
                return chapters, offset

        return {}, 0

    def _chapters_from_metadata(
        self, toc: list, num_pages: int,
    ) -> dict[str, dict[str, int]]:
        top = [(t, max(0, p - 1)) for lvl, t, p in toc if lvl <= 2 and p >= 1]
        if not top:
            return {}
        chapters: dict[str, dict[str, int]] = OrderedDict()
        for i, (title, start) in enumerate(top):
            end = top[i + 1][1] - 1 if i + 1 < len(top) else num_pages - 1
            chapters[title] = {"start": start, "end": max(start, end)}
        return chapters

    def _find_toc_pages(self, doc, num_pages):
        toc_start = toc_end = None
        for pn in range(min(num_pages, 30)):
            text = doc[pn].get_text() or ""
            if any(re.search(p, text, re.IGNORECASE) for p in self._TOC_PATTERNS):
                if toc_start is None:
                    toc_start = pn
                toc_end = pn
            elif toc_start is not None:
                break
        return toc_start, toc_end

    def _parse_toc_lines(self, text: str) -> dict[str, dict[str, int]]:
        m = re.search(
            r"^(List\s+of\s+Figures|List\s+of\s+Tables|图目录|表目录)",
            text, re.IGNORECASE | re.MULTILINE,
        )
        if m:
            text = text[: m.start()]

        pat = re.compile(
            r"^\s*(?P<title>.+?)\s*(?:\.{2,}|\s)\s*(?P<page>\d{1,5})\s*$"
        )
        entries: list[tuple[str, int]] = []
        for raw in text.splitlines():
            line = raw.strip()
            if not line or len(line) < 3 or re.fullmatch(r"\d+", line):
                continue
            m2 = pat.match(line)
            if not m2:
                continue
            title = re.sub(r"\s+", " ", m2.group("title")).strip("-_:： ")
            page = self._to_int(m2.group("page"), None)
            if not title or page is None or len(title) <= 1:
                continue
            if title.lower() in {"page", "pages", "目录", "contents"}:
                continue
            entries.append((title, page))

        if not entries:
            return {}

        dedup: OrderedDict[str, int] = OrderedDict()
        for t, p in entries:
            dedup.setdefault(t, p)

        titles = list(dedup.keys())
        pages = [dedup[t] for t in titles]
        catalog: dict[str, dict[str, int]] = OrderedDict()
        for i, t in enumerate(titles):
            s = pages[i]
            e = max(s, pages[i + 1] - 1) if i + 1 < len(pages) else s
            catalog[t] = {"start": s, "end": e}
        return catalog

    @staticmethod
    def _guess_offset(chapters: dict, toc_end: int) -> int:
        first_page = None
        for info in chapters.values():
            s = info["start"]
            if first_page is None or s < first_page:
                first_page = s
        if first_page is None:
            return 0
        return (toc_end + 1) - first_page

    # ── per-page conversion ──────────────────────────────────

    def _page_to_markdown(
        self,
        doc: fitz.Document,
        page: fitz.Page,
        page_idx: int,
        include_images: bool,
        image_dpi: int,
        max_image_bytes: int,
    ) -> str:
        parts: list[str] = []

        # ── text ──
        text = (page.get_text("text", sort=True) or "").strip()
        if text:
            parts.append(text)

        # ── tables → Markdown ──
        try:
            for tab in (page.find_tables().tables or [])[:5]:
                cells = tab.extract() or []
                if len(cells) >= 2:
                    md = self._cells_to_md_table(cells)
                    if md:
                        parts.append(md)
        except Exception:
            pass

        if not include_images:
            return "\n\n".join(parts)

        # ── embedded raster images ──
        try:
            for img_idx, img_info in enumerate(page.get_images(full=True)):
                xref = img_info[0]
                try:
                    data = doc.extract_image(xref)
                    if not data or not data.get("image"):
                        continue
                    raw = data["image"]
                    if len(raw) > max_image_bytes:
                        continue
                    # skip tiny icons (< 20x20)
                    w = data.get("width", 0)
                    h = data.get("height", 0)
                    if w < 20 and h < 20:
                        continue
                    ext = data.get("ext", "png")
                    mime = "image/jpeg" if ext in ("jpg", "jpeg") else f"image/{ext}"
                    b64 = base64.b64encode(raw).decode("ascii")
                    parts.append(
                        f"![img-p{page_idx}-{img_idx}](data:{mime};base64,{b64})"
                    )
                except Exception:
                    pass
        except Exception:
            pass

        # ── vector drawings → render as PNG ──
        try:
            drawings = page.get_drawings()
            if len(drawings) >= 3:
                valid_rects: list[fitz.Rect] = []
                for d in drawings:
                    r = d.get("rect")
                    if r:
                        try:
                            rect = fitz.Rect(r)
                            if rect.is_valid and not rect.is_empty:
                                valid_rects.append(rect)
                        except Exception:
                            pass
                if valid_rects:
                    bbox = valid_rects[0]
                    for r in valid_rects[1:]:
                        bbox |= r
                    bbox &= page.rect
                    if bbox.width > 30 and bbox.height > 30:
                        scale = image_dpi / 72
                        mat = fitz.Matrix(scale, scale)
                        pix = page.get_pixmap(matrix=mat, clip=bbox)
                        png = pix.tobytes("png")
                        if len(png) <= max_image_bytes:
                            b64 = base64.b64encode(png).decode("ascii")
                            parts.append(
                                f"![drawing-p{page_idx}](data:image/png;base64,{b64})"
                            )
        except Exception:
            pass

        return "\n\n".join(parts)

    # ── assembly ─────────────────────────────────────────────

    def _assemble_by_chapters(
        self,
        chapters: dict[str, dict[str, int]],
        page_mds: list[str],
        offset: int,
        num_pages: int,
    ) -> str:
        parts: list[str] = []
        for name, info in chapters.items():
            s = info["start"] + offset
            e = info["end"] + offset
            s = max(0, min(s, num_pages - 1))
            e = max(s, min(e, num_pages - 1))
            ch: list[str] = [f"# {name}\n"]
            for idx in range(s, e + 1):
                if idx < len(page_mds) and page_mds[idx].strip():
                    ch.append(page_mds[idx])
            parts.append("\n\n".join(ch))
        return "\n\n---\n\n".join(parts)

    # ── helpers ──────────────────────────────────────────────

    @staticmethod
    def _cells_to_md_table(cells: list) -> str:
        if not cells:
            return ""
        header = cells[0]
        ncols = len(header)
        if ncols == 0:
            return ""
        clean = lambda c: str(c or "").replace("|", "\\|").replace("\n", " ").strip()
        lines = [
            "| " + " | ".join(clean(c) for c in header) + " |",
            "| " + " | ".join("---" for _ in range(ncols)) + " |",
        ]
        for row in cells[1:]:
            padded = list(row) + [""] * max(0, ncols - len(row))
            lines.append("| " + " | ".join(clean(c) for c in padded[:ncols]) + " |")
        return "\n".join(lines)

    @staticmethod
    def _to_int(value: Any, default: int | None) -> int | None:
        try:
            if value is None or value == "":
                return default
            return int(value)
        except Exception:
            return default

    @staticmethod
    def _to_bool(value: Any, default: bool) -> bool:
        if value is None:
            return default
        if isinstance(value, bool):
            return value
        s = str(value).strip().lower()
        if s in {"1", "true", "yes", "on"}:
            return True
        if s in {"0", "false", "no", "off"}:
            return False
        return default