from yta_editor_nodes.timeline import TimelineNode
from yta_editor_nodes.timeline.serial import SerialTimelineNode
from yta_editor_time.evaluation_context import EvaluationContext
from abc import ABC, abstractmethod
from typing import Union


class _ParallelTimelineNodeAbstract(TimelineNode, ABC):
    """
    The abstract class of the parallel node.

    This class is limited and needs at least 2 nodes
    to be valid.
    """

    def __init__(
        self,
        # TODO: Put the correct class
        # TODO: Maybe I need the 'node' in the TimelineNode class
        # TODO: Maybe I need a shortcut to 'is_gpu_available'...
        name: str,
        nodes: list[SerialTimelineNode]
    ):
        if len(nodes) < 2:
            raise Exception('A ParallelTimelineNode needs 2 or more nodes to be created.')
        
        self.nodes: list[SerialTimelineNode] = nodes
        """
        The list of nodes that will be executed in parallel
        to obtain the outputs that will be combined.
        """

        super().__init__(
            name = name
        )

    def _resolve_specification(
        self,
        evaluation_context: EvaluationContext
    ) -> 'TimelineNode':
        """
        *For internal use only*

        Resolve the specification of this node by using the
        `evaluation_context` provided, to be able to know
        when the node is active and when we should process
        the input.

        This method must be called each time we modify
        something that affects to the `duration` or other
        field that makes the specifications resolved be
        different.
        """
        for node in self.nodes:
            node._resolve_specification(evaluation_context)

        return self

    # TODO: I think I need an 'is_active_at' method
    def is_active_at(
        self,
        evaluation_context: EvaluationContext
    ) -> bool:
        """
        Flag to indicate if the `evaluation_context`
        provided is in the range of this TimedNode instance, 
        which means that it is in between the limits of this
        `TimelineNode` instance.

        The formula:
        - `start_frame <= evaluation_context.frame_index < end_frame`
        """
        if not self.is_enabled:
            return False
        
        for node in self.nodes:
            if node.is_active_at(evaluation_context):
                return True
            
        return False

    def process(
        self,
        inputs: dict[str, Union['np.ndarray', 'moderngl.Texture']],
        evaluation_context: EvaluationContext,
        output_size: tuple[int, int],
        do_use_gpu: bool = True,
        **kwargs
    ) -> Union['np.ndarray', 'moderngl.Texture']:
        """
        Process the `inputs` provided for the
        `evaluation_context` given, and combine the
        outputs of the different internal nodes and
        return them as output.
        """
        active_nodes = [
            node for node in self.nodes
            if node.is_active_at(evaluation_context)
        ]

        # This shouldn't happen as we know that the parallel
        # node itself is not active, but just in case...
        if not active_nodes:
            # TODO: What do we do with the different inputs? We
            # should return only one from the dict...
            return inputs

        if len(active_nodes) == 1:
            return active_nodes[0].process(
                inputs = inputs,
                evaluation_context = evaluation_context,
                output_size = output_size,
                do_use_gpu = do_use_gpu,
                **kwargs
            )

        outputs = [
            node.process(
                inputs = inputs,
                evaluation_context = evaluation_context,
                output_size = output_size,
                do_use_gpu = do_use_gpu,
                **kwargs
            )
            for node in active_nodes
        ]

        return self._combine_outputs(outputs)
    
    @abstractmethod
    def _combine_outputs(
        self,
        outputs: list[Union['np.ndarray', 'moderngl.Texture']]
    ) -> Union['np.ndarray', 'moderngl.Texture']:
        """
        Combine the different `outputs` into a single one.

        This class must be overwritten by every specific
        class implementation.
        """
        # TODO: This must be according to the type
        pass