"""
Dataset management tools for FiftyOne MCP server.

| Copyright 2017-2026, Voxel51, Inc.
| `voxel51.com <https://voxel51.com/>`_
|
"""

import json
import logging

import fiftyone as fo
from mcp.types import Tool, TextContent

from .utils import format_response, safe_serialize, dataset_to_summary


logger = logging.getLogger(__name__)


def list_datasets():
    """Lists all available FiftyOne datasets.

    Returns:
        a dict containing list of dataset names and metadata
    """
    try:
        datasets = fo.list_datasets()
        dataset_info = []

        for name in datasets:
            try:
                dataset = fo.load_dataset(name)
                dataset_info.append(
                    {
                        "name": name,
                        "media_type": dataset.media_type,
                        "num_samples": len(dataset),
                        "persistent": dataset.persistent,
                        "tags": dataset.tags,
                    }
                )
            except Exception as e:
                logger.warning("Could not load dataset '%s': %s", name, e)
                dataset_info.append({"name": name, "error": str(e)})

        return format_response(
            {"count": len(datasets), "datasets": dataset_info}
        )

    except Exception as e:
        logger.error("Failed to list datasets: %s", e)
        return format_response(None, success=False, error=str(e))


def load_dataset(name):
    """Loads a FiftyOne dataset by name and returns basic info.

    Args:
        name: the name of the dataset to load

    Returns:
        a dict containing dataset information
    """
    try:
        dataset = fo.load_dataset(name)

        info = {
            "name": dataset.name,
            "media_type": dataset.media_type,
            "num_samples": len(dataset),
            "persistent": dataset.persistent,
            "tags": dataset.tags,
            "info": safe_serialize(dataset.info),
            "fields": list(dataset.get_field_schema().keys()),
        }

        return format_response(info)

    except Exception as e:
        logger.error("Failed to load dataset '%s': %s", name, e)
        return format_response(None, success=False, error=str(e))


def dataset_summary(name):
    """Gets detailed summary statistics for a dataset.

    Args:
        name: the name of the dataset

    Returns:
        a dict containing detailed dataset statistics
    """
    try:
        dataset = fo.load_dataset(name)
        summary = dataset_to_summary(dataset)

        summary["stats"] = {"total_samples": len(dataset), "tags": {}}

        for tag in dataset.tags:
            tagged_view = dataset.match_tags(tag)
            summary["stats"]["tags"][tag] = len(tagged_view)

        schema = dataset.get_field_schema()
        summary["value_counts"] = {}

        for field_name in schema.keys():
            if field_name in ["id", "filepath", "metadata"]:
                continue

            try:
                if hasattr(dataset, "count_values"):
                    counts = dataset.count_values(field_name)
                    if counts and len(counts) < 100:
                        summary["value_counts"][field_name] = dict(counts)
            except Exception:
                pass

        return format_response(summary)

    except Exception as e:
        logger.error("Failed to get summary for dataset '%s': %s", name, e)
        return format_response(None, success=False, error=str(e))


def get_dataset_tools():
    """Gets dataset tool definitions.

    Returns:
        a list of :class:`mcp.types.Tool` instances
    """
    return [
        Tool(
            name="list_datasets",
            description="List all available FiftyOne datasets with metadata",
            inputSchema={"type": "object", "properties": {}, "required": []},
        ),
        Tool(
            name="load_dataset",
            description="Load a FiftyOne dataset by name and return basic information",
            inputSchema={
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string",
                        "description": "Name of the dataset to load",
                    }
                },
                "required": ["name"],
            },
        ),
        Tool(
            name="dataset_summary",
            description="Get detailed summary statistics and metadata for a dataset",
            inputSchema={
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    }
                },
                "required": ["name"],
            },
        ),
    ]


async def handle_tool_call(name, arguments):
    """Handles tool calls for dataset operations.

    Args:
        name: the name of the tool
        arguments: a dict of arguments for the tool

    Returns:
        a list of :class:`mcp.types.TextContent` instances
    """
    try:
        if name == "list_datasets":
            result = list_datasets()
        elif name == "load_dataset":
            dataset_name = arguments.get("name")
            if not dataset_name:
                result = format_response(
                    None, success=False, error="Dataset name is required"
                )
            else:
                result = load_dataset(dataset_name)
        elif name == "dataset_summary":
            dataset_name = arguments.get("name")
            if not dataset_name:
                result = format_response(
                    None, success=False, error="Dataset name is required"
                )
            else:
                result = dataset_summary(dataset_name)
        else:
            result = format_response(
                None, success=False, error="Unknown tool: %s" % name
            )

        return [TextContent(type="text", text=json.dumps(result, indent=2))]

    except Exception as e:
        logger.error("Error handling tool call '%s': %s", name, e)
        error_result = format_response(None, success=False, error=str(e))
        return [
            TextContent(type="text", text=json.dumps(error_result, indent=2))
        ]
