"""
Sample manipulation 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 _get_view, format_response, safe_serialize


logger = logging.getLogger(__name__)


def add_samples(dataset_name, samples):
    """Adds new samples to a dataset.

    Args:
        dataset_name: the name of the dataset
        samples: a list of dicts, each with a ``filepath`` key and any
            additional field values to set on the sample

    Returns:
        a dict with the IDs of the added samples
    """
    try:
        dataset = fo.load_dataset(dataset_name)

        fo_samples = []
        for d in samples:
            d = dict(d)
            filepath = d.pop("filepath", None)
            if not filepath:
                return format_response(
                    None,
                    success=False,
                    error=("Each sample dict must contain a 'filepath' key"),
                )
            fo_samples.append(fo.Sample(filepath=filepath, **d))

        sample_ids = dataset.add_samples(fo_samples)

        return format_response(
            {
                "dataset_name": dataset_name,
                "added_count": len(sample_ids),
                "sample_ids": [str(sid) for sid in sample_ids],
            }
        )

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


def set_values(dataset_name, field, values, key_field=None):
    """Sets the values of a field across samples.

    Accepts either a list of values (one per sample in default order)
    or a dict mapping sample IDs to values for explicit assignment.

    Args:
        dataset_name: the name of the dataset
        field: the field path to set values for
        values: a list of values (one per sample) or a
            ``{sample_id: value}`` dict for explicit ID-based assignment
        key_field (None): when ``values`` is a list, the sample field
            used as the key. Defaults to ``"id"``

    Returns:
        a dict with the count of updated samples
    """
    try:
        dataset = fo.load_dataset(dataset_name)

        if isinstance(values, dict):
            ids = dataset.values("id")
            id_to_value = {str(k): v for k, v in values.items()}
            ordered_values = [id_to_value.get(str(sid)) for sid in ids]
            dataset.set_values(field, ordered_values)
            updated_count = sum(1 for v in ordered_values if v is not None)
        else:
            kf = key_field or "id"
            dataset.set_values(field, values, key_field=kf)
            updated_count = len(values)

        return format_response(
            {
                "dataset_name": dataset_name,
                "field": field,
                "updated_count": updated_count,
            }
        )

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


def tag_samples(dataset_name, tags, view_stages=None, sample_ids=None):
    """Adds tags to samples in a dataset.

    Args:
        dataset_name: the name of the dataset
        tags: a list of tags to add
        view_stages (None): an optional list of serialized view stage
            dicts to select which samples to tag
        sample_ids (None): an optional list of sample IDs to tag.
            Takes precedence over ``view_stages`` when provided

    Returns:
        a dict with the count of tagged samples
    """
    try:
        dataset = fo.load_dataset(dataset_name)

        if sample_ids is not None:
            view = dataset.select(sample_ids)
        else:
            view = _get_view(dataset, view_stages)

        count = len(view)
        view.tag_samples(tags)

        return format_response(
            {
                "dataset_name": dataset_name,
                "tags": tags,
                "tagged_count": count,
            }
        )

    except Exception as e:
        logger.error("Failed to tag samples in '%s': %s", dataset_name, e)
        return format_response(None, success=False, error=str(e))


def untag_samples(dataset_name, tags, view_stages=None, sample_ids=None):
    """Removes tags from samples in a dataset.

    Args:
        dataset_name: the name of the dataset
        tags: a list of tags to remove
        view_stages (None): an optional list of serialized view stage
            dicts to select which samples to untag
        sample_ids (None): an optional list of sample IDs to untag.
            Takes precedence over ``view_stages`` when provided

    Returns:
        a dict with the count of untagged samples
    """
    try:
        dataset = fo.load_dataset(dataset_name)

        if sample_ids is not None:
            view = dataset.select(sample_ids)
        else:
            view = _get_view(dataset, view_stages)

        count = len(view)
        view.untag_samples(tags)

        return format_response(
            {
                "dataset_name": dataset_name,
                "tags": tags,
                "untagged_count": count,
            }
        )

    except Exception as e:
        logger.error("Failed to untag samples in '%s': %s", dataset_name, e)
        return format_response(None, success=False, error=str(e))


def count_sample_tags(dataset_name):
    """Counts how many samples have each tag.

    Args:
        dataset_name: the name of the dataset

    Returns:
        a dict mapping tags to their sample counts
    """
    try:
        dataset = fo.load_dataset(dataset_name)
        tag_counts = dataset.count_sample_tags()

        return format_response(
            {
                "dataset_name": dataset_name,
                "tags": safe_serialize(tag_counts),
                "num_tags": len(tag_counts),
            }
        )

    except Exception as e:
        logger.error(
            "Failed to count sample tags in '%s': %s", dataset_name, e
        )
        return format_response(None, success=False, error=str(e))


def get_sample_tools():
    """Gets the list of sample manipulation MCP tools.

    Returns:
        a list of :class:`mcp.types.Tool` instances
    """
    return [
        Tool(
            name="add_samples",
            description=(
                "Add new samples to an existing dataset. Each sample "
                "requires a 'filepath' key and can include any additional "
                "field values. Returns the IDs of the added samples. "
                "Use this for programmatic sample creation from Python "
                "dicts — use the import_samples operator for file/dir "
                "imports instead."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "samples": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "filepath": {
                                    "type": "string",
                                    "description": (
                                        "Path to the sample media file"
                                    ),
                                }
                            },
                            "required": ["filepath"],
                        },
                        "description": (
                            "List of sample dicts, each with 'filepath' "
                            "and any extra field values "
                            '(e.g., [{"filepath": "/path/img.jpg", '
                            '"label": "cat"}])'
                        ),
                    },
                },
                "required": ["dataset_name", "samples"],
            },
        ),
        Tool(
            name="set_values",
            description=(
                "Bulk-assign values to a field across multiple samples. "
                "Accepts either a list of values (one per sample in "
                "dataset order) or a {sample_id: value} dict for explicit "
                "assignment. Use this when you need to write computed "
                "scores, labels, or metadata back to the dataset. "
                "Note: the edit_field_values operator does value remapping "
                "(old→new), not bulk assignment."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "field": {
                        "type": "string",
                        "description": (
                            "Field path to write values to "
                            "(e.g., 'my_score', 'predicted_label')"
                        ),
                    },
                    "values": {
                        "description": (
                            "Either a list of values (one per sample) or "
                            "a {sample_id: value} dict for explicit "
                            "assignment by sample ID. The dict form is "
                            "recommended when targeting specific samples"
                        ),
                    },
                    "key_field": {
                        "type": "string",
                        "description": (
                            "When values is a list, the sample field used "
                            "as the key. Defaults to 'id'"
                        ),
                    },
                },
                "required": ["dataset_name", "field", "values"],
            },
        ),
        Tool(
            name="tag_samples",
            description=(
                "Add tags to samples in a dataset. Optionally filter "
                "which samples to tag using view_stages or sample_ids. "
                "Returns the count of samples that were tagged. Tags "
                "can be used to mark subsets for review, training splits, "
                "or any custom categorization."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "tags": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "List of tags to add to samples",
                    },
                    "view_stages": {
                        "type": "array",
                        "items": {"type": "object"},
                        "description": (
                            "Optional list of serialized view stage dicts "
                            "to select which samples to tag"
                        ),
                    },
                    "sample_ids": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": (
                            "Optional list of specific sample IDs to tag. "
                            "Takes precedence over view_stages"
                        ),
                    },
                },
                "required": ["dataset_name", "tags"],
            },
        ),
        Tool(
            name="untag_samples",
            description=(
                "Remove tags from samples in a dataset. Optionally filter "
                "which samples to untag using view_stages or sample_ids. "
                "Returns the count of samples that were processed."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                    "tags": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": ("List of tags to remove from samples"),
                    },
                    "view_stages": {
                        "type": "array",
                        "items": {"type": "object"},
                        "description": (
                            "Optional list of serialized view stage dicts "
                            "to select which samples to untag"
                        ),
                    },
                    "sample_ids": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": (
                            "Optional list of specific sample IDs to "
                            "untag. Takes precedence over view_stages"
                        ),
                    },
                },
                "required": ["dataset_name", "tags"],
            },
        ),
        Tool(
            name="count_sample_tags",
            description=(
                "Count how many samples have each tag in a dataset. "
                "Returns a {tag: count} dict and total number of distinct "
                "tags. Useful for understanding tag distributions before "
                "filtering or after tagging operations."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "Name of the dataset",
                    },
                },
                "required": ["dataset_name"],
            },
        ),
    ]


_TOOL_NAMES = {
    "add_samples",
    "set_values",
    "tag_samples",
    "untag_samples",
    "count_sample_tags",
}

_TOOL_HANDLERS = {
    "add_samples": lambda a: add_samples(
        a["dataset_name"],
        a["samples"],
    ),
    "set_values": lambda a: set_values(
        a["dataset_name"],
        a["field"],
        a["values"],
        key_field=a.get("key_field"),
    ),
    "tag_samples": lambda a: tag_samples(
        a["dataset_name"],
        a["tags"],
        view_stages=a.get("view_stages"),
        sample_ids=a.get("sample_ids"),
    ),
    "untag_samples": lambda a: untag_samples(
        a["dataset_name"],
        a["tags"],
        view_stages=a.get("view_stages"),
        sample_ids=a.get("sample_ids"),
    ),
    "count_sample_tags": lambda a: count_sample_tags(
        a["dataset_name"],
    ),
}

_REQUIRED_ARGS = {
    "add_samples": ["dataset_name", "samples"],
    "set_values": ["dataset_name", "field", "values"],
    "tag_samples": ["dataset_name", "tags"],
    "untag_samples": ["dataset_name", "tags"],
    "count_sample_tags": ["dataset_name"],
}


async def handle_tool_call(name, arguments):
    """Handles sample manipulation tool calls.

    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 not in _TOOL_NAMES:
            result = format_response(
                None, success=False, error=f"Unknown tool: {name}"
            )
        else:
            missing = [
                arg for arg in _REQUIRED_ARGS[name] if arg not in arguments
            ]
            if missing:
                result = format_response(
                    None,
                    success=False,
                    error=(
                        f"{missing[0]} is required"
                        if len(missing) == 1
                        else f"Required arguments missing: "
                        f"{', '.join(missing)}"
                    ),
                )
            else:
                result = _TOOL_HANDLERS[name](arguments)

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

    except Exception as e:
        logger.error("Error handling sample tool '%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))
        ]
