"""
Query execution tools

Tools for executing DAX queries against PowerBI datasets.
"""

import json
import logging
from typing import Optional

from mcp.server import FastMCP
from mcp.server.fastmcp import Context

from ..models import ResponseFormat
from ..formatters import format_query_results_markdown, truncate_response
from ..validation import validate_dax_query

logger = logging.getLogger(__name__)


def register_query_tools(mcp: FastMCP) -> None:
    """Register query-related tools with the MCP server"""

    @mcp.tool()
    async def query_dataset(
        dataset_id: str,
        dax_query: str,
        workspace_id: Optional[str] = None,
        format: str = "markdown",
        ctx: Context = None
    ) -> str:
        """
        Execute DAX (Data Analysis Expressions) queries against a PowerBI dataset.

        This tool allows you to query dataset tables using DAX, PowerBI's formula language.
        DAX queries can retrieve data, perform calculations, and apply filters.

        Use this when you need to:
        - Retrieve data from dataset tables
        - Perform calculations and aggregations
        - Filter and analyze dataset data
        - Get specific rows or aggregated results

        DAX Query Syntax:
        - Basic table query: EVALUATE TableName
        - Filtered query: EVALUATE FILTER(TableName, TableName[Column] = "Value")
        - Aggregation: EVALUATE SUMMARIZE(TableName, TableName[Column], "Total", SUM(TableName[Amount]))
        - Top N: EVALUATE TOPN(10, TableName, TableName[Column], DESC)

        Parameters:
        - dataset_id (required): The unique identifier of the dataset to query
        - dax_query (required): DAX query to execute (must start with EVALUATE)
        - workspace_id (optional): Workspace (group) ID. Omit for "My workspace"
        - format: Response format - "json" or "markdown" (default: "markdown")

        Returns:
        Query results as a table with rows and columns. Markdown format presents results
        as formatted tables. Results are limited to 100 rows in markdown view.

        Example queries:
        1. Get all rows: "EVALUATE 'Sales'"
        2. Filter data: "EVALUATE FILTER('Sales', 'Sales'[Year] = 2024)"
        3. Aggregate: "EVALUATE SUMMARIZE('Sales', 'Sales'[Category], "Total", SUM('Sales'[Amount]))"
        4. Top 10: "EVALUATE TOPN(10, 'Sales', 'Sales'[Amount], DESC)"

        Error handling:
        - For syntax errors, check your DAX query starts with EVALUATE
        - If table/column not found, use get_dataset to see available schema
        - For large result sets, add filters or use TOPN to limit rows
        - Timeout errors suggest simplifying the query or adding filters
        """
        try:
            if not dataset_id:
                return "Error: dataset_id is required"

            if not dax_query:
                return "Error: dax_query is required"

            # Validate DAX query
            try:
                validate_dax_query(dax_query)
            except ValueError as e:
                return f"Error: {str(e)}"

            format_type = ResponseFormat(format.lower())

            # Get client from context
            client = ctx.request_context.lifespan_context["client"]

            if workspace_id:
                await ctx.info(f"Executing query on dataset {dataset_id} in workspace {workspace_id}")
            else:
                await ctx.info(f"Executing query on dataset {dataset_id} in My workspace")

            await ctx.report_progress(0.3, 1.0, "Executing DAX query...")

            # Execute query
            queries = [{"query": dax_query}]
            result = await client.execute_queries(
                dataset_id=dataset_id,
                queries=queries,
                workspace_id=workspace_id
            )

            await ctx.report_progress(1.0, 1.0, "Query completed")
            await ctx.info("Query executed successfully")

            # Format response
            if format_type == ResponseFormat.MARKDOWN:
                response = format_query_results_markdown(result)
            else:
                response = json.dumps(result, indent=2)

            return truncate_response(response)

        except ValueError as e:
            error_msg = f"Invalid parameter value: {str(e)}"
            await ctx.error(error_msg)
            return f"Error: {error_msg}"
        except Exception as e:
            error_msg = str(e)
            await ctx.error(f"Failed to execute query: {error_msg}")
            return f"Error executing query_dataset: {error_msg}"

    logger.info("Query tools registered")
