"""Tools for manipulating Docling documents."""

from dataclasses import dataclass
from typing import Annotated

from pydantic import Field

from docling_core.types.doc.document import (
    DocItem,
    GroupItem,
    RefItem,
    SectionHeaderItem,
    TextItem,
    TitleItem,
)

from docling_mcp.logger import setup_logger
from docling_mcp.shared import local_document_cache, mcp

# Create a default project logger
logger = setup_logger()


# TODO: Provide a proper structure instead of a single string
@dataclass
class DocumentAnchorOutput:
    """Output of the get_overview_of_document_anchors tool."""

    structure: Annotated[
        str,
        Field(
            description=(
                "A string containing the hierarchical structure of the document with "
                "indentation to show nesting levels, along with anchor references."
            )
        ),
    ]


@mcp.tool(title="Get overview of Docling document anchors")
def get_overview_of_document_anchors(
    document_key: Annotated[
        str,
        Field(description="The unique identifier of the document in the local cache."),
    ],
) -> DocumentAnchorOutput:
    """Retrieve a structured overview of a document from the local document cache.

    This tool returns a text representation of the Docling document's structure,
    showing the hierarchy and types of elements within the document. Each line in the
    output includes the document anchor reference and item label.
    """
    if document_key not in local_document_cache:
        doc_keys = ", ".join(local_document_cache.keys())
        raise ValueError(
            f"document-key: {document_key} is not found. Existing document-keys are: "
            f"{doc_keys}"
        )

    doc = local_document_cache[document_key]

    lines = []
    slevel = 0
    for item, level in doc.iterate_items():
        ref = item.get_ref()

        if isinstance(item, DocItem):
            if isinstance(item, TitleItem):
                lines.append(f"[anchor:{ref.cref}] {item.label}: {item.text}")

            elif isinstance(item, SectionHeaderItem):
                slevel = item.level
                indent = "  " * (level + slevel)
                lines.append(
                    f"{indent}[anchor:{ref.cref}] {item.label}-{level}: {item.text}"
                )

            else:
                indent = "  " * (level + slevel + 1)
                lines.append(f"{indent}[anchor:{ref.cref}] {item.label}")

        elif isinstance(item, GroupItem):
            indent = "  " * (level + slevel + 1)
            lines.append(f"{indent}[anchor:{ref.cref}] {item.label}")

    return DocumentAnchorOutput("\n".join(lines))


@dataclass
class DocumentItemText:
    """Text content of a Docling document item."""

    text: Annotated[
        str,
        Field(description="The text content of an item in a Docling document."),
    ]


@mcp.tool(title="Get text of Docling document item at anchor")
def get_text_of_document_item_at_anchor(
    document_key: Annotated[
        str,
        Field(description="The unique identifier of the document in the local cache."),
    ],
    document_anchor: Annotated[
        str,
        Field(
            description=(
                "The anchor reference that identifies the specific item within the "
                "document."
            ),
            examples=["#/texts/2"],
        ),
    ],
) -> DocumentItemText:
    """Retrieve the text content of a specific document item identified by its anchor.

    This tool extracts the text from a Docling document item at the specified anchor
    location within a document that exists in the local document cache.
    """
    if document_key not in local_document_cache:
        doc_keys = ", ".join(local_document_cache.keys())
        raise ValueError(
            f"document-key: {document_key} is not found. Existing document-keys are: {doc_keys}"
        )

    doc = local_document_cache[document_key]

    ref = RefItem(cref=document_anchor)
    item = ref.resolve(doc=doc)

    if isinstance(item, TextItem):
        text = item.text
    else:
        raise ValueError(
            f"Item at {document_anchor} for document-key: {document_key} is not a textual item."
        )

    return DocumentItemText(text)


@mcp.tool(title="Update text of Docling document item at anchor")
def update_text_of_document_item_at_anchor(
    document_key: Annotated[
        str,
        Field(description="The unique identifier of the document in the local cache."),
    ],
    document_anchor: Annotated[
        str,
        Field(
            description=(
                "The anchor reference that identifies the specific item within the "
                "document."
            ),
            examples=["#/texts/6"],
        ),
    ],
    updated_text: Annotated[
        str,
        Field(description="The new text content to replace the existing content."),
    ],
) -> bool:
    """Update the text content of a specific document item identified by its anchor.

    This tool modifies the text of an existing document item at the specified anchor
    location within a document that exists in the local document cache. It returns
    True if the update was successful.
    """
    if document_key not in local_document_cache:
        doc_keys = ", ".join(local_document_cache.keys())
        raise ValueError(
            f"document-key: {document_key} is not found. Existing document-keys are: "
            f"{doc_keys}"
        )

    doc = local_document_cache[document_key]

    ref = RefItem(cref=document_anchor)
    item = ref.resolve(doc=doc)

    if isinstance(item, TextItem):
        item.text = updated_text
    else:
        raise ValueError(
            f"Item at {document_anchor} for document-key: {document_key} is not a "
            "textual item."
        )

    return True


@mcp.tool(title="Delete Docling document items at anchors")
def delete_document_items_at_anchors(
    document_key: Annotated[
        str,
        Field(description="The unique identifier of the document in the local cache."),
    ],
    document_anchors: Annotated[
        list[str],
        Field(
            description=(
                "A list of anchor references identifying the items to be deleted from the "
                "document."
            ),
            examples=["#/texts/2", "#/tables/1"],
        ),
    ],
) -> bool:
    """Delete multiple document items identified by their anchors.

    This tool removes specified items from a Docling document that exists in the local
    document cache, based on their anchor references. It returns True if the all the
    items were successfully removed.
    """
    if document_key not in local_document_cache:
        doc_keys = ", ".join(local_document_cache.keys())
        raise ValueError(
            f"document-key: {document_key} is not found. Existing document-keys are: {doc_keys}"
        )

    doc = local_document_cache[document_key]

    items = []
    for _ in document_anchors:
        ref = RefItem(cref=_)
        items.append(ref.resolve(doc=doc))

    doc.delete_items(node_items=items)

    return True
