from __future__ import annotations

import json
import logging
from typing import Any

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

from cog_mcp.config import Config
from cog_mcp.formatting import (
    format_aggregation,
    format_instances,
    format_query_results,
)
from cog_mcp.resources import FILTER_DOCS, QUERY_DOCS, fetch_view_schema, fetch_views

logger = logging.getLogger(__name__)


def _parse_json_param(value: str | dict | list | None) -> Any | None:
    """Parse a JSON parameter that may arrive as a string or already-parsed object.

    Some MCP clients auto-parse JSON strings into dicts/lists before passing
    them to tool functions. This helper handles both cases.
    """
    if value is None:
        return None
    if isinstance(value, str):
        return json.loads(value)
    return value


def register_tools(mcp: FastMCP, client: CogniteClient, config: Config) -> None:
    """Register all MCP tools on the FastMCP server."""

    # ── Discovery tools ──────────────────────────────────────────────

    @mcp.tool()
    def list_views() -> str:
        """List all available views with their space, externalId, and version.

        Call this first to discover valid view coordinates before using
        list_instances, search_instances, or other tools.
        """
        return fetch_views(client, config)

    @mcp.tool()
    def get_view_schema(
        view_space: str,
        view_external_id: str,
        view_version: str,
    ) -> str:
        """Get the full schema of a view including all property names, types, and relations.

        Call this before filtering, sorting, or aggregating to get correct property
        names and relation targets. Get the space, externalId, and version from list_views.
        """
        return fetch_view_schema(client, view_space, view_external_id, view_version)

    @mcp.tool()
    def get_filter_docs() -> str:
        """Get reference documentation for CDF instance filter syntax.

        Call this when you need to construct filter expressions for
        list_instances, search_instances, or aggregate_instances.
        """
        return FILTER_DOCS

    @mcp.tool()
    def get_query_docs() -> str:
        """Get reference documentation for the CDF graph query API.

        Call this when you need to construct queries for the
        query_instances tool, including traversals and multi-result-set queries.
        """
        return QUERY_DOCS

    # ── Instance tools ───────────────────────────────────────────────

    @mcp.tool()
    def list_instances(
        view_space: str,
        view_external_id: str,
        view_version: str,
        filter: str | dict | None = None,
        limit: int = 100,
        sort: str | list | dict | None = None,
        instance_type: str = "node",
    ) -> str:
        """List instances of a view with optional filtering.

        Use the list_views tool to find valid view coordinates. For filter/sort property
        names, call get_view_schema.

        Instance space filters are automatically applied from configuration.
        The SDK handles pagination internally — use `limit` to control how many
        instances are returned. Use aggregate_instances first to check total counts.

        For filter syntax, call the get_filter_docs tool.

        Args:
            view_space: Space of the view to list instances from.
            view_external_id: External ID of the view.
            view_version: Version of the view.
            filter: Optional JSON filter expression (e.g. {"equals": {"property": ["viewSpace", "viewId/version", "prop"], "value": "x"}}).
            limit: Max instances to return. Default 100. Use -1 for ALL instances (auto-paginates).
            sort: Optional JSON sort spec (e.g. [{"property": ["space", "viewId/version", "prop"], "direction": "ascending"}]).
            instance_type: "node" or "edge". Default "node".
        """
        view_id = ViewId(space=view_space, external_id=view_external_id, version=view_version)

        # Parse filter
        parsed_filter = _parse_json_param(filter)

        # Parse sort
        parsed_sort = None
        if sort:
            raw_sort = _parse_json_param(sort)
            if isinstance(raw_sort, list):
                parsed_sort = [
                    InstanceSort(
                        property=s["property"],
                        direction=s.get("direction", "ascending"),
                        nulls_first=s.get("nullsFirst"),
                    )
                    for s in raw_sort
                ]
            else:
                parsed_sort = InstanceSort(
                    property=raw_sort["property"],
                    direction=raw_sort.get("direction", "ascending"),
                )

        # Resolve limit: -1 means all
        effective_limit = None if limit == -1 else limit

        result = client.data_modeling.instances.list(
            instance_type=instance_type,
            sources=[view_id],
            space=config.instance_spaces,
            filter=parsed_filter,
            limit=effective_limit,
            sort=parsed_sort,
        )

        # Get total count for messaging
        total = None
        if effective_limit is not None and len(result) >= effective_limit:
            try:
                agg = client.data_modeling.instances.aggregate(
                    view=view_id,
                    aggregates=aggregations.Count("externalId"),
                    space=config.instance_spaces,
                    filter=parsed_filter,
                )
                if hasattr(agg, "value"):
                    total = int(agg.value)
                elif isinstance(agg, list) and agg:
                    total = int(agg[0].value)
            except Exception:
                pass

        return format_instances(list(result), total=total)

    @mcp.tool()
    def search_instances(
        view_space: str,
        view_external_id: str,
        view_version: str,
        query: str,
        properties: str | list | None = None,
        filter: str | dict | None = None,
        limit: int = 100,
        instance_type: str = "node",
        operator: str | None = None,
    ) -> str:
        """Full-text search across instances of a view.

        Use the list_views tool to find valid view coordinates. For property names to
        search or filter on, call get_view_schema.

        Instance space filters are automatically applied from configuration.
        Note: Search is eventually consistent — new/updated data may take a few seconds to appear.

        For filter syntax, call the get_filter_docs tool.

        Args:
            view_space: Space of the view to search in.
            view_external_id: External ID of the view.
            view_version: Version of the view.
            query: Search text. Tokens are matched case-insensitively. Last token uses prefix matching.
            properties: Optional JSON list of property names to search (e.g. ["name", "description"]). If not specified, all text fields are searched.
            filter: Optional JSON filter to narrow results (e.g. {"range": {"property": [...], "gt": 2015}}).
            limit: Max results. Default 100. Use -1 for max (1000).
            instance_type: "node" or "edge". Default "node".
            operator: "AND" (all terms must match) or "OR" (any term matches). Default: API default.
        """
        view_id = ViewId(space=view_space, external_id=view_external_id, version=view_version)

        parsed_filter = _parse_json_param(filter)
        parsed_properties = _parse_json_param(properties)
        effective_limit = None if limit == -1 else limit

        result = client.data_modeling.instances.search(
            view=view_id,
            query=query,
            instance_type=instance_type,
            properties=parsed_properties,
            space=config.instance_spaces,
            filter=parsed_filter,
            limit=effective_limit,
            operator=operator,
        )

        return format_instances(list(result))

    @mcp.tool()
    def aggregate_instances(
        view_space: str,
        view_external_id: str,
        view_version: str,
        aggregate: str,
        property: str | None = None,
        group_by: str | None = None,
        filter: str | dict | None = None,
        query: str | None = None,
        instance_type: str = "node",
        histogram_interval: float | None = None,
    ) -> str:
        """Aggregate instances of a view (count, sum, avg, min, max, histogram).

        Use the list_views tool to find valid view coordinates. Call get_view_schema for
        valid property and group_by names.

        Instance space filters are automatically applied. Use this to understand data volumes
        before listing — especially important for large datasets.

        Tip: To get the total number of instances, use aggregate="count" with property="externalId".

        For filter syntax, call the get_filter_docs tool.

        Args:
            view_space: Space of the view.
            view_external_id: External ID of the view.
            view_version: Version of the view.
            aggregate: One of "count", "sum", "avg", "min", "max", "histogram".
            property: Property to aggregate. Required for all aggregates. Use "externalId" for total count.
            group_by: Optional property name to group by (e.g. "status"). For multiple, use comma-separated (e.g. "status,type").
            filter: Optional JSON filter expression.
            query: Optional search query to filter before aggregating.
            instance_type: "node" or "edge". Default "node".
            histogram_interval: Bucket interval for histogram aggregation. Required if aggregate is "histogram".
        """
        view_id = ViewId(space=view_space, external_id=view_external_id, version=view_version)

        # Build aggregation
        agg_map = {
            "count": aggregations.Count,
            "sum": aggregations.Sum,
            "avg": aggregations.Avg,
            "min": aggregations.Min,
            "max": aggregations.Max,
        }

        if aggregate == "histogram":
            if not property or histogram_interval is None:
                return json.dumps({"error": "histogram requires 'property' and 'histogram_interval'"})
            agg = aggregations.Histogram(property=property, interval=histogram_interval)
        elif aggregate in agg_map:
            if not property:
                return json.dumps({"error": f"'{aggregate}' requires 'property'"})
            agg = agg_map[aggregate](property=property)
        else:
            return json.dumps({"error": f"Unknown aggregate: {aggregate}. Use one of: count, sum, avg, min, max, histogram"})

        parsed_filter = _parse_json_param(filter)

        # group_by: accept a property name or comma-separated names
        parsed_group_by: str | list[str] | None = None
        if group_by:
            parts = [p.strip() for p in group_by.split(",")]
            parsed_group_by = parts if len(parts) > 1 else parts[0]

        result = client.data_modeling.instances.aggregate(
            view=view_id,
            aggregates=agg,
            group_by=parsed_group_by,
            instance_type=instance_type,
            query=query,
            space=config.instance_spaces,
            filter=parsed_filter,
        )

        # If group_by returned no results, provide a diagnostic hint
        if group_by and not result:
            try:
                check = client.data_modeling.instances.aggregate(
                    view=view_id,
                    aggregates=aggregations.Count(property=property or "externalId"),
                    instance_type=instance_type,
                    space=config.instance_spaces,
                    filter=parsed_filter,
                )
                total = check[0].value if check else 0
                hint = (
                    f"group_by '{group_by}' returned no results. "
                    f"A count on '{property or 'externalId'}' (without group_by) shows {total} values. "
                )
                if total == 0:
                    hint += "This is because no instances match the provided filter."
                else:
                    hint += (
                        f"This suggests that while {total} instances match your filter, "
                        f"none of them have a populated value for the '{group_by}' property."
                    )
                return format_aggregation(result, hint=hint)
            except Exception as e:
                logger.warning(f"Failed to get diagnostic count for empty group_by: {e}")

        return format_aggregation(result)

    @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)

        # Build with_ dict
        with_: dict[str, NodeResultSetExpression | EdgeResultSetExpression] = {}
        for name, expr_def in raw.get("with", {}).items():
            if "nodes" in expr_def:
                node_def = expr_def["nodes"]
                node_filter = node_def.get("filter")

                # Inject space filter
                space_filter = filters.In(["node", "space"], config.instance_spaces)
                if node_filter:
                    combined = filters.And(space_filter, filters.Filter.load(node_filter))
                else:
                    combined = space_filter

                with_[name] = NodeResultSetExpression(
                    from_=node_def.get("from"),
                    filter=combined,
                    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"]
                edge_filter = edge_def.get("filter")

                space_filter = filters.In(["edge", "space"], config.instance_spaces)
                if edge_filter:
                    combined = filters.And(space_filter, filters.Filter.load(edge_filter))
                else:
                    combined = space_filter

                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"),
                )

        # Build select dict
        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)

    @mcp.tool()
    def retrieve_instances(
        nodes: str | list | None = None,
        edges: str | list | None = None,
        view_space: str | None = None,
        view_external_id: str | None = None,
        view_version: str | None = None,
    ) -> str:
        """Retrieve specific instances by their IDs.

        Use the list_views tool to find valid view coordinates if scoping by view.

        Args:
            nodes: JSON list of node IDs (e.g. [{"space": "mySpace", "externalId": "node1"}]).
            edges: JSON list of edge IDs (e.g. [{"space": "mySpace", "externalId": "edge1"}]).
            view_space: Optional view space to scope returned properties.
            view_external_id: Optional view external ID.
            view_version: Optional view version.
        """
        parsed_nodes = None
        if nodes:
            raw_nodes = _parse_json_param(nodes)
            parsed_nodes = [
                NodeId(space=n["space"], external_id=n["externalId"])
                for n in raw_nodes
            ]

        parsed_edges = None
        if edges:
            raw_edges = _parse_json_param(edges)
            parsed_edges = [
                EdgeId(space=e["space"], external_id=e["externalId"])
                for e in raw_edges
            ]

        sources = None
        if view_space and view_external_id and view_version:
            sources = [ViewId(space=view_space, external_id=view_external_id, version=view_version)]

        result = client.data_modeling.instances.retrieve(
            nodes=parsed_nodes,
            edges=parsed_edges,
            sources=sources,
        )

        # InstancesResult has .nodes and .edges
        all_items = []
        if result.nodes:
            all_items.extend(result.nodes)
        if result.edges:
            all_items.extend(result.edges)

        return format_instances(all_items)
