from typing import List, Tuple

from pydantic import Field

from highlighter.agent.capabilities.base_capability import Capability, StreamEvent
from highlighter.core.data_models.data_sample import DataSample

__all__ = ["SelectNthSample"]


class SelectNthSample(Capability):
    """
    Capability for selecting the nth data_sample from a list of data_samples.
    Data samples are ordered by recorded_at timestamp before selection.

    Example element definition:
    {
        "name": "Select Nth Sample",
        "input": [{ "name": "data_samples", "type": "List[DataSample]" }],
        "output": [{ "name": "data_samples", "type": "List[DataSample]" }],
        "deploy": { "local": { "module": "highlighter.agent.capabilities.select_nth_sample",
                               "class_name": "SelectNthSample" } },
        "parameters": {"n": 1, "fallback": false}
    }

    Stream parameters:
        n: int
            The 0-based index of the sample to select. If n is negative or
            exceeds the number of available samples and fallback is false,
            an error is returned.
        fallback: bool
            If true (default), select the highest index that doesn't exceed "n".
            E.g. if "n" is 2 and there is only one image, select that image.
    """

    class StreamParameters(Capability.StreamParameters):
        n: int = Field(description="0-based index of the sample to select")
        fallback: bool = Field(
            description="Whether to run on lower index if 'n' is not available", default=True
        )

    def process_frame(self, stream, data_samples: List[DataSample]) -> Tuple[StreamEvent, dict]:
        """
        Select the nth data_sample from the provided list (zero-indexed).

        Args:
            stream: The stream context
            data_samples: List of DataSample objects to select from

        Returns:
            Tuple of (StreamEvent, dict) where dict contains:
                - data_samples: List containing the selected DataSample if successful
                - diagnostic: Error message if selection failed
        """
        n = self.stream_parameters(stream.stream_id).n
        fallback = self.stream_parameters(stream.stream_id).fallback

        if n < 0:
            return StreamEvent.ERROR, {"diagnostic": f"Index n must be non-negative, got {n}"}

        if len(data_samples) == 0 and fallback:
            return StreamEvent.OKAY, {"data_samples": []}
        elif len(data_samples) <= n:
            if fallback:
                n = len(data_samples) - 1
            else:
                return StreamEvent.ERROR, {
                    "diagnostic": f"Index n={n} out of range for {len(data_samples)} samples"
                }

        # Sort by recorded_at timestamp
        sorted_samples = sorted(data_samples, key=lambda s: s.recorded_at)

        selected_sample = sorted_samples[n]

        return StreamEvent.OKAY, {"data_samples": [selected_sample]}
