dbt-mcp

import logging
from dataclasses import dataclass

from mcp.server.fastmcp import FastMCP
from pydantic import Field

from dbt_mcp.config.config_providers import ConfigProvider, DiscoveryConfig
from dbt_mcp.discovery.client import (
    AppliedResourceType,
    ExposuresFetcher,
    LineageFetcher,
    MacrosFetcher,
    MetadataAPIClient,
    ModelPerformanceFetcher,
    ModelsFetcher,
    PaginatedResourceFetcher,
    ResourceDetailsFetcher,
    SourcesFetcher,
)
from dbt_mcp.prompts.prompts import get_prompt
from dbt_mcp.tools.definitions import dbt_mcp_tool
from dbt_mcp.tools.fields import (
    DEPTH_FIELD,
    NAME_FIELD,
    TYPES_FIELD,
    UNIQUE_ID_FIELD,
    UNIQUE_ID_REQUIRED_FIELD,
)
from dbt_mcp.tools.parameters import LineageResourceType
from dbt_mcp.tools.register import register_tools
from dbt_mcp.tools.tool_names import ToolName
from dbt_mcp.tools.toolsets import Toolset

logger = logging.getLogger(__name__)


@dataclass
class DiscoveryToolContext:
    models_fetcher: ModelsFetcher
    exposures_fetcher: ExposuresFetcher
    sources_fetcher: SourcesFetcher
    macros_fetcher: MacrosFetcher
    resource_details_fetcher: ResourceDetailsFetcher
    lineage_fetcher: LineageFetcher
    model_performance_fetcher: ModelPerformanceFetcher

    def __init__(self, config_provider: ConfigProvider[DiscoveryConfig]):
        api_client = MetadataAPIClient(config_provider=config_provider)
        self.models_fetcher = ModelsFetcher(
            api_client=api_client,
            paginator=PaginatedResourceFetcher(
                api_client=api_client,
                edges_path=("data", "environment", "applied", "models", "edges"),
                page_info_path=("data", "environment", "applied", "models", "pageInfo"),
            ),
        )
        self.exposures_fetcher = ExposuresFetcher(
            api_client=api_client,
            paginator=PaginatedResourceFetcher(
                api_client=api_client,
                edges_path=("data", "environment", "definition", "exposures", "edges"),
                page_info_path=(
                    "data",
                    "environment",
                    "definition",
                    "exposures",
                    "pageInfo",
                ),
            ),
        )
        self.sources_fetcher = SourcesFetcher(
            api_client=api_client,
            paginator=PaginatedResourceFetcher(
                api_client,
                edges_path=("data", "environment", "applied", "sources", "edges"),
                page_info_path=(
                    "data",
                    "environment",
                    "applied",
                    "sources",
                    "pageInfo",
                ),
            ),
        )
        self.macros_fetcher = MacrosFetcher(
            api_client=api_client,
            paginator=PaginatedResourceFetcher(
                api_client,
                edges_path=("data", "environment", "applied", "resources", "edges"),
                page_info_path=(
                    "data",
                    "environment",
                    "applied",
                    "resources",
                    "pageInfo",
                ),
            ),
        )
        self.resource_details_fetcher = ResourceDetailsFetcher(api_client=api_client)
        self.lineage_fetcher = LineageFetcher(api_client=api_client)
        self.model_performance_fetcher = ModelPerformanceFetcher(
            api_client=api_client,
            resource_details_fetcher=self.resource_details_fetcher,
        )


@dbt_mcp_tool(
    description=get_prompt("discovery/get_mart_models"),
    title="Get Mart Models",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_mart_models(context: DiscoveryToolContext) -> list[dict]:
    mart_models = await context.models_fetcher.fetch_models(
        model_filter={"modelingLayer": "marts"}
    )
    return [m for m in mart_models if m["name"] != "metricflow_time_spine"]


@dbt_mcp_tool(
    description=get_prompt("discovery/get_all_models"),
    title="Get All Models",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_all_models(context: DiscoveryToolContext) -> list[dict]:
    return await context.models_fetcher.fetch_models()


@dbt_mcp_tool(
    description=get_prompt("discovery/get_model_details"),
    title="Get Model Details",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_model_details(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
) -> list[dict]:
    return await context.resource_details_fetcher.fetch_details(
        resource_type=AppliedResourceType.MODEL,
        unique_id=unique_id,
        name=name,
    )


@dbt_mcp_tool(
    description=get_prompt("discovery/get_model_parents"),
    title="Get Model Parents",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_model_parents(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
) -> list[dict]:
    return await context.models_fetcher.fetch_model_parents(name, unique_id)


@dbt_mcp_tool(
    description=get_prompt("discovery/get_model_children"),
    title="Get Model Children",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_model_children(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
) -> list[dict]:
    return await context.models_fetcher.fetch_model_children(name, unique_id)


@dbt_mcp_tool(
    description=get_prompt("discovery/get_model_health"),
    title="Get Model Health",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_model_health(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
) -> list[dict]:
    return await context.models_fetcher.fetch_model_health(name, unique_id)


@dbt_mcp_tool(
    description=get_prompt("discovery/get_model_performance"),
    title="Get Model Performance",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_model_performance(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
    num_runs: int = Field(
        default=1,
        description="Number of historical runs to return. Default is 1 (latest run only). "
        "Use values > 1 to analyze performance trends over time.",
        ge=1,
        le=100,
    ),
    include_tests: bool = Field(
        default=False,
        description="If True, include test execution history (name, status, executionTime) for each run. "
        "Useful for analyzing test performance alongside model execution. "
        "Default is False to reduce response size.",
    ),
) -> list[dict]:
    """Get model execution performance metrics from historical runs."""
    results = await context.model_performance_fetcher.fetch_performance(
        name=name,
        unique_id=unique_id,
        num_runs=num_runs,
        include_tests=include_tests,
    )

    return results


@dbt_mcp_tool(
    description=get_prompt("discovery/get_lineage"),
    title="Get Lineage",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_lineage(
    context: DiscoveryToolContext,
    unique_id: str = UNIQUE_ID_REQUIRED_FIELD,
    types: list[LineageResourceType] | None = TYPES_FIELD,
    depth: int = DEPTH_FIELD,
) -> list[dict]:
    return await context.lineage_fetcher.fetch_lineage(
        unique_id=unique_id, types=types, depth=depth
    )


@dbt_mcp_tool(
    description=get_prompt("discovery/get_exposures"),
    title="Get Exposures",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_exposures(context: DiscoveryToolContext) -> list[dict]:
    return await context.exposures_fetcher.fetch_exposures()


@dbt_mcp_tool(
    description=get_prompt("discovery/get_exposure_details"),
    title="Get Exposure Details",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_exposure_details(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
) -> list[dict]:
    return await context.resource_details_fetcher.fetch_details(
        resource_type=AppliedResourceType.EXPOSURE,
        unique_id=unique_id,
        name=name,
    )


@dbt_mcp_tool(
    description=get_prompt("discovery/get_all_sources"),
    title="Get All Sources",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_all_sources(
    context: DiscoveryToolContext,
    source_names: list[str] | None = None,
    unique_ids: list[str] | None = None,
) -> list[dict]:
    return await context.sources_fetcher.fetch_sources(source_names, unique_ids)


@dbt_mcp_tool(
    description=get_prompt("discovery/get_source_details"),
    title="Get Source Details",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_source_details(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
) -> list[dict]:
    return await context.resource_details_fetcher.fetch_details(
        resource_type=AppliedResourceType.SOURCE,
        unique_id=unique_id,
        name=name,
    )


@dbt_mcp_tool(
    description=get_prompt("discovery/get_all_macros"),
    title="Get All Macros",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_all_macros(
    context: DiscoveryToolContext,
    package_names: list[str] | None = Field(
        default=None,
        description="Optional list of package names to filter macros by "
        "(e.g., ['my_project', 'my_package']).",
    ),
    return_package_names_only: bool = Field(
        default=False,
        description="If True, returns only the unique package names instead of "
        "full macro details. Use this to discover available packages first, "
        "then filter by specific packages to get macro details.",
    ),
    include_default_dbt_packages: bool = Field(
        default=False,
        description="If True, includes the default dbt macros that "
        "are maintained by dbt Labs.",
    ),
) -> list[dict] | list[str]:
    return await context.macros_fetcher.fetch_macros(
        package_names=package_names,
        return_package_names_only=return_package_names_only,
        include_default_dbt_packages=include_default_dbt_packages,
    )


@dbt_mcp_tool(
    description=get_prompt("discovery/get_macro_details"),
    title="Get Macro Details",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_macro_details(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
) -> list[dict]:
    return await context.resource_details_fetcher.fetch_details(
        resource_type=AppliedResourceType.MACRO,
        unique_id=unique_id,
        name=name,
    )


@dbt_mcp_tool(
    description=get_prompt("discovery/get_seed_details"),
    title="Get Seed Details",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_seed_details(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
) -> list[dict]:
    return await context.resource_details_fetcher.fetch_details(
        resource_type=AppliedResourceType.SEED,
        unique_id=unique_id,
        name=name,
    )


@dbt_mcp_tool(
    description=get_prompt("discovery/get_semantic_model_details"),
    title="Get Semantic Model Details",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_semantic_model_details(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
) -> list[dict]:
    return await context.resource_details_fetcher.fetch_details(
        resource_type=AppliedResourceType.SEMANTIC_MODEL,
        unique_id=unique_id,
        name=name,
    )


@dbt_mcp_tool(
    description=get_prompt("discovery/get_snapshot_details"),
    title="Get Snapshot Details",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_snapshot_details(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
) -> list[dict]:
    return await context.resource_details_fetcher.fetch_details(
        resource_type=AppliedResourceType.SNAPSHOT,
        unique_id=unique_id,
        name=name,
    )


@dbt_mcp_tool(
    description=get_prompt("discovery/get_test_details"),
    title="Get Test Details",
    read_only_hint=True,
    destructive_hint=False,
    idempotent_hint=True,
)
async def get_test_details(
    context: DiscoveryToolContext,
    name: str | None = NAME_FIELD,
    unique_id: str | None = UNIQUE_ID_FIELD,
) -> list[dict]:
    return await context.resource_details_fetcher.fetch_details(
        resource_type=AppliedResourceType.TEST,
        unique_id=unique_id,
        name=name,
    )


DISCOVERY_TOOLS = [
    get_mart_models,
    get_all_models,
    get_model_details,
    get_model_parents,
    get_model_children,
    get_model_health,
    get_model_performance,
    get_lineage,
    get_exposures,
    get_exposure_details,
    get_all_sources,
    get_source_details,
    get_all_macros,
    get_macro_details,
    get_seed_details,
    get_semantic_model_details,
    get_snapshot_details,
    get_test_details,
]


def register_discovery_tools(
    dbt_mcp: FastMCP,
    discovery_config_provider: ConfigProvider[DiscoveryConfig],
    *,
    disabled_tools: set[ToolName],
    enabled_tools: set[ToolName] | None,
    enabled_toolsets: set[Toolset],
    disabled_toolsets: set[Toolset],
) -> None:
    def bind_context() -> DiscoveryToolContext:
        return DiscoveryToolContext(config_provider=discovery_config_provider)

    register_tools(
        dbt_mcp,
        tool_definitions=[tool.adapt_context(bind_context) for tool in DISCOVERY_TOOLS],
        disabled_tools=disabled_tools,
        enabled_tools=enabled_tools,
        enabled_toolsets=enabled_toolsets,
        disabled_toolsets=disabled_toolsets,
    )