"""
Aggregation tools for FiftyOne MCP server.

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

import logging

import fiftyone as fo
from mcp.types import Tool

from .utils import SDK, _get_view, format_response, mcp_tool, safe_serialize


logger = logging.getLogger(__name__)


@mcp_tool(SDK)
def count_values(ctx, dataset_name, field, view_stages=None):
    """Counts the occurrences of each value for a field.

    Args:
        ctx: an optional
            :class:`fiftyone.operators.executor.ExecutionContext`
        dataset_name: the name of the dataset
        field: the field path to count values for
        view_stages (None): an optional list of serialized view stage
            dicts to pre-filter the dataset

    Returns:
        a dict with field value counts
    """
    try:
        dataset = fo.load_dataset(dataset_name)
        view = _get_view(dataset, view_stages)
        counts = view.count_values(field)

        return format_response(
            {
                "field": field,
                "counts": safe_serialize(counts),
                "total": sum(counts.values()) if counts else 0,
                "num_values": len(counts),
            }
        )

    except Exception as e:
        logger.error("Failed to count values for '%s': %s", field, e)
        return format_response(None, success=False, error=str(e))


@mcp_tool(SDK)
def distinct(ctx, dataset_name, field, view_stages=None):
    """Gets the distinct values for a field.

    Args:
        ctx: an optional
            :class:`fiftyone.operators.executor.ExecutionContext`
        dataset_name: the name of the dataset
        field: the field path to get distinct values for
        view_stages (None): an optional list of serialized view stage
            dicts to pre-filter the dataset

    Returns:
        a dict with the list of unique values
    """
    try:
        dataset = fo.load_dataset(dataset_name)
        view = _get_view(dataset, view_stages)
        values = view.distinct(field)

        return format_response(
            {
                "field": field,
                "values": safe_serialize(values),
                "count": len(values),
            }
        )

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


@mcp_tool(SDK)
def compute_bounds(ctx, dataset_name, field, view_stages=None):
    """Gets the min and max values for a numeric field.

    Args:
        ctx: an optional
            :class:`fiftyone.operators.executor.ExecutionContext`
        dataset_name: the name of the dataset
        field: the numeric field path
        view_stages (None): an optional list of serialized view stage
            dicts to pre-filter the dataset

    Returns:
        a dict with min and max values for the field
    """
    try:
        dataset = fo.load_dataset(dataset_name)
        view = _get_view(dataset, view_stages)
        mn, mx = view.bounds(field)

        return format_response({"field": field, "min": mn, "max": mx})

    except Exception as e:
        logger.error("Failed to compute bounds for '%s': %s", field, e)
        return format_response(None, success=False, error=str(e))


@mcp_tool(SDK)
def compute_mean(ctx, dataset_name, field, view_stages=None):
    """Computes the mean value of a numeric field.

    Args:
        ctx: an optional
            :class:`fiftyone.operators.executor.ExecutionContext`
        dataset_name: the name of the dataset
        field: the numeric field path
        view_stages (None): an optional list of serialized view stage
            dicts to pre-filter the dataset

    Returns:
        a dict with the mean value for the field
    """
    try:
        dataset = fo.load_dataset(dataset_name)
        view = _get_view(dataset, view_stages)
        result = view.mean(field)

        return format_response({"field": field, "mean": result})

    except Exception as e:
        logger.error("Failed to compute mean for '%s': %s", field, e)
        return format_response(None, success=False, error=str(e))


@mcp_tool(SDK)
def compute_sum(ctx, dataset_name, field, view_stages=None):
    """Computes the sum of a numeric field across all samples.

    Args:
        ctx: an optional
            :class:`fiftyone.operators.executor.ExecutionContext`
        dataset_name: the name of the dataset
        field: the numeric field path
        view_stages (None): an optional list of serialized view stage
            dicts to pre-filter the dataset

    Returns:
        a dict with the sum value for the field
    """
    try:
        dataset = fo.load_dataset(dataset_name)
        view = _get_view(dataset, view_stages)
        result = view.sum(field)

        return format_response({"field": field, "sum": result})

    except Exception as e:
        logger.error("Failed to compute sum for '%s': %s", field, e)
        return format_response(None, success=False, error=str(e))


@mcp_tool(SDK)
def compute_std(ctx, dataset_name, field, view_stages=None):
    """Computes the standard deviation of a numeric field.

    Args:
        ctx: an optional
            :class:`fiftyone.operators.executor.ExecutionContext`
        dataset_name: the name of the dataset
        field: the numeric field path
        view_stages (None): an optional list of serialized view stage
            dicts to pre-filter the dataset

    Returns:
        a dict with the standard deviation for the field
    """
    try:
        dataset = fo.load_dataset(dataset_name)
        view = _get_view(dataset, view_stages)
        result = view.std(field)

        return format_response({"field": field, "std": result})

    except Exception as e:
        logger.error("Failed to compute std for '%s': %s", field, e)
        return format_response(None, success=False, error=str(e))


@mcp_tool(SDK)
def histogram_values(
    ctx,
    dataset_name,
    field,
    bins=50,
    value_range=None,
    view_stages=None,
):
    """Computes a histogram of values for a field.

    Args:
        ctx: an optional
            :class:`fiftyone.operators.executor.ExecutionContext`
        dataset_name: the name of the dataset
        field: the field path
        bins (50): the number of histogram bins
        value_range (None): an optional ``[min, max]`` range for the
            histogram. If None, the full range of values is used
        view_stages (None): an optional list of serialized view stage
            dicts to pre-filter the dataset

    Returns:
        a dict with counts, edges, and out-of-range count
    """
    try:
        dataset = fo.load_dataset(dataset_name)
        view = _get_view(dataset, view_stages)

        kwargs = {"bins": bins}
        if value_range is not None:
            kwargs["range"] = value_range

        counts, edges, other = view.histogram_values(field, **kwargs)

        return format_response(
            {
                "field": field,
                "counts": counts,
                "edges": edges,
                "other": other,
                "bins": bins,
            }
        )

    except Exception as e:
        logger.error("Failed to compute histogram for '%s': %s", field, e)
        return format_response(None, success=False, error=str(e))


@mcp_tool(SDK)
def get_values(ctx, dataset_name, field, view_stages=None, limit=10000):
    """Gets the values of a field across all samples.

    Args:
        ctx: an optional
            :class:`fiftyone.operators.executor.ExecutionContext`
        dataset_name: the name of the dataset
        field: the field path
        view_stages (None): an optional list of serialized view stage
            dicts to pre-filter the dataset
        limit (10000): maximum number of values to return. Use 0 for
            no limit

    Returns:
        a dict with the list of field values
    """
    try:
        dataset = fo.load_dataset(dataset_name)
        view = _get_view(dataset, view_stages)

        truncated = False
        if limit and len(view) > limit:
            view = view.limit(limit)
            truncated = True

        values = view.values(field)

        return format_response(
            {
                "field": field,
                "values": safe_serialize(values),
                "count": len(values),
                "truncated": truncated,
            }
        )

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


def register_tools(registry):
    """Registers all aggregation tools with the registry.

    Args:
        registry: a :class:`fiftyone_mcp.registry.ToolRegistry`
    """
    registry.register(
        Tool(
            name="count_values",
            description=(
                "Count occurrences of each value for a field "
                "across a dataset. Returns a {value: count} dict. "
                "Useful for understanding label distributions, "
                "tag frequencies, or any categorical field. "
                "Optionally filter with view_stages before counting."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "field": {
                        "type": "string",
                        "description": (
                            "Field path to count values for "
                            "(e.g., 'tags', "
                            "'ground_truth.label')"
                        ),
                    },
                    "view_stages": {
                        "type": "array",
                        "items": {"type": "object"},
                        "description": (
                            "Optional list of serialized view "
                            "stage dicts to pre-filter the dataset"
                        ),
                    },
                },
                "required": ["dataset_name", "field"],
            },
        ),
        count_values,
    )

    registry.register(
        Tool(
            name="distinct",
            description=(
                "Get the list of unique values for a field "
                "across a dataset. Returns an array of distinct "
                "values with a count. Useful for discovering all "
                "classes, categories, or unique identifiers in a "
                "field."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "field": {
                        "type": "string",
                        "description": (
                            "Field path to get distinct values "
                            "for (e.g., 'ground_truth.label')"
                        ),
                    },
                    "view_stages": {
                        "type": "array",
                        "items": {"type": "object"},
                        "description": (
                            "Optional list of serialized view "
                            "stage dicts to pre-filter the dataset"
                        ),
                    },
                },
                "required": ["dataset_name", "field"],
            },
        ),
        distinct,
    )

    registry.register(
        Tool(
            name="bounds",
            description=(
                "Get the min and max values for a numeric field. "
                "Returns {min, max}. Useful for understanding the "
                "range of confidence scores, image dimensions, or "
                "any numeric field."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "field": {
                        "type": "string",
                        "description": (
                            "Numeric field path (e.g., "
                            "'ground_truth.detections"
                            ".confidence')"
                        ),
                    },
                    "view_stages": {
                        "type": "array",
                        "items": {"type": "object"},
                        "description": (
                            "Optional list of serialized view "
                            "stage dicts to pre-filter the dataset"
                        ),
                    },
                },
                "required": ["dataset_name", "field"],
            },
        ),
        compute_bounds,
    )

    registry.register(
        Tool(
            name="mean",
            description=(
                "Compute the mean (average) of a numeric field "
                "across all samples. Returns a single float. "
                "Useful for average confidence, uniqueness "
                "scores, etc."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "field": {
                        "type": "string",
                        "description": "Numeric field path",
                    },
                    "view_stages": {
                        "type": "array",
                        "items": {"type": "object"},
                        "description": (
                            "Optional list of serialized view "
                            "stage dicts to pre-filter the dataset"
                        ),
                    },
                },
                "required": ["dataset_name", "field"],
            },
        ),
        compute_mean,
    )

    registry.register(
        Tool(
            name="sum",
            description=(
                "Compute the sum of a numeric field across all "
                "samples. Returns a single number."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "field": {
                        "type": "string",
                        "description": "Numeric field path",
                    },
                    "view_stages": {
                        "type": "array",
                        "items": {"type": "object"},
                        "description": (
                            "Optional list of serialized view "
                            "stage dicts to pre-filter the dataset"
                        ),
                    },
                },
                "required": ["dataset_name", "field"],
            },
        ),
        compute_sum,
    )

    registry.register(
        Tool(
            name="std",
            description=(
                "Compute the standard deviation of a numeric "
                "field across all samples. Returns a single float."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "field": {
                        "type": "string",
                        "description": "Numeric field path",
                    },
                    "view_stages": {
                        "type": "array",
                        "items": {"type": "object"},
                        "description": (
                            "Optional list of serialized view "
                            "stage dicts to pre-filter the dataset"
                        ),
                    },
                },
                "required": ["dataset_name", "field"],
            },
        ),
        compute_std,
    )

    registry.register(
        Tool(
            name="histogram_values",
            description=(
                "Compute a histogram of values for a numeric "
                "field. Returns bin counts, bin edges, and the "
                "count of values outside the range. Useful for "
                "visualizing distributions of confidence scores, "
                "object sizes, etc."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "field": {
                        "type": "string",
                        "description": "Numeric field path",
                    },
                    "bins": {
                        "type": "integer",
                        "description": (
                            "Number of histogram bins. " "Default is 50"
                        ),
                        "default": 50,
                    },
                    "value_range": {
                        "type": "array",
                        "items": {"type": "number"},
                        "minItems": 2,
                        "maxItems": 2,
                        "description": (
                            "Optional [min, max] range for the "
                            "histogram. Values outside this range "
                            "are counted in 'other'. If omitted, "
                            "full range is used"
                        ),
                    },
                    "view_stages": {
                        "type": "array",
                        "items": {"type": "object"},
                        "description": (
                            "Optional list of serialized view "
                            "stage dicts to pre-filter the dataset"
                        ),
                    },
                },
                "required": ["dataset_name", "field"],
            },
        ),
        histogram_values,
    )

    registry.register(
        Tool(
            name="get_values",
            description=(
                "Get the raw values of a field for all samples "
                "in a dataset. Returns a list of values in sample "
                "order. Useful for reading custom field data, "
                "scores, or metadata. Capped at 10,000 samples "
                "by default to avoid large responses."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "field": {
                        "type": "string",
                        "description": (
                            "Field path to read "
                            "(e.g., 'uniqueness', 'my_score')"
                        ),
                    },
                    "view_stages": {
                        "type": "array",
                        "items": {"type": "object"},
                        "description": (
                            "Optional list of serialized view "
                            "stage dicts to pre-filter the dataset"
                        ),
                    },
                    "limit": {
                        "type": "integer",
                        "description": (
                            "Maximum number of values to return. "
                            "Default is 10000. Use 0 for no limit"
                        ),
                        "default": 10000,
                    },
                },
                "required": ["dataset_name", "field"],
            },
        ),
        get_values,
    )
