"""PDF parsing and image extraction utilities for kbase."""

from __future__ import annotations

import base64
import contextlib
import json
import logging
import shutil
import time
from pathlib import Path
from typing import Any, Dict, Literal, Mapping, Optional

logger = logging.getLogger(__name__)

ParserBackend = Literal["auto", "docling", "pymupdf"]
ImageBackend = Literal["auto", "pymupdf", "pypdfium2", "none"]
OCRBackend = Literal["none", "tesseract", "rapidocr"]

DEFAULT_FALLBACK_POLICY: Dict[str, Any] = {
    "min_text_chars": 180,
    "render_on_no_images": True,
    "max_render_pages": 2,
}


def _ensure_backend(name: str, value: str, choices: set[str]) -> None:
    if value not in choices:
        allowed = ", ".join(sorted(choices))
        raise ValueError(f"{name} 必须为以下之一: {allowed}，当前: {value}")


def _import_docling_converter():
    from docling.document_converter import DocumentConverter

    return DocumentConverter


@contextlib.contextmanager
def _docling_artifacts_ctx(docling_artifacts_path: Optional[str]):
    if not docling_artifacts_path:
        yield
        return

    artifacts_dir = Path(docling_artifacts_path).expanduser()
    if not artifacts_dir.exists() or not artifacts_dir.is_dir():
        raise FileNotFoundError(f"Docling 模型目录不存在: {artifacts_dir}")

    from docling.datamodel.settings import settings

    previous = settings.artifacts_path
    settings.artifacts_path = artifacts_dir
    try:
        yield
    finally:
        settings.artifacts_path = previous


def _parse_text_with_docling(
    file_path: Path,
    docling_artifacts_path: Optional[str] = None,
) -> Dict[str, Any]:
    converter_cls = _import_docling_converter()
    with _docling_artifacts_ctx(docling_artifacts_path):
        converter = converter_cls()
        result = converter.convert(str(file_path))
    doc = result.document
    title = getattr(doc, "title", None) or file_path.stem
    pages = len(getattr(doc, "pages", []) or [])
    text = doc.export_to_markdown() or ""
    return {
        "text": text,
        "metadata": {
            "title": title,
            "pages": pages,
            "source": str(file_path),
        },
    }


def _parse_text_with_pymupdf(file_path: Path) -> Dict[str, Any]:
    import fitz

    text_parts: list[str] = []
    with fitz.open(str(file_path)) as doc:
        for page in doc:
            text_parts.append(page.get_text("text") or "")
        metadata = doc.metadata or {}
        title = metadata.get("title") or file_path.stem
        pages = len(doc)
    return {
        "text": "\n\n".join(part.strip() for part in text_parts if part.strip()),
        "metadata": {
            "title": title,
            "pages": pages,
            "source": str(file_path),
        },
    }


def _extract_images_with_pymupdf(
    file_path: str | Path,
    output_dir: str | Path,
    return_base64: bool = False,
    save_images: bool = True,
) -> list[Dict[str, Any]]:
    """Extract embedded images with PyMuPDF."""
    import fitz

    src = Path(file_path)
    out = Path(output_dir)
    if save_images:
        out.mkdir(parents=True, exist_ok=True)

    results: list[Dict[str, Any]] = []
    seen_xrefs: set[int] = set()

    with fitz.open(str(src)) as doc:
        for page_idx, page in enumerate(doc):
            page_images = page.get_images(full=True)
            for image_idx, image_info in enumerate(page_images):
                xref = int(image_info[0])
                if xref in seen_xrefs:
                    continue
                seen_xrefs.add(xref)
                extracted = doc.extract_image(xref)
                if not extracted:
                    continue

                image_bytes = extracted.get("image", b"")
                image_ext = (extracted.get("ext") or "png").lower()
                image_name = f"page_{page_idx + 1:04d}_{image_idx + 1:03d}.{image_ext}"
                image_path = out / image_name

                image_item: Dict[str, Any] = {
                    "page": page_idx + 1,
                    "index": image_idx,
                    "ext": image_ext,
                    "width": extracted.get("width"),
                    "height": extracted.get("height"),
                    "size_bytes": len(image_bytes),
                }

                if save_images:
                    image_path.write_bytes(image_bytes)
                    image_item["path"] = str(image_path)

                if return_base64:
                    image_item["base64"] = base64.b64encode(image_bytes).decode("ascii")

                results.append(image_item)
    return results


def _extract_images_with_pypdfium2(
    file_path: str | Path,
    output_dir: str | Path,
    return_base64: bool = False,
    save_images: bool = True,
    max_pages: int = 2,
) -> list[Dict[str, Any]]:
    """Render PDF pages as fallback images with pypdfium2."""
    import pypdfium2 as pdfium

    src = Path(file_path)
    out = Path(output_dir)
    if save_images:
        out.mkdir(parents=True, exist_ok=True)

    results: list[Dict[str, Any]] = []
    document = pdfium.PdfDocument(str(src))
    total_pages = len(document)
    limit = min(max_pages, total_pages)

    for page_idx in range(limit):
        page = document[page_idx]
        bitmap = page.render(scale=2.0)
        image_name = f"rendered_page_{page_idx + 1:04d}.png"
        image_path = out / image_name

        image_bytes: bytes = b""
        pil_image = None
        try:
            pil_image = bitmap.to_pil()
            if save_images:
                pil_image.save(str(image_path), format="PNG")
                image_bytes = image_path.read_bytes()
            elif return_base64:
                from io import BytesIO

                memory = BytesIO()
                pil_image.save(memory, format="PNG")
                image_bytes = memory.getvalue()
        except Exception as exc:  # pragma: no cover - depends on PIL runtime
            logger.warning("pypdfium2 渲染保存失败: %s", exc)
            continue
        finally:
            if pil_image is not None:
                pil_image.close()
            bitmap.close()
            page.close()

        item: Dict[str, Any] = {
            "page": page_idx + 1,
            "index": 0,
            "ext": "png",
            "path": str(image_path) if save_images else "",
            "size_bytes": len(image_bytes),
            "rendered": True,
        }
        if return_base64 and image_bytes:
            item["base64"] = base64.b64encode(image_bytes).decode("ascii")
        results.append(item)

    document.close()
    return results


def _merge_fallback_policy(policy: Optional[Mapping[str, Any]]) -> Dict[str, Any]:
    merged = dict(DEFAULT_FALLBACK_POLICY)
    if policy:
        merged.update(dict(policy))
    return merged


def _parse_text(
    file_path: Path,
    parser_backend: ParserBackend,
    fallback_policy: Mapping[str, Any],
    docling_artifacts_path: Optional[str] = None,
) -> tuple[str, Dict[str, Any], str]:
    if parser_backend == "docling":
        parsed = _parse_text_with_docling(file_path, docling_artifacts_path=docling_artifacts_path)
        return parsed["text"], parsed["metadata"], "docling"
    if parser_backend == "pymupdf":
        parsed = _parse_text_with_pymupdf(file_path)
        return parsed["text"], parsed["metadata"], "pymupdf"

    # auto strategy: docling first, fallback to pymupdf
    min_text_chars = int(fallback_policy.get("min_text_chars", 180))
    try:
        parsed = _parse_text_with_docling(file_path, docling_artifacts_path=docling_artifacts_path)
        text = parsed["text"]
        if len(text.strip()) >= min_text_chars:
            return text, parsed["metadata"], "docling"
        logger.info("Docling 文本较短，切换到 PyMuPDF: %s", file_path)
    except Exception as exc:
        logger.info("Docling 不可用或解析失败，切换到 PyMuPDF: %s", exc)

    parsed = _parse_text_with_pymupdf(file_path)
    return parsed["text"], parsed["metadata"], "pymupdf"


def parse_single_pdf(
    file_path: str,
    output_dir: Optional[str] = None,
    *,
    parser_backend: ParserBackend = "auto",
    extract_images: bool = True,
    image_backend: ImageBackend = "auto",
    save_images: bool = True,
    return_base64: bool = False,
    ocr_backend: OCRBackend = "none",
    fallback_policy: Optional[Mapping[str, Any]] = None,
    docling_artifacts_path: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Parse one PDF and optionally persist markdown/json outputs.

    Raises:
        FileNotFoundError: when file does not exist.
        ValueError: when backend option is invalid.
    """
    _ensure_backend("parser_backend", parser_backend, {"auto", "docling", "pymupdf"})
    _ensure_backend("image_backend", image_backend, {"auto", "pymupdf", "pypdfium2", "none"})
    _ensure_backend("ocr_backend", ocr_backend, {"none", "tesseract", "rapidocr"})

    source = Path(file_path)
    if not source.exists():
        raise FileNotFoundError(f"PDF 文件不存在: {source}")
    if source.suffix.lower() != ".pdf":
        raise ValueError(f"仅支持 PDF 文件，当前: {source}")

    merged_policy = _merge_fallback_policy(fallback_policy)
    started_at = time.time()

    text, metadata, parser_used = _parse_text(
        source,
        parser_backend,
        merged_policy,
        docling_artifacts_path=docling_artifacts_path,
    )
    image_list: list[Dict[str, Any]] = []
    image_backend_used = "none"

    if extract_images and image_backend != "none":
        image_out_dir = Path(output_dir) / "images" if output_dir else source.parent / f"{source.stem}_images"
        selected_backend = "pymupdf" if image_backend == "auto" else image_backend
        try:
            if selected_backend == "pymupdf":
                image_list = _extract_images_with_pymupdf(
                    source,
                    image_out_dir,
                    return_base64=return_base64,
                    save_images=save_images,
                )
                image_backend_used = "pymupdf"
            elif selected_backend == "pypdfium2":
                image_list = _extract_images_with_pypdfium2(
                    source,
                    image_out_dir,
                    return_base64=return_base64,
                    save_images=save_images,
                    max_pages=int(merged_policy.get("max_render_pages", 2)),
                )
                image_backend_used = "pypdfium2"
        except Exception as exc:
            logger.warning("图片提取失败: %s", exc)
            image_list = []

        # auto mode fallback: render pages when no embedded image is found
        render_on_no_images = bool(merged_policy.get("render_on_no_images", True))
        if image_backend == "auto" and not image_list and render_on_no_images:
            try:
                image_list = _extract_images_with_pypdfium2(
                    source,
                    image_out_dir,
                    return_base64=return_base64,
                    save_images=save_images,
                    max_pages=int(merged_policy.get("max_render_pages", 2)),
                )
                image_backend_used = "pypdfium2"
            except Exception as exc:
                logger.info("pypdfium2 回退未启用: %s", exc)

    elapsed = time.time() - started_at
    result: Dict[str, Any] = {
        "text": text,
        "images": image_list,
        "tables": [],
        "metadata": {
            **metadata,
            "parser_backend": parser_used,
            "image_backend": image_backend_used,
            "ocr_backend": ocr_backend,
            "elapsed_seconds": round(elapsed, 4),
            "text_length": len(text),
            "image_count": len(image_list),
            "docling_artifacts_path": docling_artifacts_path or "",
        },
    }

    if output_dir:
        out = Path(output_dir)
        out.mkdir(parents=True, exist_ok=True)
        stem = source.stem
        (out / f"{stem}_text.md").write_text(text, encoding="utf-8")
        (out / f"{stem}_parsed.json").write_text(
            json.dumps(result, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )

    return result


def parse_pdfs(
    input_dir: str,
    output_dir: str,
    *,
    parser_backend: ParserBackend = "auto",
    extract_images: bool = True,
    image_backend: ImageBackend = "auto",
    save_images: bool = True,
    return_base64: bool = False,
    ocr_backend: OCRBackend = "none",
    fallback_policy: Optional[Mapping[str, Any]] = None,
    move_source_pdf: bool = False,
    docling_artifacts_path: Optional[str] = None,
) -> Dict[str, Any]:
    """Parse all PDFs in a directory and emit structured folder outputs."""
    source_dir = Path(input_dir)
    target_dir = Path(output_dir)

    if not source_dir.exists():
        raise FileNotFoundError(f"输入目录不存在: {source_dir}")

    target_dir.mkdir(parents=True, exist_ok=True)
    pdf_files = sorted(source_dir.glob("*.pdf"))
    if not pdf_files:
        return {
            "total": 0,
            "success": 0,
            "failed": 0,
            "output_dir": str(target_dir),
            "items": [],
        }

    success = 0
    failed = 0
    items: list[Dict[str, Any]] = []

    for index, pdf_path in enumerate(pdf_files, start=1):
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        doc_folder = target_dir / f"{index:02d}_paper_{timestamp}"
        try:
            parsed = parse_single_pdf(
                str(pdf_path),
                output_dir=str(doc_folder),
                parser_backend=parser_backend,
                extract_images=extract_images,
                image_backend=image_backend,
                save_images=save_images,
                return_base64=return_base64,
                ocr_backend=ocr_backend,
                fallback_policy=fallback_policy,
                docling_artifacts_path=docling_artifacts_path,
            )
            destination_pdf = doc_folder / pdf_path.name
            if move_source_pdf:
                shutil.move(str(pdf_path), str(destination_pdf))
            else:
                shutil.copy2(str(pdf_path), str(destination_pdf))
            success += 1
            items.append(
                {
                    "source": str(pdf_path),
                    "output": str(doc_folder),
                    "status": "ok",
                    "metadata": parsed.get("metadata", {}),
                }
            )
        except Exception as exc:
            failed += 1
            logger.exception("解析失败: %s", pdf_path)
            items.append(
                {
                    "source": str(pdf_path),
                    "output": str(doc_folder),
                    "status": "failed",
                    "error": str(exc),
                }
            )

    return {
        "total": len(pdf_files),
        "success": success,
        "failed": failed,
        "output_dir": str(target_dir),
        "items": items,
    }


__all__ = [
    "parse_pdfs",
    "parse_single_pdf",
    "_extract_images_with_pymupdf",
]
