from enum import IntEnum


class DscCompressionInfo:
    """

    Class contains information about DSC compression used on frame.

    """
    class DscColorFormat(IntEnum):
        """
        Contains values of possible color format.
        """
        CF_NONE = -1
        CF_RGB = 0
        CF_YCbCr_422 = 1
        CF_YCbCr_444 = 2
        CF_YCbCr_420 = 3
        CF_Simple_422 = 4

    def __init__(self):
        self.color_format = DscCompressionInfo.DscColorFormat.CF_NONE
        self.bpp = 0
        self.is_block_prediction_enabled = False
        self.h_slice_size = 0
        self.v_slice_size = 0
        self.buffer_bit_depth = 0
        self.version = (0, 0)
        self.is_simple_as_444 = False

    def is_valid(self) -> bool:
        """
        Return state of the video frame and check color_format, bpp, h and v slice_size and DSC version.
        If everything ok, return True, otherwise - False.

        Returns:
            object of `bool` type.
        """
        return self.color_format is not None and\
               self.bpp > 0 and\
               self.h_slice_size > 0 and\
               self.v_slice_size > 0 and\
               self.version in [(1, 2), (1, 1)]


def create_from_pps(pps_bytearray: bytearray) -> DscCompressionInfo:
    """
    Fill structure 'DscCompressionInfo' from PPS header of the DSC image.

    Returns:
        object of `DscCompressionInfo` type.
    """
    if b'DSCF' not in pps_bytearray:
        if len(pps_bytearray) < 128:
            raise ValueError("Incorrect PPS size!")
        pps = pps_bytearray
    else:
        if len(pps_bytearray) < 132:
            raise ValueError("Incorrect PPS size!")
        pps = pps_bytearray[4:]

    is_yuv = ((pps[4] >> 4) & 1) == 0
    is_simple422 = ((pps[4] >> 3) & 1) == 1
    is_native422 = (pps[88] & 1) == 1
    is_native420 = ((pps[88] >> 1) & 1) == 1
    width = (pps[8] << 8) | pps[9]
    height = (pps[6] << 8) | pps[7]
    slice_height = (pps[10] << 8) | pps[11]
    slice_width = (pps[12] << 8) | pps[13]

    info = DscCompressionInfo()

    info.version = (pps[0] >> 4 & 0xf, pps[0] & 0xf)
    info.buffer_bit_depth = pps[3] & 0xf
    info.is_block_prediction_enabled = bool(pps[4] >> 5 & 0x1)
    info.h_slice_size = int(width / slice_width)
    info.v_slice_size = int(height / slice_height)

    if is_native422 or is_native420:
        info.bpp = int(pps[4] & 0x3 << 8 | pps[5]) / 2
    else:
        info.bpp = int(pps[4] & 0x3 << 8 | pps[5])

    if is_yuv:
        if is_simple422:
            info.color_format = DscCompressionInfo.DscColorFormat.CF_Simple_422
        elif is_native422:
            info.color_format = DscCompressionInfo.DscColorFormat.CF_YCbCr_422
        elif is_native420:
            info.color_format = DscCompressionInfo.DscColorFormat.CF_YCbCr_420
        else:
            info.color_format = DscCompressionInfo.DscColorFormat.CF_YCbCr_444
    else:
        info.color_format = DscCompressionInfo.DscColorFormat.CF_RGB

    return info