"""
MCP tool for querying TopSQL by CPU time.

Ranks SQL statements by CPU consumption over a time range, joining with
topsql_sql_meta to resolve sql_digest → SQL text.

"""

import json
from typing import Any

from pydantic import ValidationError as PydanticValidationError

from src.models.query_topsql import TOPSQL_CPU_TIME_METRIC, TOPSQL_SQL_META_METRIC, TopSQLRequest
from src.services.auth import get_valid_server_auth
from src.services.prometheus import get_prometheus_client_for_cluster
from src.tools.query_range import mcp
from src.utils.errors import (
    AuthorizationError,
    AuthServiceError,
    NetworkError,
    PrometheusAPIError,
    PrometheusTimeoutError,
    ValidationError,
)
from src.utils.logging_utils import get_logger
from src.utils.tool_response import format_tool_error, tool_text_response

logger = get_logger("tool.query_topsql")

TOPSQL_TIMEOUT = 20.0


def _build_topsql_query(request: TopSQLRequest) -> str:
    duration = request.compute_range_duration()
    metric = TOPSQL_CPU_TIME_METRIC
    limit = request.limit

    if request.instance_type:
        cpu_expr = (
            f"topk({limit}, sum by (sql_digest) "
            f'(sum_over_time({metric}{{instance_type="{request.instance_type}"}}[{duration}])))'
        )
    else:
        cpu_expr = f"topk({limit}, sum by (sql_digest) (sum_over_time({metric}[{duration}])))"

    meta_expr = f"max by (sql_digest, normalized_sql) ({TOPSQL_SQL_META_METRIC})"
    return f"{cpu_expr} * on(sql_digest) group_left(normalized_sql) {meta_expr}"


def _extract_topsql_entries(
    result_data: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    raw: list[tuple[str, str | None, float]] = []
    for item in result_data:
        metric = item.get("metric", {})
        sql_digest = metric.get("sql_digest", "")
        if not sql_digest:
            continue
        normalized_sql = metric.get("normalized_sql") or None
        value = item.get("value")
        cpu_time = 0.0
        if isinstance(value, (list, tuple)) and len(value) >= 2:
            try:
                cpu_time = float(value[1])
            except (ValueError, TypeError):
                pass
        raw.append((sql_digest, normalized_sql, cpu_time))

    raw.sort(key=lambda x: x[2], reverse=True)

    entries = []
    for rank, (sql_digest, sql_text, cpu_time) in enumerate(raw, 1):
        entries.append(
            {
                "rank": rank,
                "sql_digest": sql_digest,
                "sql_text": sql_text,
                "cpu_time_ms": cpu_time,
            }
        )
    return entries


@mcp.tool()
async def query_topsql(
    cluster_id: str,
    env: str,
    deploy_type: str,
    start: str,
    end: str,
    instance_type: str | None = None,
    limit: int = 20,
) -> list[dict[str, Any]]:
    """
    Rank SQL statements by CPU time using TiDB TopSQL metrics. Always prefer the tool over raw PromQL for common TopSQL queries unless it does not satisfy your needs.

    Follow the WORKFLOW: `list_metrics` -> `list_metric_labels` -> `list_metric_label_values` -> `query_instant` or `query_range` for custom TopSQL queries
    **important**: the caridinality of raw TopSQL metrics can be very high, you MUST add the label filters.

    TopSQL Metric Schema (all metrics share the 'sql_digest' label as join key):

    - topsql_cpu_time_ms: CPU time consumed by each SQL on a TiDB instance (milliseconds).

    - topsql_stmt_exec_count: Total execution count per SQL digest.

    - topsql_stmt_duration_sum_ns: Total wall-clock duration per SQL (nanoseconds, includes network/IO/lock waits).

    - topsql_read_keys: Number of TiKV keys read per SQL digest (high values suggest missing indexes).

    - topsql_sql_meta: SQL metadata dictionary, contains the SQL text in the label value.

    - topsql_cpu_time_ms:sum_over_time_5m:sum_sql_digest & topsql_stmt_exec_count:sum_over_time_5m:sum_sql_digest_instance_type: Pre-aggregated, use when raw metric cardinality is high.
    """
    try:
        logger.info(
            "tool_call_start name=query_topsql cluster_id=%s env=%s deploy_type=%s",
            cluster_id,
            env,
            deploy_type,
        )
        await get_valid_server_auth()
        try:
            request = TopSQLRequest.model_validate(
                {
                    "cluster_id": str(cluster_id).strip(),
                    "env": env.strip(),
                    "deploy_type": deploy_type.strip(),
                    "start": start,
                    "end": end,
                    "instance_type": instance_type,
                    "limit": limit,
                }
            )
        except (ValueError, PydanticValidationError) as e:
            raise ValidationError(f"Invalid TopSQL query parameters: {str(e)}") from e

        client = await get_prometheus_client_for_cluster(
            cluster_id=request.cluster_id,
            env=request.env,
            deploy_type=request.deploy_type,
            timeout=TOPSQL_TIMEOUT,
        )

        query = _build_topsql_query(request)
        response = await client.query_instant(
            query=query,
            time=request.end,
            timeout=TOPSQL_TIMEOUT,
        )
        results = response.get("data", {}).get("result", [])
        entries = _extract_topsql_entries(results)

        result = {"status": "success", "data": entries}
        logger.info(
            "tool_call_success name=query_topsql cluster_id=%s entries=%s",
            request.cluster_id,
            len(entries),
        )
        return tool_text_response(json.dumps(result, ensure_ascii=False))

    except (
        AuthServiceError,
        AuthorizationError,
        ValidationError,
        NetworkError,
        PrometheusTimeoutError,
        PrometheusAPIError,
    ) as e:
        logger.warning("tool_call_error name=query_topsql type=%s", e.__class__.__name__)
        return tool_text_response(format_tool_error(e))
    except Exception as e:
        logger.exception("tool_call_unexpected_error name=query_topsql")
        return tool_text_response(format_tool_error(e))