from cognite.client import CogniteClient
from mcp.server.fastmcp import FastMCP

from cog_mcp.config import Config
from cog_mcp.formatting import format_document_answer, format_document_summary
from cog_mcp.parsing import parse_json_param, parse_node_id, parse_node_ids


def register_document_tools(mcp: FastMCP, client: CogniteClient, config: Config) -> None:
    """Register AI document tools (question answering, summarization)."""

    @mcp.tool()
    def ask_documents_question(
        question: str,
        file_instance_ids: str | list | dict,
        language: str = "English",
        additional_context: str | None = None,
        ignore_unknown_ids: bool = False,
    ) -> str:
        """Ask a question about one or more PDF documents using document instance IDs.

        This tool only supports document instance IDs to stay consistent with other
        data modeling tools in this MCP server.

        Args:
            question: The natural-language question to ask.
            file_instance_ids: JSON array (or single object) of document instance IDs:
                [{"space": "my_space", "externalId": "doc_123"}]
            language: Answer language. Default "English".
            additional_context: Optional additional context for the LLM.
            ignore_unknown_ids: If true, skip missing/unprocessed documents.
        """
        instance_ids = parse_node_ids(file_instance_ids, field_name="file_instance_ids")

        answer = client.ai.tools.documents.ask_question(
            question=question,
            instance_id=instance_ids,
            language=language,
            additional_context=additional_context,
            ignore_unknown_ids=ignore_unknown_ids,
        )
        return format_document_answer(answer)

    @mcp.tool()
    def summarize_document(
        instance_id: str | dict,
    ) -> str:
        """Summarize a single PDF document using document instance ID.

        Args:
            instance_id: JSON object with document instance ID:
                {"space": "my_space", "externalId": "doc_123"}
        """
        if instance_id is None:
            raise ValueError("instance_id is required")
        parsed_instance_id = parse_node_id(parse_json_param(instance_id), "instance_id")

        summary = client.ai.tools.documents.summarize(
            instance_id=parsed_instance_id,
        )
        return format_document_summary(summary)
