from UniTAP.common import ColorInfo, VideoMode
from UniTAP.libs.lib_tsi.tsi_types import TSI_IMF_PK444_RGB_W8_C8, TSI_IMF_PK444_RGB_W16_C12, \
    TSI_IMF_PK444_RGB_W16_C16, TSI_IMF_PK444_YCbCr_W8_C8, TSI_IMF_PK444_YCbCr_W16_C16, \
    TSI_IMF_PK422_CbYCrY_W16_C8, TSI_IMF_PK422_CbYCrY_W16_C16, TSI_IMF_PK420_CxYY0_W8_C8, \
    TSI_IMF_PK420_CxYY0_W8_C10, TSI_IMF_PK420_CxYY_W16_C12, TSI_IMF_PK420_CxYY_W16_C16
from .types import PGColorInfo, PGColorDepth, PGColorimetry, PGDynamicRange


def pg_colorformat_to_ci_colorformat(color_mode: PGColorInfo) -> ColorInfo.ColorFormat:
    if color_mode == PGColorInfo.RGB:
        return ColorInfo.ColorFormat.CF_RGB
    elif color_mode == PGColorInfo.YCbCr444:
        return ColorInfo.ColorFormat.CF_YCbCr_444
    elif color_mode == PGColorInfo.YCbCr422:
        return ColorInfo.ColorFormat.CF_YCbCr_422
    elif color_mode == PGColorInfo.YCbCr420:
        return ColorInfo.ColorFormat.CF_YCbCr_420
    elif color_mode == PGColorInfo.Y_only:
        return ColorInfo.ColorFormat.CF_Y_ONLY
    elif color_mode == PGColorInfo.RAW:
        return ColorInfo.ColorFormat.CF_RAW
    else:
        return ColorInfo.ColorFormat.CF_UNKNOWN


def pg_colorformat_from_ci_colorformat(color_mode: ColorInfo.ColorFormat) -> PGColorInfo:
    if color_mode == ColorInfo.ColorFormat.CF_RGB:
        return PGColorInfo.RGB
    elif color_mode == ColorInfo.ColorFormat.CF_YCbCr_444:
        return PGColorInfo.YCbCr444
    elif color_mode == ColorInfo.ColorFormat.CF_YCbCr_422:
        return PGColorInfo.YCbCr422
    elif color_mode == ColorInfo.ColorFormat.CF_YCbCr_420:
        return PGColorInfo.YCbCr420
    elif color_mode == ColorInfo.ColorFormat.CF_Y_ONLY:
        return PGColorInfo.Y_only
    elif color_mode == ColorInfo.ColorFormat.CF_RAW:
        return PGColorInfo.RAW
    else:
        return PGColorInfo.Unknown


def pg_colorimetry_to_ci_colorimetry(color_mode: PGColorInfo, colorimetry: PGColorimetry)\
        -> ColorInfo.Colorimetry:
    if color_mode == PGColorInfo.RGB:
        return ColorInfo.Colorimetry.CM_sRGB
    elif color_mode in [PGColorInfo.YCbCr444, PGColorInfo.YCbCr422, PGColorInfo.YCbCr420]:
        if colorimetry == PGColorimetry.ITU601:
            return ColorInfo.Colorimetry.CM_ITUR_BT601
        elif colorimetry == PGColorimetry.ITU709:
            return ColorInfo.Colorimetry.CM_ITUR_BT709
        else:
            return ColorInfo.Colorimetry.CM_RESERVED
    else:
        return ColorInfo.Colorimetry.CM_RESERVED


def pg_colorimetry_from_ci_colorimetry(colorimetry: ColorInfo.Colorimetry) -> PGColorimetry:
    if colorimetry == ColorInfo.Colorimetry.CM_ITUR_BT601:
        return PGColorimetry.ITU601
    elif colorimetry == ColorInfo.Colorimetry.CM_ITUR_BT709:
        return PGColorimetry.ITU709
    else:
        return PGColorimetry.ITU601


def pg_bpc_to_ci_bpc(bpc: PGColorDepth) -> int:
    if bpc == PGColorDepth.BPC6:
        return 6
    elif bpc == PGColorDepth.BPC8:
        return 8
    elif bpc == PGColorDepth.BPC10:
        return 10
    elif bpc == PGColorDepth.BPC12:
        return 12
    elif bpc == PGColorDepth.BPC16:
        return 16
    elif bpc == PGColorDepth.BPC7:
        return 7
    elif bpc == PGColorDepth.BPC14:
        return 14
    return 0


def pg_bpc_from_ci_bpc(bpc: int) -> PGColorDepth:
    if bpc == 6:
        return PGColorDepth.BPC6
    elif bpc == 8:
        return PGColorDepth.BPC8
    elif bpc == 10:
        return PGColorDepth.BPC10
    elif bpc == 12:
        return PGColorDepth.BPC12
    elif bpc == 16:
        return PGColorDepth.BPC16
    elif bpc == 7:
        return PGColorDepth.BPC7
    elif bpc == 14:
        return PGColorDepth.BPC14
    return PGColorDepth.Unknown


def pg_pixel_format_from_vm(vm: VideoMode) -> int:
    bpc = vm.color_info.bpc
    color_format = vm.color_info.color_format

    if color_format == ColorInfo.ColorFormat.CF_RGB:
        if bpc == 8:
            return TSI_IMF_PK444_RGB_W8_C8
        elif bpc in [10, 12]:
            return TSI_IMF_PK444_RGB_W16_C12
        elif bpc == 16:
            return TSI_IMF_PK444_RGB_W16_C16
    elif color_format == ColorInfo.ColorFormat.CF_YCbCr_444:
        if bpc == 8:
            return TSI_IMF_PK444_YCbCr_W8_C8
        elif bpc in [10, 12, 16]:
            return TSI_IMF_PK444_YCbCr_W16_C16
    elif color_format == ColorInfo.ColorFormat.CF_YCbCr_422:
        if bpc == 8:
            return TSI_IMF_PK422_CbYCrY_W16_C8
        elif bpc in [10, 12, 16]:
            return TSI_IMF_PK422_CbYCrY_W16_C16
    elif color_format == ColorInfo.ColorFormat.CF_YCbCr_420:
        if bpc == 8:
            return TSI_IMF_PK420_CxYY0_W8_C8
        elif bpc == 10:
            return TSI_IMF_PK420_CxYY0_W8_C10
        elif bpc == 12:
            return TSI_IMF_PK420_CxYY_W16_C12
        elif bpc == 16:
            return TSI_IMF_PK420_CxYY_W16_C16

    return 0xffffffff


def pg_dynamicrange_to_ci_dynamicrange(dynamic_range: PGDynamicRange) -> ColorInfo.DynamicRange:
    if dynamic_range == PGDynamicRange.CTA:
        return ColorInfo.DynamicRange.DR_CTA
    elif dynamic_range == PGDynamicRange.VESA:
        return ColorInfo.DynamicRange.DR_VESA
    else:
        return ColorInfo.DynamicRange.DR_UNKNOWN


def pg_dynamicrange_from_ci_dynamicrange(dynamic_range: ColorInfo.DynamicRange) -> PGDynamicRange:
    if dynamic_range == ColorInfo.DynamicRange.DR_CTA:
        return PGDynamicRange.CTA
    else:
        return PGDynamicRange.VESA


def pg_timingflags_from_vm(vm: VideoMode) -> int:
    timing_flags = 0

    timing_flags |= pg_bpc_from_ci_bpc(vm.color_info.bpc)
    timing_flags |= (1 << 9) if vm.timing.hswidth < 0 else 0
    timing_flags |= (1 << 10) if vm.timing.vswidth < 0 else 0
    timing_flags |= pg_colorformat_from_ci_colorformat(vm.color_info.color_format) << 11
    timing_flags |= pg_colorimetry_from_ci_colorimetry(vm.color_info.colorimetry) << 16
    timing_flags |= pg_dynamicrange_from_ci_dynamicrange(vm.color_info.dynamic_range) << 31

    return timing_flags