from UniTAP.libs.lib_dscl.dscl_types import *
from UniTAP.libs.lib_dscl.dscl import (DSCL_Encode, DSCL_Decode, DSCL_CalculateDSCCRC, DSCL_ExtractPPSFromData,
                                       DSCL_GeneratePPS)
from UniTAP.common import VideoFrameDSC, CompressionInfo
from UniTAP.libs.lib_uicl.uicl_utils import *


def dscl_cf_from_dsc_cf(color_format: CompressionInfo.DscColorFormat) -> DSCL_Colorformat:
    assert isinstance(color_format, CompressionInfo.DscColorFormat)
    if color_format == CompressionInfo.DscColorFormat.CF_RGB:
        return DSCL_Colorformat.Colorformat_RGB
    elif color_format == CompressionInfo.DscColorFormat.CF_YCbCr_444:
        return DSCL_Colorformat.Colorformat_YCbCr444
    elif color_format == CompressionInfo.DscColorFormat.CF_YCbCr_422:
        return DSCL_Colorformat.Colorformat_YCbCr422
    elif color_format == CompressionInfo.DscColorFormat.CF_YCbCr_420:
        return DSCL_Colorformat.Colorformat_YCbCr420
    elif color_format == CompressionInfo.DscColorFormat.CF_Simple_422:
        return DSCL_Colorformat.Colorformat_Simple422
    else:
        return DSCL_Colorformat.Colorformat_MaxValue


def dscl_cf_to_dsc_cf(color_format: DSCL_Colorformat) -> CompressionInfo.DscColorFormat:
    if color_format == DSCL_Colorformat.Colorformat_RGB:
        return CompressionInfo.DscColorFormat.CF_RGB
    elif color_format == DSCL_Colorformat.Colorformat_YCbCr444:
        return CompressionInfo.DscColorFormat.CF_YCbCr_444
    elif color_format == DSCL_Colorformat.Colorformat_YCbCr422:
        return CompressionInfo.DscColorFormat.CF_YCbCr_422
    elif color_format == DSCL_Colorformat.Colorformat_YCbCr420:
        return CompressionInfo.DscColorFormat.CF_YCbCr_420
    elif color_format == DSCL_Colorformat.Colorformat_Simple422:
        return CompressionInfo.DscColorFormat.CF_Simple_422
    else:
        return CompressionInfo.DscColorFormat.CF_NONE


def vm_cf_to_dscl_cf(color_format: ColorInfo.ColorFormat) -> DSCL_Colorformat:
    if color_format == ColorInfo.ColorFormat.CF_RGB:
        return DSCL_Colorformat.Colorformat_RGB
    elif color_format == ColorInfo.ColorFormat.CF_YCbCr_444:
        return DSCL_Colorformat.Colorformat_YCbCr444
    elif color_format == ColorInfo.ColorFormat.CF_YCbCr_422:
        return DSCL_Colorformat.Colorformat_YCbCr422
    elif color_format == ColorInfo.ColorFormat.CF_YCbCr_420:
        return DSCL_Colorformat.Colorformat_YCbCr420
    else:
        return DSCL_Colorformat.Colorformat_MaxValue


def vm_cf_to_dsc_cf(color_format: ColorInfo.ColorFormat) -> CompressionInfo.DscColorFormat:
    if color_format == ColorInfo.ColorFormat.CF_RGB:
        return CompressionInfo.DscColorFormat.CF_RGB
    elif color_format == ColorInfo.ColorFormat.CF_YCbCr_444:
        return CompressionInfo.DscColorFormat.CF_YCbCr_444
    elif color_format == ColorInfo.ColorFormat.CF_YCbCr_422:
        return CompressionInfo.DscColorFormat.CF_YCbCr_422
    elif color_format == ColorInfo.ColorFormat.CF_YCbCr_420:
        return CompressionInfo.DscColorFormat.CF_YCbCr_420
    else:
        return CompressionInfo.DscColorFormat.CF_NONE


def dscl_pps_cf_to_compression_info_cf(dscl_pps: DSCL_PPS) -> CompressionInfo.DscColorFormat:
    if dscl_pps.convert_rgb:
        return CompressionInfo.DscColorFormat.CF_RGB
    elif dscl_pps.native_422:
        return CompressionInfo.DscColorFormat.CF_YCbCr_422
    elif dscl_pps.native_420:
        return CompressionInfo.DscColorFormat.CF_YCbCr_420
    elif dscl_pps.simple_422:
        return CompressionInfo.DscColorFormat.CF_Simple_422
    else:
        return CompressionInfo.DscColorFormat.CF_YCbCr_444


def is_simple_as_444(dscl_pps: DSCL_PPS) -> bool:
    return dscl_pps_cf_to_compression_info_cf(dscl_pps) == CompressionInfo.DscColorFormat.CF_YCbCr_444


def dscl_pps_to_compression_info(dscl_pps: DSCL_PPS) -> CompressionInfo:
    compression_info = CompressionInfo()

    compression_info.color_format = dscl_pps_cf_to_compression_info_cf(dscl_pps)
    compression_info.bpp = ((dscl_pps.bits_per_pixel_high << 8) | dscl_pps.bits_per_pixel_low)
    compression_info.is_block_prediction_enabled = dscl_pps.block_pred_enable
    compression_info.h_slice_size = ((dscl_pps.slice_height_high << 8) | dscl_pps.slice_height_low)
    compression_info.v_slice_size = ((dscl_pps.slice_width_high << 8) | dscl_pps.slice_width_low)
    compression_info.buffer_bit_depth = dscl_pps.linebuf_depth
    compression_info.version = (dscl_pps.dsc_version_major, dscl_pps.dsc_version_minor)
    compression_info.is_simple_as_444 = is_simple_as_444(dscl_pps)

    return compression_info


def image_from_dsc_vf(video_frame: VideoFrameDSC) -> DSCL_Image:
    image = DSCL_Image()
    image.DataSize = len(video_frame.data)

    if isinstance(video_frame.data, bytearray):
        image.DataPtr = (c_uint8 * image.DataSize)(*video_frame.data)
    elif isinstance(video_frame.data, bytes):
        image.DataPtr = (c_uint8 * image.DataSize)(*bytearray(video_frame.data))
    else:
        image.DataPtr = video_frame.data

    return image


def dscl_image_from_vf(video_frame: VideoFrame, parameters: CompressionInfo) -> DSCL_Image:
    image = DSCL_Image()
    image.DataSize = len(video_frame.data)

    if isinstance(video_frame.data, bytearray):
        image.DataPtr = (c_uint8 * image.DataSize).from_buffer(video_frame.data)
    elif isinstance(video_frame.data, bytes):
        image.DataPtr = (c_uint8 * image.DataSize).from_buffer(bytearray(video_frame.data))
    else:
        image.DataPtr = video_frame.data

    return image


def dscl_image_to_dsc_vf(image: DSCL_Image, width: int = 0, height: int = 0, bpc: int = 0) -> VideoFrameDSC:
    video_frame = VideoFrameDSC()

    video_frame.data = bytearray(image.DataPtr[:image.DataSize])
    video_frame.width = width
    video_frame.height = height
    video_frame.color_info.bpc = bpc
    video_frame.color_info.color_format = ColorInfo.ColorFormat.CF_DSC

    return video_frame


def encode_vf(video_frame: VideoFrame, parameters: CompressionInfo) -> VideoFrameDSC:
    if video_frame.color_info.color_format == ColorInfo.ColorFormat.CF_DSC:
        raise AssertionError(f"Image must not have color format DSC. Current format: "
                             f"{video_frame.color_info.color_format.name}")

    uicl_image = image_from_vf(video_frame)

    width, height, color_info, data_info = image_params_to_size_and_ci(uicl_image.Parameters)
    horizontal_slice_count = calculate_slice_number(width, parameters.h_slice_size, color_info.color_format)
    vertical_slice_count = calculate_slice_number(height, parameters.v_slice_size, color_info.color_format)
    dsc_version = DSCL_DscVersion()
    dsc_version.major = parameters.version[0]
    dsc_version.minor = parameters.version[1]
    target_pps = DSCL_GeneratePPS(decoded_image_parameters=uicl_image.Parameters,
                                  is_simple_422=parameters.color_format == CompressionInfo.DscColorFormat.CF_Simple_422,
                                  is_block_prediction_enabled=parameters.is_block_prediction_enabled,
                                  bits_per_pixel=parameters.bpp,
                                  horizontal_slice_count=horizontal_slice_count,
                                  vertical_slice_count=vertical_slice_count,
                                  buffer_bit_depth=parameters.buffer_bit_depth,
                                  dsc_version=dsc_version)
    encoded_image = DSCL_Encode(uicl_image, target_pps)
    encoded_vf = dscl_image_to_dsc_vf(encoded_image, width, height, color_info.bpc)
    encoded_vf.compression_info = parameters

    return encoded_vf


def decode_vf(video_frame: VideoFrameDSC) -> VideoFrame:
    if bytearray("DSCF".encode()) not in video_frame.data:
        video_frame.data = bytearray("DSCF".encode()) + video_frame.data
    encoded_image = image_from_dsc_vf(video_frame)
    decoded_image = DSCL_Decode(encoded_image, decode_simple422_to_444=video_frame.compression_info.is_simple_as_444)
    decoded_vf = image_to_vf(decoded_image)

    return decoded_vf


def dscl_image_calculate_crc(video_frame: VideoFrameDSC):
    image = DSCL_Image()
    image.DataSize = len(video_frame.data)

    if isinstance(video_frame.data, bytearray):
        image.DataPtr = (c_uint8 * image.DataSize)(*video_frame.data)
    elif isinstance(video_frame.data, bytes):
        image.DataPtr = (c_uint8 * image.DataSize)(*bytearray(video_frame.data))
    else:
        image.DataPtr = video_frame.data
    crc = DSCL_CalculateDSCCRC(image)
    return crc.eng0, crc.eng1, crc.eng2


def dscl_data_processing(data: bytearray) -> VideoFrameDSC:
    vf = VideoFrameDSC()
    pps_data = DSCL_ExtractPPSFromData(data)
    vf.compression_info = dscl_pps_to_compression_info(pps_data)
    vf.height = ((pps_data.pic_height_high << 8) | pps_data.pic_height_low)
    vf.width = ((pps_data.pic_width_high << 8) | pps_data.pic_width_low)
    vf.color_info.bpc = pps_data.bits_per_component
    vf.color_info.color_format = ColorInfo.ColorFormat.CF_DSC
    vf.data = data
    return vf


def calculate_slice_size(value: int, slice_number: int,
                         color_format: ColorInfo.ColorFormat = ColorInfo.ColorFormat.CF_RGB):
    slice_size = value / slice_number + (1 if value % slice_number else 0)

    if color_format in [ColorInfo.ColorFormat.CF_YCbCr_422, ColorInfo.ColorFormat.CF_YCbCr_420,
                        ColorInfo.ColorFormat.CF_Y_ONLY]:
        slice_size += (1 if slice_size % 2 else 0)

    return slice_size


def calculate_slice_number(value: int, slice_size: int,
                           color_format: ColorInfo.ColorFormat = ColorInfo.ColorFormat.CF_RGB):
    if color_format in [ColorInfo.ColorFormat.CF_YCbCr_422, ColorInfo.ColorFormat.CF_YCbCr_420,
                        ColorInfo.ColorFormat.CF_Y_ONLY]:
        slice_size -= (1 if slice_size % 2 else 0)

    slice_number = (value + slice_size - 1) // slice_size
    return slice_number