# -*- coding: utf-8 -*-
# filename: pdf.py
# @Time    : 2025-01-20
# @Author  : Claude Code
# @Email   : noreply@anthropic.com
# @Software: PyCharm
"""PDF file handler for file preview tool"""
import logging
import tempfile
from pathlib import Path

import pypdf

from tf_tools_mcp.tools.file_preview_tools.file_preview_constants import (
    get_prompt_template_for_file_type,
)
from tf_tools_mcp.tools.file_preview_tools.file_preview_types import FilePreviewOptions
from tf_tools_mcp.tools.file_preview_tools.handlers.base import (
    BaseFileHandler,
    format_file_size,
)

logger = logging.getLogger(__name__)


class PDFHandler(BaseFileHandler):
    """PDF 处理器"""

    supported_extensions = {".pdf"}
    supported_mime_types = {"application/pdf"}

    async def extract_metadata(self, file_path: str) -> dict:
        """提取 PDF 元数据

        Args:
            file_path: PDF 文件路径

        Returns:
            PDF 元数据字典
        """
        with open(file_path, "rb") as f:
            pdf_reader = pypdf.PdfReader(f)
            page_count = len(pdf_reader.pages)

            # 提取前 N 页文本
            preview_text = ""
            max_pages = min(self.options.max_pages, page_count)

            for i in range(max_pages):
                page = pdf_reader.pages[i]
                try:
                    text = page.extract_text()
                    if text:
                        preview_text += f"\n--- 第 {i + 1} 页 ---\n{text}"
                except Exception as e:
                    logger.warning(f"Failed to extract text from page {i + 1}: {e}")

            return {
                "page_count": page_count,
                "preview_text": preview_text.strip(),
            }

    async def generate_summary(self, file_path: str) -> str | None:
        """生成 PDF 内容描述（强制使用PDF转图片+视觉模型）

        Args:
            file_path: PDF 文件路径

        Returns:
            AI生成的PDF分析内容或文本预览（不包含文件基本信息）
        """
        # 导入依赖
        try:
            from pdf2image import convert_from_path
        except ImportError:
            logger.warning("pdf2image not installed, using text extraction as fallback")
            # 回退到文本提取
            format_metadata = await self.extract_metadata(file_path)
            preview_text = format_metadata.get("preview_text", "")

            if preview_text:
                # 限制内容预览长度
                content = preview_text[:1000]
                if len(preview_text) > 1000:
                    content += "..."
                return f"内容预览：\n{content}"
            return None

        # 导入 ImageUnderstandingTool
        from tf_tools_mcp.tools.file_preview_tools.image_understanding import (
            ImageUnderstandingTool,
        )

        # 自动选择 prompt 模板
        prompt_template = get_prompt_template_for_file_type("pdf")

        try:
            # 将 PDF 前几页转换为图片
            max_pages = self.options.max_pages
            images = convert_from_path(file_path, first_page=1, last_page=max_pages)

            # 保存为临时文件
            temp_files = []
            for i, img in enumerate(images):
                temp_file = tempfile.NamedTemporaryFile(
                    suffix=f"_page_{i + 1}.png", delete=False
                )
                img.save(temp_file.name, "PNG")
                temp_files.append(temp_file.name)

            try:
                # 调用 ImageUnderstandingTool，限制max_tokens确保输出≤1000字
                tool = ImageUnderstandingTool(model="glm-4.6v-flashx", max_tokens=600)
                result = tool.run(images=temp_files, prompt_template=prompt_template)

                # 只返回AI生成的描述内容
                return result.description[:1000]
            finally:
                # 清理临时文件
                import os

                for temp_file in temp_files:
                    try:
                        os.unlink(temp_file)
                    except Exception:
                        pass

        except Exception as e:
            logger.error(f"Failed to generate PDF summary: {e}")
            # 回退到文本提取，保持与 ImportError 分支一致的格式
            format_metadata = await self.extract_metadata(file_path)
            preview_text = format_metadata.get("preview_text", "")

            if preview_text:
                # 限制内容预览长度
                content = preview_text[:1000]
                if len(preview_text) > 1000:
                    content += "..."
                return f"内容预览：\n{content}"
            return None
