from ctypes import Structure, c_uint32, c_uint8, c_uint16
from enum import IntEnum
from .link_rx_types import CableCapabilitiesEnum


class DPRXHWCaps(Structure):
    _fields_ = [
        ('mst', c_uint32, 1),
        ('hdcp_1_x', c_uint32, 1),
        ('hdcp_2_x', c_uint32, 1),
        ('fec', c_uint32, 1),
        ('dsc', c_uint32, 1),
        ('alpm', c_uint32, 1),
        ('lanecount3', c_uint32, 1),
        ('edp', c_uint32, 1),
        ('mst_stream_count', c_uint32, 4),
        ('pr', c_uint32, 1),
        ('psr', c_uint32, 1),
        ('', c_uint32, 1),
        ('display_id_capability', c_uint32, 1),
        ('max_link_rate', c_uint32, 8),
        ('force_link_congiguration', c_uint32, 1),
        ('cirdan', c_uint32, 1),
        ('dp_power', c_uint32, 1),
        ('aux_swing', c_uint32, 1),
        ('dp2_custom_rates', c_uint32, 1),
        ('custom_rate_ex', c_uint32, 1),
        ('fec2', c_uint32, 1),
        ('dsc2', c_uint32, 1),
        ('dp2_support_rates', c_uint32, 8),
        ('', c_uint32, 24),
        ('link_pat_log', c_uint32, 1),
        ('vbid_hw_log', c_uint32, 1),
        ('msa_hw_log', c_uint32, 1),
        ('aux_bw_log', c_uint32, 1),
        ('', c_uint32, 28),
        ('scrambler_seed', c_uint32, 1),
        ('sdp_err_counters', c_uint32, 1),
        ('sink_cnt_config', c_uint32, 1),
        ('', c_uint32, 29)
    ]


class DPRXLinkControl(Structure):
    _fields_ = [
        ('tps4_type', c_uint32, 1),
        ('alpm_aux_wake', c_uint32, 1),
        ('old_dp2_lt', c_uint32, 1),
        ('edp_aux_preamble', c_uint32, 1)
    ]

    def value(self) -> int:
        return self.tps4_type << 0 | self.alpm_aux_wake << 1 | self.old_dp2_lt << 2 | self.edp_aux_preamble << 3


class CableTypeEnum(IntEnum):
    Unknown = 0
    Passive = 1
    LRD = 2
    Retimer = 3


class PortCableAttributesEnum(IntEnum):
    DP_CABLE_ATTR_UHBR_10_20_CAPABILITY_POS = 0
    DP_CABLE_ATTR_UHBR_10_20_CAPABILITY_MASK = 0x03 << DP_CABLE_ATTR_UHBR_10_20_CAPABILITY_POS
    DP_CABLE_ATTR_UHBR_10_20_NOT_CAPABLE = 0x00 << DP_CABLE_ATTR_UHBR_10_20_CAPABILITY_POS
    DP_CABLE_ATTR_UHBR_10_CAPABLE = 0x01 << DP_CABLE_ATTR_UHBR_10_20_CAPABILITY_POS
    DP_CABLE_ATTR_UHBR_10_20_CAPABLE = 0x02 << DP_CABLE_ATTR_UHBR_10_20_CAPABILITY_POS

    DP_CABLE_ATTR_UHBR_13_5_CAPABILITY_POS = 2
    DP_CABLE_ATTR_UHBR_13_5_CAPABLE = 0x01 << DP_CABLE_ATTR_UHBR_13_5_CAPABILITY_POS

    DP_CABLE_ATTR_CABLE_TYPE_POS = 3
    DP_CABLE_ATTR_CABLE_TYPE_MASK = 0x7 << DP_CABLE_ATTR_CABLE_TYPE_POS


class PortCableAttributesStruct(Structure):
    _fields_ = [
        ('data', c_uint8)
    ]

    def isUHBR20Supported(self) -> bool:
        return (self.data & PortCableAttributesEnum.DP_CABLE_ATTR_UHBR_10_20_CAPABILITY_MASK) == PortCableAttributesEnum.DP_CABLE_ATTR_UHBR_10_20_CAPABLE

    def isUHBR13_5Supported(self) -> bool:
        return (self.data & PortCableAttributesEnum.DP_CABLE_ATTR_UHBR_13_5_CAPABLE) != 0

    def isUHBR10Supported(self) -> bool:
        return ((self.data & PortCableAttributesEnum.DP_CABLE_ATTR_UHBR_10_20_CAPABILITY_MASK) == PortCableAttributesEnum.DP_CABLE_ATTR_UHBR_10_CAPABLE) or self.isUHBR20Supported()

    def get_cable_caps(self) -> CableCapabilitiesEnum:
        if self.isUHBR20Supported():
            return CableCapabilitiesEnum.DP80
        elif self.isUHBR13_5Supported():
            return CableCapabilitiesEnum.DP54
        elif self.isUHBR10Supported():
            return CableCapabilitiesEnum.DP40
        else:
            return CableCapabilitiesEnum.Unknown

    def get_cable_type(self) -> CableTypeEnum:
        return CableTypeEnum(self.data & PortCableAttributesEnum.DP_CABLE_ATTR_CABLE_TYPE_MASK >> PortCableAttributesEnum.DP_CABLE_ATTR_CABLE_TYPE_POS)


class CableAttributesStruct(Structure):
    _fields_ = [
        ('currentPort', c_uint8),
        ('anotherPort', c_uint8),
        ('reserved', c_uint16)
    ]


class RoutedLTStatusPrivate(Structure):
    _fields_ = [
        ('enabled', c_uint32, 1),
        ('dp2OldLt', c_uint32, 1),
        ('state', c_uint32, 2),
        ('success', c_uint32, 1),
        ('', c_uint32, 11),
        ('step', c_uint32, 9),
        ('', c_uint32, 7)
    ]
