from yta_editor.utils import VideoUtils
from yta_validation.parameter import ParameterValidator
from av.video.frame import VideoFrame
from av.audio.frame import AudioFrame
from typing import Union

import numpy as np


IS_FROM_EMPTY_PART_METADATA = 'is_from_empty_part'
"""
Metadata key to indicate if the frame
has been generated by an empty part
and should be ignored when trying to
combine with others.
"""

class _FrameWrappedBase:
    """
    Class to wrap video and audio frames from
    the pyav library but to support a metadata
    field to inject some information we need
    when processing and combining them.
    """

    @property
    def frame(
        self
    ) -> VideoFrame:
        """
        Get the original pyav VideoFrame.
        """
        return self._frame

    @property
    def is_from_empty_part(
        self
    ) -> bool:
        """
        Flag to indicate if the frame comes from
        an empty part or not, that will be done
        by checking the 'is_from_empty_part' 
        attribute in the metadata.

        This will return True only if the value
        is set and it is set as 'True'.
        """
        is_from_empty_part = self.metadata.get(IS_FROM_EMPTY_PART_METADATA, None)

        return (
            is_from_empty_part == 'True'
            if is_from_empty_part is not None else
            is_from_empty_part
        )

    def __init__(
        self,
        frame: Union[VideoFrame, AudioFrame],
        metadata: dict = {}
    ):
        ParameterValidator.validate_mandatory_instance_of('frame', frame, [VideoFrame, AudioFrame])
        ParameterValidator.validate_mandatory_dict('metadata', metadata)

        self._frame: Union[VideoFrame, AudioFrame] = frame
        """
        The VideoFrame or AudioFrame pyav instance.
        """
        self.metadata: dict = metadata or {}
        """
        The metadata we want to include with the
        frame.
        """

    def __getattr__(
        self,
        name
    ):
        return getattr(self._frame, name)

    def __setattr__(
        self,
        name,
        value
    ):
        super().__setattr__(name, value)
        #setattr(self._frame, name, value)
        # if name in ('_frame', 'metadata'):
        #     super().__setattr__(name, value)
        # else:
        #     setattr(self._frame, name, value)

    def __repr__(
        self
    ):
        cname = self.__class__.__name__
        return f'<{cname} metadata={self.metadata} frame={self._frame!r}>'
    
    def set_is_from_empty_part(
        self,
        is_from_empty_part: bool = True
    ) -> None:
        """
        Add the metadata information to indicate
        that this is a frame that comes from an 
        empty part.
        """
        self.metadata[IS_FROM_EMPTY_PART_METADATA] = (
            'True'
            if is_from_empty_part else
            'False'
        )
    
class VideoFrameWrapped(_FrameWrappedBase):
    """
    Class to wrap video frames from the pyav
    library but to support a metadata field
    to inject some information we need when
    processing and combining them.
    """

    @property
    def has_alpha_layer(
        self
    ) -> bool:
        """
        Flag to indicate if the frame has an alpha
        layer or not, that is defined by an 'a' in
        its format name.
        """
        return VideoUtils.video.alpha.videoframe_has_alpha_layer(self._frame)
    
    @property
    def has_alpha_pixels(
        self
    ) -> bool:
        """
        Flag to indicate if the frame comes with
        at least one pixel that has some kind of
        transparency and has to be considered for
        the alpha blending process because of that.
        """
        return (
            self.has_alpha_layer and
            VideoUtils.video.alpha.numpy_videoframe_has_transparent_pixels(self.as_rgba_numpy)
        )
    
    @property
    def size(
        self
    ) -> tuple[int, int]:
        """
        The size of the frame expressed as (width,
        height).
        """
        return (self.frame.width, self.frame.height)

    @property
    def as_rgb_numpy(
        self
    ) -> 'np.ndarray':
        """
        Get the video frame but as a numpy array with
        the 'rgb24' format, that doesn't include the
        alpha channel.

        The 'rgb24' format is perfect for video
        edition including not the alpha channel.
        """
        return self.frame.to_ndarray(format = 'rgb24')
    
    @property
    def as_rgba_numpy(
        self
    ) -> 'np.ndarray':
        """
        Get the video frame but as a numpy array with
        the 'rgba' format, that includes the alpha
        channel.

        The 'rgba' format is perfect for video edition
        including the alpha channel.
        """
        return self.frame.to_ndarray(format = 'rgba')

    def __init__(
        self,
        frame: VideoFrame,
        metadata: dict = {},
        is_from_empty_part: Union[bool, None] = None,
    ):
        ParameterValidator.validate_mandatory_instance_of('frame', frame, VideoFrame)

        super().__init__(
            frame = frame,
            metadata = metadata
        )

        if is_from_empty_part is not None:
            self.set_is_from_empty_part(is_from_empty_part)
        
# TODO: Maybe move? But careful with
# cyclic import issues
class AudioFrameWrapped(_FrameWrappedBase):
    """
    Class to wrap audio frames from the pyav
    library but to support a metadata field
    to inject some information we need when
    processing and combining them.
    """

    def __init__(
        self,
        frame: AudioFrame,
        metadata: dict = {},
        is_from_empty_part: bool = False
    ):
        ParameterValidator.validate_mandatory_instance_of('frame', frame, AudioFrame)

        super().__init__(
            frame = frame,
            metadata = metadata
        )

        if is_from_empty_part:
            self.set_is_from_empty_part()