"""Builder for LangChain tools from DataHub MCP tools."""

from typing import TYPE_CHECKING

from datahub_agent_context.mcp_tools import get_me
from datahub_agent_context.mcp_tools.assertions import get_dataset_assertions
from datahub_agent_context.mcp_tools.documents import grep_documents, search_documents
from datahub_agent_context.mcp_tools.domains import remove_domains, set_domains
from datahub_agent_context.mcp_tools.structured_properties import (
    add_structured_properties,
    remove_structured_properties,
)
from datahub_agent_context.utils import create_context_wrapper

if TYPE_CHECKING:
    from datahub.sdk.main_client import DataHubClient

try:
    from langchain_core.tools import tool  # type: ignore[import-not-found]
    from langchain_core.tools.base import BaseTool  # type: ignore[import-not-found]
except ImportError as e:
    raise ImportError(
        "langchain-core is required for LangChain tools. "
        "Install with: pip install 'datahub-agent-context[langchain]'"
    ) from e

from datahub_agent_context.mcp_tools.descriptions import update_description
from datahub_agent_context.mcp_tools.entities import get_entities, list_schema_fields
from datahub_agent_context.mcp_tools.lineage import (
    get_lineage,
    get_lineage_paths_between,
)
from datahub_agent_context.mcp_tools.owners import add_owners, remove_owners
from datahub_agent_context.mcp_tools.queries import get_dataset_queries
from datahub_agent_context.mcp_tools.save_document import save_document
from datahub_agent_context.mcp_tools.search import search
from datahub_agent_context.mcp_tools.tags import add_tags, remove_tags
from datahub_agent_context.mcp_tools.terms import (
    add_glossary_terms,
    remove_glossary_terms,
)


def build_langchain_tools(
    client: "DataHubClient",
    include_mutations: bool = False,
) -> list[BaseTool]:
    """Build LangChain tools with automatic context management.

    Each tool is wrapped to automatically set the DataHubGraph in context
    before execution, allowing tools to retrieve it using get_graph().

    Args:
        client: DataHubClient instance
        include_mutations: Whether to include mutation tools (default: False)

    Returns:
        List of LangChain BaseTool instances

    Example:
        from datahub.sdk.main_client import DataHubClient
        from datahub_agent_context.langchain_tools import build_langchain_tools

        client = DataHubClient(...)
        tools = build_langchain_tools(client, include_mutations=True)

        # Use with LangChain agents - context is managed automatically
        agent = create_react_agent(llm, tools, prompt)
        result = agent.invoke({"input": "search for datasets"})
    """
    tools = []

    # Wrap each tool to manage context automatically
    tools.append(tool(create_context_wrapper(search_documents, client)))
    tools.append(tool(create_context_wrapper(grep_documents, client)))

    tools.append(tool(create_context_wrapper(get_entities, client)))
    tools.append(tool(create_context_wrapper(list_schema_fields, client)))
    tools.append(tool(create_context_wrapper(get_me, client)))
    tools.append(tool(create_context_wrapper(get_lineage, client)))
    tools.append(tool(create_context_wrapper(get_lineage_paths_between, client)))

    tools.append(tool(create_context_wrapper(get_dataset_queries, client)))
    tools.append(tool(create_context_wrapper(get_dataset_assertions, client)))
    tools.append(tool(create_context_wrapper(search, client)))

    if include_mutations:
        tools.append(tool(create_context_wrapper(update_description, client)))
        tools.append(tool(create_context_wrapper(set_domains, client)))
        tools.append(tool(create_context_wrapper(remove_domains, client)))
        tools.append(tool(create_context_wrapper(add_owners, client)))
        tools.append(tool(create_context_wrapper(remove_owners, client)))
        tools.append(tool(create_context_wrapper(add_structured_properties, client)))
        tools.append(tool(create_context_wrapper(remove_structured_properties, client)))
        tools.append(tool(create_context_wrapper(add_tags, client)))
        tools.append(tool(create_context_wrapper(remove_tags, client)))
        tools.append(tool(create_context_wrapper(add_glossary_terms, client)))
        tools.append(tool(create_context_wrapper(remove_glossary_terms, client)))
        tools.append(tool(create_context_wrapper(save_document, client)))

    return tools
