from collections.abc import Generator
from io import BytesIO
from typing import Any

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


class PdfSinglePageTool(Tool):
    def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage]:
        file = tool_parameters.get("file")
        page = tool_parameters.get("page", 0)

        if not file:
            yield self.create_text_message("Error: file is required")
            return

        pdf_bytes = file.blob
        reader = PyPDF2.PdfReader(BytesIO(pdf_bytes))
        num_pages = len(reader.pages)

        page_index = int(page)
        if page_index < 0:
            page_index = 0
        if page_index >= num_pages:
            page_index = num_pages - 1

        selected_page = reader.pages[page_index]
        text = selected_page.extract_text() or ""

        yield self.create_json_message({
            "start": page_index,
            "end": page_index,
            "pages": [text]
        })