"""
Dataset management tools

Tools for listing and inspecting 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, DetailLevel
from ..formatters import (
    format_datasets_markdown,
    format_dataset_markdown,
    format_refresh_history_markdown,
    format_parameters_markdown,
    truncate_response
)
from ..validation import validate_refresh_top
from ..constants import REFRESH_HISTORY_DEFAULT

logger = logging.getLogger(__name__)


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

    @mcp.tool()
    async def get_datasets(
        workspace_id: Optional[str] = None,
        format: str = "markdown",
        detail: str = "concise",
        ctx: Context = None
    ) -> str:
        """
        Get list of datasets from a specific workspace or My workspace.

        Datasets in PowerBI contain the data model, including tables, columns, relationships,
        and measures. This tool retrieves all datasets accessible in the specified workspace.

        Use this when you need to:
        - List all datasets in a workspace
        - Find a dataset ID for querying
        - Discover available data sources
        - Check dataset refresh status and configuration

        Parameters:
        - workspace_id (optional): Workspace (group) ID. If not provided, returns datasets from "My workspace"
        - format: Response format - "json" or "markdown" (default: "markdown")
        - detail: Detail level - "concise" or "detailed" (default: "concise")

        Returns:
        List of datasets with their IDs, names, and optionally detailed metadata including
        refresh status, storage mode, and configuration details.

        Example usage:
        - Get datasets from specific workspace: workspace_id="abc123..."
        - Get datasets from My workspace: (omit workspace_id)
        - Get detailed info: detail="detailed"

        Error handling:
        - If workspace_id not found, verify the ID is correct using get_workspaces
        - For permission errors, ensure service principal has read access to the workspace
        - Empty list means no datasets in the workspace or no access
        """
        try:
            format_type = ResponseFormat(format.lower())
            detail_level = DetailLevel(detail.lower())

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

            if workspace_id:
                await ctx.info(f"Fetching datasets from workspace {workspace_id}")
            else:
                await ctx.info("Fetching datasets from My workspace")

            # Fetch datasets from API
            result = await client.get_datasets(workspace_id=workspace_id)
            datasets = result.get("value", [])

            await ctx.info(f"Found {len(datasets)} datasets")

            # Format response
            if format_type == ResponseFormat.MARKDOWN:
                response = format_datasets_markdown(datasets, detail_level)
            else:
                response = json.dumps(datasets, 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 get datasets: {error_msg}")
            return f"Error executing get_datasets: {error_msg}"

    @mcp.tool()
    async def get_dataset(
        dataset_id: str,
        workspace_id: Optional[str] = None,
        format: str = "json",
        ctx: Context = None
    ) -> str:
        """
        Get detailed information about a specific dataset.

        Retrieves comprehensive metadata about a PowerBI dataset including configuration,
        refresh settings, and data source information.

        Use this when you need to:
        - Get detailed metadata about a specific dataset
        - Check dataset configuration and capabilities
        - Verify dataset refresh settings
        - Understand dataset storage mode and requirements

        Parameters:
        - dataset_id (required): The unique identifier of the dataset
        - workspace_id (optional): Workspace (group) ID. Omit for datasets in "My workspace"
        - format: Response format - "json" or "markdown" (default: "json")

        Returns:
        Detailed dataset information including:
        - Dataset name and ID
        - Configuration details
        - Refresh capabilities and requirements
        - Storage mode
        - Creation date
        - Identity requirements

        Example usage:
        - Get dataset from My workspace: dataset_id="dataset123"
        - Get dataset from specific workspace: dataset_id="dataset123", workspace_id="workspace456"

        Error handling:
        - If dataset_id not found (404), verify the ID using get_datasets
        - For permission errors, ensure service principal has read access
        - Check workspace_id matches the workspace containing the dataset
        """
        try:
            if not dataset_id:
                return "Error: dataset_id is required"

            format_type = ResponseFormat(format.lower())

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

            if workspace_id:
                await ctx.info(f"Fetching dataset {dataset_id} from workspace {workspace_id}")
            else:
                await ctx.info(f"Fetching dataset {dataset_id} from My workspace")

            # Fetch dataset from API
            result = await client.get_dataset(
                dataset_id=dataset_id,
                workspace_id=workspace_id
            )

            await ctx.info(f"Retrieved dataset: {result.get('name', 'Unknown')}")

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

            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 get dataset: {error_msg}")
            return f"Error executing get_dataset: {error_msg}"

    @mcp.tool()
    async def get_refresh_history(
        dataset_id: str,
        workspace_id: Optional[str] = None,
        top: int = REFRESH_HISTORY_DEFAULT,
        format: str = "markdown",
        ctx: Context = None
    ) -> str:
        """
        Get refresh history for a PowerBI dataset.

        Shows recent refresh operations including:
        - Status (Completed, Failed, Unknown)
        - Start and end times
        - Request ID and refresh type
        - Error messages (for failed refreshes)

        Useful for monitoring data freshness and troubleshooting refresh failures.

        Parameters:
        - dataset_id (required): The dataset ID
        - workspace_id (optional): Workspace ID. Omit for datasets in "My workspace"
        - top: Number of refresh records to return (default: 5, max: 60)
        - format: Response format - "markdown" or "json" (default: "markdown")

        Returns:
        Formatted refresh history with status, timestamps, and error details for failed refreshes.

        Example usage:
        - Get last 5 refreshes: dataset_id="abc123"
        - Get last 10 refreshes: dataset_id="abc123", top=10
        - Check specific workspace: dataset_id="abc123", workspace_id="workspace456"

        Error handling:
        - If dataset_id not found, verify the ID using get_datasets
        - For permission errors, ensure service principal has read access
        """
        try:
            if not dataset_id:
                return "Error: dataset_id is required"

            # Validate top parameter
            validate_refresh_top(top)
            format_type = ResponseFormat(format.lower())

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

            if workspace_id:
                await ctx.info(f"Fetching refresh history for dataset {dataset_id} from workspace {workspace_id}")
            else:
                await ctx.info(f"Fetching refresh history for dataset {dataset_id} from My workspace")

            # Fetch refresh history from API
            result = await client.get_refresh_history(
                dataset_id=dataset_id,
                workspace_id=workspace_id,
                top=top
            )
            refresh_history = result.get("value", [])

            await ctx.info(f"Retrieved {len(refresh_history)} refresh records")

            # Get dataset name for formatting
            dataset_info = await client.get_dataset(dataset_id=dataset_id, workspace_id=workspace_id)
            dataset_name = dataset_info.get("name", dataset_id)

            # Format response
            if format_type == ResponseFormat.MARKDOWN:
                response = format_refresh_history_markdown(refresh_history, dataset_name)
            else:
                response = json.dumps(refresh_history, 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 get refresh history: {error_msg}")
            return f"Error executing get_refresh_history: {error_msg}"

    @mcp.tool()
    async def get_parameters(
        dataset_id: str,
        workspace_id: Optional[str] = None,
        format: str = "markdown",
        detail: str = "concise",
        ctx: Context = None
    ) -> str:
        """
        Get parameters defined in a PowerBI dataset.

        Returns parameter definitions including:
        - Name and data type
        - Current value
        - Whether parameter is required
        - Suggested values (if defined)

        Useful for discovering available parameters before querying parameterized datasets.

        Note: Not supported for datasets with SQL, Oracle, Teradata, SAP HANA DirectQuery
        connections or datasets modified via XMLA endpoint.

        Parameters:
        - dataset_id (required): The dataset ID
        - workspace_id (optional): Workspace ID. Omit for datasets in "My workspace"
        - format: Response format - "markdown" or "json" (default: "markdown")
        - detail: Detail level - "concise" or "detailed" (default: "concise")

        Returns:
        Formatted parameter information with names, types, values, and suggested values.

        Example usage:
        - Get all parameters: dataset_id="abc123"
        - Get detailed info: dataset_id="abc123", detail="detailed"
        - Check specific workspace: dataset_id="abc123", workspace_id="workspace456"

        Error handling:
        - If dataset_id not found, verify the ID using get_datasets
        - "Not supported" errors indicate dataset type limitations
        - Empty result means no parameters are defined
        """
        try:
            if not dataset_id:
                return "Error: dataset_id is required"

            format_type = ResponseFormat(format.lower())
            detail_level = DetailLevel(detail.lower())

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

            if workspace_id:
                await ctx.info(f"Fetching parameters for dataset {dataset_id} from workspace {workspace_id}")
            else:
                await ctx.info(f"Fetching parameters for dataset {dataset_id} from My workspace")

            # Fetch parameters from API
            result = await client.get_parameters(
                dataset_id=dataset_id,
                workspace_id=workspace_id
            )
            parameters = result.get("value", [])

            await ctx.info(f"Retrieved {len(parameters)} parameters")

            # Get dataset name for formatting
            dataset_info = await client.get_dataset(dataset_id=dataset_id, workspace_id=workspace_id)
            dataset_name = dataset_info.get("name", dataset_id)

            # Format response
            if format_type == ResponseFormat.MARKDOWN:
                response = format_parameters_markdown(parameters, dataset_name, detail_level)
            else:
                response = json.dumps(parameters, 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 get parameters: {error_msg}")
            return f"Error executing get_parameters: {error_msg}"

    logger.info("Dataset tools registered")
