from cognite.client import CogniteClient
from cognite.client.data_classes.data_modeling import ViewId
from cognite.client.data_classes.data_modeling.instances import InstanceSort
from cognite.client.data_classes.data_modeling.query import (
    EdgeResultSetExpression,
    NodeResultSetExpression,
    Query,
    Select,
    SourceSelector,
)
from mcp.server.fastmcp import FastMCP

from cog_mcp.config import Config
from cog_mcp.formatting import format_query_results
from cog_mcp.parsing import inject_space_filter, parse_json_param, parse_through


def register_query_tools(mcp: FastMCP, client: CogniteClient, config: Config) -> None:
    """Register the graph query tool."""

    @mcp.tool()
    def query_instances(
        query: str | dict,
    ) -> str:
        """Execute a graph query across related instances using the DM query API.

        Use list_views and get_view_schema tools for view coordinates, property names,
        and relation identifiers.

        This is the most powerful query tool — it supports traversing relationships,
        joining data across views, and complex filtering. Instance space filters are
        automatically injected into all result set expressions.

        For full query syntax and examples, call the get_query_docs tool.

        Args:
            query: JSON query definition with "with" (result set expressions) and
                "select" (properties to return). Example — list 10 maintenance orders:
                {"with": {"orders": {"nodes": {"filter": {"hasData": [{"space": "cdf_idm", "externalId": "CogniteMaintenanceOrder", "version": "v1", "type": "view"}]}, "limit": 10}}}, "select": {"orders": {"sources": [{"source": {"space": "cdf_idm", "externalId": "CogniteMaintenanceOrder", "version": "v1"}, "properties": ["name", "status", "priority"]}]}}}
        """
        raw = parse_json_param(query)

        with_: dict[str, NodeResultSetExpression | EdgeResultSetExpression] = {}
        for name, expr_def in raw.get("with", {}).items():
            if "nodes" in expr_def:
                node_def = expr_def["nodes"]
                combined = inject_space_filter(
                    "node", node_def.get("filter"), config.instance_spaces
                )

                with_[name] = NodeResultSetExpression(
                    from_=node_def.get("from"),
                    filter=combined,
                    through=parse_through(node_def.get("through")),
                    direction=node_def.get("direction", "outwards"),
                    chain_to=node_def.get("chainTo", "destination"),
                    limit=node_def.get("limit"),
                    sort=[
                        InstanceSort(
                            property=s["property"],
                            direction=s.get("direction", "ascending"),
                        )
                        for s in node_def.get("sort", [])
                    ]
                    if node_def.get("sort")
                    else None,
                )
            elif "edges" in expr_def:
                edge_def = expr_def["edges"]
                combined = inject_space_filter(
                    "edge", edge_def.get("filter"), config.instance_spaces
                )

                with_[name] = EdgeResultSetExpression(
                    from_=edge_def.get("from"),
                    filter=combined,
                    direction=edge_def.get("direction", "outwards"),
                    chain_to=edge_def.get("chainTo", "destination"),
                    limit=edge_def.get("limit"),
                    max_distance=edge_def.get("maxDistance"),
                )

        select: dict[str, Select] = {}
        for name, sel_def in raw.get("select", {}).items():
            sources = []
            for src in sel_def.get("sources", []):
                src_ref = src["source"]
                view_id = ViewId(
                    space=src_ref["space"],
                    external_id=src_ref["externalId"],
                    version=src_ref.get("version"),
                )
                sources.append(SourceSelector(source=view_id, properties=src.get("properties")))
            select[name] = Select(sources=sources, limit=sel_def.get("limit"))

        q = Query(
            with_=with_,
            select=select,
            parameters=raw.get("parameters"),
            cursors=raw.get("cursors"),
        )

        result = client.data_modeling.instances.query(q)
        return format_query_results(result)
