# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""classes containing the workflow steps supported by the high level interface"""

import collections.abc
import numpy as np
from gammapy.utils.scripts import make_name, requires_module

try:
    import ray
except ImportError:
    ray = None  # Fallback if ray is not available


class Product:
    """Represents a product generated during a processing step in a workflow."""

    def __init__(self, name=None, step_name=None, data=None):
        """
        Initialize a Product instance.

        Parameters
        ----------
        name : str, optional
            A label or identifier for the product.
            Default is None, and the name is randomly generated.
        step_name : str, optional
            The name of the processing step that produced this product.
            Default is None.
        data : object, optional
            The data associated with the product, it can be any object.
        """

        if step_name is None:
            step_name = ""
        self.step_name = step_name
        self.data = data
        if name is None:
            name = make_name()
        self.name = name
        self.pid = make_name()  # unique id generated for each product

    @property
    @requires_module("ray")
    def needs_update(self):
        """
        Check whether the product's data needs to be updated.
        Only available if `ray` is available.

        Returns
        -------
        bool
            True if `data` is a `ray.ObjectRef`, indicating that the data is a reference
            to a remote object and may need to be resolved or updated.
        """
        return isinstance(self.data, ray.ObjectRef)

    def copy(self):
        """
        Create a shallow copy of the product with a new unique identifier.

        Returns
        -------
        Product
            A new `Product` instance with the same `tag`, `step_name`, and `data`,
            but a newly generated `name`.
        """
        return self.__class__(name=self.name, step_name=self.step_name, data=self.data)


class Products(collections.abc.MutableSequence):
    """
    Container for managing a list of `Product` instances generated by the workflow.

    This class provides list-like behavior and additional utilities for
    selecting, copying, and updating products, including support for
    Ray-based asynchronous data resolution.

    Parameters
    ----------
    products : list of `Product`, optional
        Initial list of products. Defaults to an empty list.
    """

    def __init__(self, products=None):
        if products is None:
            products = []
        self._products = products

    @property
    def names(self):
        """List of product names."""
        return [p.name for p in self._products]

    @property
    def pids(self):
        """List of product names."""
        return [p.pid for p in self._products]

    @property
    def unique_names(self):
        """List of unique product names. Return formatted as step_name.product_name.product_id."""
        names = []
        for p in self._products:
            components = [p.step_name, p.name, p.pid]
            name = ".".join(components)
            names.append(name)
        return names

    @property
    def data(self):
        """List of `Product` data."""
        return [p.data for p in self._products]

    def index(self, key):
        """
        Get the index of a product by pid or position.

        Parameters
        ----------
        key : int, slice, or str
            Index or name of the product.

        Returns
        -------
        int or slice
            Index corresponding to the key.
        """
        if isinstance(key, (int, slice)):
            return key
        elif isinstance(key, str):
            return self.names.index(key)

    def copy(self):
        """Create a deep copy of the products container."""
        return self.__class__([p.copy() for p in self])

    def __len__(self):
        return len(self._products)

    def __getitem__(self, key):
        if isinstance(key, np.ndarray) and key.dtype == bool:
            return self.__class__(list(np.array(self._products)[key]))
        else:
            return self._products[self.index(key)]

    def __delitem__(self, key):
        del self._products[self.index(key)]

    def insert(self, idx, model):
        self._products.insert(idx, model)

    def __setitem__(self, key, product):
        if isinstance(product, Product):
            if isinstance(key, np.ndarray) and key.dtype == bool:
                ind = np.nonzero(key)[0]
                for idx in ind:
                    self._products[int(idx)] = product
            else:
                self._products[self.index(key)] = product
        else:
            raise TypeError(f"Invalid type: {product!r}")

    def select_from_list(self, requirements):
        """
        Select products based on a list of criteria.

        Parameters
        ----------
        requirements : list of dict
            List of selection criteria dictionaries passed to select method.

        Returns
        -------
        Products
            A new `Products` instance with selected products.
        """
        selection = Products([])
        for req in requirements:
            selection.extend(self.select(**req))
        return selection

    def select(self, name=None, step_name=None, pid=None):
        """
        Select products matching the given criteria.

        Parameters
        ----------
        name : str, optional
            Name to match. Default is None.
            Default is None, and name filter is not applied.
        step_name : str, optional
            Step name to match.
            Default is None, and `step_name` filter is not applied.
        pid : str, optional
            Product pid to match.
            Default is None, and `pid` filter is not applied.

        Returns
        -------
        Products
            A new `Products` instance with matching products.
        """

        mask = self.selection_mask(name=name, step_name=step_name, pid=pid)
        return self[mask]

    def selection_mask(self, name=None, step_name=None, pid=None):
        """
        Create a boolean mask for selected products.

        Parameters
        ----------
        name : str, optional
            Names to match.
            Default is None, and name filter is not applied.
        step_name : str, optional
            Step names to match.
            Default is None, and `step_name` filter is not applied.
        pid : str, optional
            Product pid to match.
            Default is None, and `pid` filter is not applied.

        Returns
        -------
        mask : np.ndarray
            Boolean array indicating selected products.
        """

        mask = np.ones(len(self), dtype=bool)
        for idx, p in enumerate(self._products):
            if name is not None:
                mask[idx] &= p.name == name
            if step_name is not None:
                mask[idx] &= p.step_name == step_name
            if pid is not None:
                mask[idx] &= p.pid == pid
        return mask

    @requires_module("ray")
    def get(self):
        """
        Resolve and update products whose data is only a reference (`ray.ObjectRef`).
        Only available if `ray` is available.

        This method fetches remote data using `ray.get()` and updates
        the corresponding products in-place.

        """
        ind = np.nonzero([p.needs_update for p in self._products])[0]
        refs_to_run = {data for data in np.array(self.data)[ind]}  # unique objects
        if refs_to_run:
            results = ray.get(list(refs_to_run))
            products = [item for sublist in results for item in sublist]  # flatten
            for product in products:
                mask = self.selection_mask(pid=product.pid)
                self[mask] = product
