import warnings
from typing import List, Optional, Union, Type

from .link_rx_types import LinkCapabilities, LinkEDPCapabilities, DisplayPortLinkCaps
from .private_link_rx_types import DPRXHWCaps, DPRXLinkControl
from UniTAP.libs.lib_tsi import PortIO
from UniTAP.libs.lib_tsi.tsi_types import *


class LinkDisplayPortCaps:
    """
    Class `LinkDisplayPortCaps` allows settings link capabilities on Sink (RX - receiver) side.
    - Set configuration `set`.
    - Get current configuration on link `link_caps_status`.
    """
    def __init__(self, port_io: PortIO, caps: DPRXHWCaps):
        self.__io = port_io
        self.__caps = caps

    def set(self, capabilities: Union[LinkCapabilities, LinkEDPCapabilities]):
        """
        Set new settings on link. Only those values will be written that were specified.

        Args:
            capabilities (`LinkCapabilities` or `LinkEDPCapabilities`)
        """
        if isinstance(capabilities, LinkCapabilities):
            self.__set_max_lanes(capabilities.max_lane)
            self.__set_max_bitrate(capabilities.bit_rate)
            self.__set_force_cable_status_to_plugged(capabilities.force_cable_status_to_plugged)
            self.__set_dp_link_control(old_dp2_lt=capabilities.old_dp_2_0_lt)
            self.__set_dp_128_132_bitrates(capabilities.dp_128_132_bitrates)
            self.__set_bitrate_override(10.0, capabilities.override_10g)
            self.__set_dsc(capabilities.dsc)
            self.__set_ss_sbm(capabilities.ss_sbm)
            self.__set_fec(capabilities.fec)
            self.__set_tps3(capabilities.tps3)
            self.__set_tps4(capabilities.tps4)
            self.__set_mst(capabilities.mst)
            self.__set_edp_support(False)
            self.__set_mst_sink_count(capabilities.mst_sink_count)
        elif isinstance(capabilities, LinkEDPCapabilities):
            self.__set_max_lanes(capabilities.max_lane)
            self.__set_edp_sel_rates(capabilities.eDp_cur_rate)
            self.__set_dp_link_control(edp_aux_preamble=capabilities.eDp_aux_preamble)
            self.__set_edp_support(capabilities.eDp_support)
        else:
            raise TypeError("Incorrect LinkCapabilities format!")

    def link_caps_status(self, config_type: Union[Type[DisplayPortLinkCaps], None] = None) -> DisplayPortLinkCaps:
        """
        Returns current configuration on link `LinkCapabilities`.

        Returns:
            object of `LinkCapabilities` type
        """
        if config_type is None or issubclass(config_type, LinkCapabilities):
            link_status = LinkCapabilities()
            link_status.max_lane = self.__get_max_lanes()
            link_status.bit_rate = self.__get_max_bitrate()
            link_status.force_cable_status_to_plugged = self.__get_force_cable_status_to_plugged()
            link_status.old_dp_2_0_lt = self.__get_dp_link_control().old_dp2_lt
            link_status.dp_128_132_bitrates = self.__get_dp_128_132_bitrates()
            link_status.override_10g = self.__get_bitrate_override(10.0)
            link_status.dsc = self.__get_dsc()
            link_status.ss_sbm = self.__get_ss_sbm()
            link_status.fec = self.__get_fec()
            link_status.tps3 = self.__get_tps3()
            link_status.tps4 = self.__get_tps4()
            link_status.mst = self.__get_mst()
            link_status.mst_sink_count = self.__get_mst_sink_count()
            return link_status
        elif issubclass(config_type, LinkEDPCapabilities):
            link_status = LinkEDPCapabilities()
            link_status.max_lane = self.__get_max_lanes()
            link_status.eDp_cur_rate = self.__get_edp_sel_rates()
            link_status.eDp_supported_rates = self.__get_edp_caps_rates()
            link_status.eDp_aux_preamble = self.__get_dp_link_control().edp_aux_preamble
            link_status.eDp_support = self.__get_edp_support()
            return link_status
        else:
            raise TypeError("Incorrect LinkCapabilities format!")

    def __set_max_lanes(self, lane_count: Optional[int]):
        if lane_count is None:
            return
        if lane_count not in [1, 2, 4]:
            raise ValueError(f"Incorrect lane count number {lane_count}. Must be from list: 1, 2, 4")
        self.__io.set(TSI_DPRX_MAX_LANES, lane_count, c_int)

    def __get_max_lanes(self) -> int:
        return self.__io.get(TSI_DPRX_MAX_LANES)[1]

    def __set_max_bitrate(self, bitrate: Optional[float]):
        if bitrate is None:
            return
        bitrate = round(bitrate / 0.27)
        if bitrate not in [6, 10, 20, 25, 30]:
            raise ValueError(f"Incorrect bit rate number {bitrate}. Must be from list: 1.62, 2.7, 5.4, 6.75, 8.1")
        self.__io.set(TSI_DPRX_MAX_LINK_RATE, bitrate, c_int)

    def __get_max_bitrate(self) -> int:
        return round(self.__io.get(TSI_DPRX_MAX_LINK_RATE)[1] * 0.27, 2)

    def __set_bitrate_override(self, to_override: float, with_override: Optional[float]):
        if with_override is None:
            return

        if not isinstance(to_override, float) or not isinstance(with_override, float):
            raise TypeError("DP 128b/132b bitrate override settings must float type!")

        list_value = self.__io.get(TSI_DP2RX_CUSTOM_RATE_MAP, c_uint, 3)[1]

        if with_override not in [2.5, 2.7, 5.0, 5.4]:
            raise ValueError(f"Incorrect bit rate number {with_override}. "
                             f"Must be from list: 2.5, 2.7, 5.0, 5.4")
        list_value[0] = int(with_override * 1000000000 / 200000)
        self.__io.set(TSI_DP2RX_CUSTOM_RATE_MAP, list_value, c_uint, 3)

    def __get_bitrate_override(self, to_override: float) -> Optional[float]:
        if not isinstance(to_override, float):
            raise TypeError("DP 128b/132b bitrate override settings must float type!")

        override_list = self.__io.get(TSI_DP2RX_CUSTOM_RATE_MAP, c_uint, 3)[1]
        if override_list.count(0) == len(override_list):
            return None

        if to_override == 10.0 and override_list[0] > 0:
            return round((override_list[0] * 200000 / 1000000000), 2)

        return None

    def __set_dp_link_control(self, old_dp2_lt: Optional[bool] = None, edp_aux_preamble: Optional[bool] = None):
        ctrl = self.__get_dp_link_control()

        if old_dp2_lt is not None:
            ctrl.old_dp2_lt = old_dp2_lt
        if edp_aux_preamble is not None:
            ctrl.edp_aux_preamble = edp_aux_preamble

        self.__io.set(TSI_DPRX_LINK_CONTROL, ctrl.value(), c_uint32)

    def __get_dp_link_control(self) -> DPRXLinkControl:
        return self.__io.get(TSI_DPRX_LINK_CONTROL, DPRXLinkControl)[1]

    def __set_dp_128_132_bitrates(self, rates: Optional[List[float]]):
        flags = self.__io.get(TSI_DPRX_LINK_FLAGS)[1] & ~0xc0
        if rates is None:
            return
        elif isinstance(rates, list) and len(rates) == 0:
            flags &= ~(1 << 4)
            self.__io.set(TSI_DPRX_LINK_FLAGS, flags, c_int)
        else:
            flags |= (1 << 4)
            self.__io.set(TSI_DPRX_LINK_FLAGS, flags, c_int)
            val = 0
            if 10.0 in rates:
                val |= 0x01
            if 13.5 in rates:
                val |= 0x04
            if 20.0 in rates:
                val |= 0x02
            self.__io.set(TSI_DP2RX_LINK_RATE_CAPS, val, c_int)

    def __get_dp_128_132_bitrates(self) -> Optional[List[float]]:
        val = (self.__io.get(TSI_DPRX_LINK_FLAGS)[1] >> 4) & 0x1
        if val:
            rates = self.__io.get(TSI_DP2RX_LINK_RATE_CAPS)[1]
            list_rates = []
            if rates & 0x01:
                list_rates.append(10.0)
            if rates & 0x04:
                list_rates.append(13.5)
            if rates & 0x02:
                list_rates.append(20.0)
            return list_rates
        else:
            return None

    def __set_force_cable_status_to_plugged(self, enable: Optional[bool]):
        if enable is None:
            return
        val = 0x3C if enable else 0
        self.__io.set(TSI_DPRX_HPD_FORCE, val, c_int)

    def __get_force_cable_status_to_plugged(self) -> bool:
        return bool(self.__io.get(TSI_DPRX_HPD_FORCE)[1])

    def __set_dsc(self, enable: Optional[bool]):
        if enable is None:
            return
        val = 0x1 if enable else 0x0
        if self.__caps.dsc:
            self.__io.set(TSI_DPRX_DSC_CONTROL, val, c_int)

    def __get_dsc(self) -> Optional[bool]:
        if self.__io.get(TSI_DPRX_DSC_STATUS_R)[1] & 0x1:
            return bool(self.__io.get(TSI_DPRX_DSC_CONTROL)[1] & 0x1)

    def __set_fec(self, enable: Optional[bool]):
        if enable is None:
            return
        val = 0x1 if enable else 0x0
        if self.__caps.fec:
            self.__io.set(TSI_DPRX_FEC_CONTROL, val, c_int)

    def __get_fec(self) -> Optional[bool]:
        if self.__caps.fec:
            return bool(self.__io.get(TSI_DPRX_FEC_CONTROL)[1] & 0x1)

    def __set_mst(self, enable: Optional[bool]):
        if enable is None:
            return
        flags = self.__io.get(TSI_DPRX_LINK_FLAGS)[1] & ~0xc0
        if enable:
            flags |= 1
        else:
            flags &= ~1
        if self.__caps.mst:
            self.__io.set(TSI_DPRX_LINK_FLAGS, flags, c_int)

    def __get_mst(self) -> Optional[bool]:
        if self.__caps.mst:
            return bool(self.__io.get(TSI_DPRX_LINK_FLAGS)[1] & 0x1)

    def __set_mst_sink_count(self, count: int):
        if count is None:
            return
        if self.__caps.sink_cnt_config and self.__caps.mst:
            if count > self.__caps.mst_stream_count:
               warnings.warn(f"Maximum sink count: {self.__caps.mst_stream_count}")
            else:
                self.__io.set(TSI_DPRX_MST_SINK_COUNT, count, c_uint32)
        else:
            warnings.warn("MST or Custom Sink Count are not supported")

    def __get_mst_sink_count(self) -> int:
        if self.__caps.sink_cnt_config and self.__caps.mst:
            return int(self.__io.get(TSI_DPRX_MST_SINK_COUNT)[1])
        else:
            return 0

    def __set_tps3(self, enable: Optional[bool]):
        if enable is None:
            return
        flags = self.__io.get(TSI_DPRX_LINK_FLAGS)[1] & ~0xc0
        if enable:
            flags |= (1 << 1)
        else:
            flags &= ~(1 << 1)
        if self.__caps.fec:
            self.__io.set(TSI_DPRX_LINK_FLAGS, flags, c_int)

    def __get_tps3(self) -> Optional[bool]:
        if self.__caps.fec:
            return bool((self.__io.get(TSI_DPRX_LINK_FLAGS)[1] >> 1) & 0x1)

    def __set_tps4(self, enable: Optional[bool]):
        if enable is None:
            return
        flags = self.__io.get(TSI_DPRX_LINK_FLAGS)[1] & ~0xc0
        if enable:
            flags |= (1 << 2)
        else:
            flags &= ~(1 << 2)
        if self.__caps.fec:
            self.__io.set(TSI_DPRX_LINK_FLAGS, flags, c_int)

    def __get_tps4(self) -> Optional[bool]:
        if self.__caps.fec:
            return bool((self.__io.get(TSI_DPRX_LINK_FLAGS)[1] >> 2) & 0x1)

    def __set_ss_sbm(self, enable: bool):
        if enable is None:
            return
        if self.__caps.dp2_support_rates:
            flags = self.__io.get(TSI_DPRX_LINK_FLAGS)[1] & ~0xc0
            if enable:
                flags |= (1 << 5)
            else:
                flags &= ~(1 << 5)
            self.__io.set(TSI_DPRX_LINK_FLAGS, flags, c_int)

    def __get_ss_sbm(self) -> Optional[bool]:
        if self.__caps.dp2_support_rates:
            return bool((self.__io.get(TSI_DPRX_LINK_FLAGS)[1] >> 5) & 0x1)

    def __set_edp_support(self, enable: bool):
        if self.__caps.edp:
            flags = self.__io.get(TSI_DPRX_LINK_FLAGS)[1] & ~0xc0
            if enable:
                flags |= (1 << 3)
            else:
                flags &= ~(1 << 3)
            self.__io.set(TSI_DPRX_LINK_FLAGS, flags, c_int)

    def __get_edp_support(self) -> Optional[bool]:
        if self.__caps.edp:
            return bool((self.__io.get(TSI_DPRX_LINK_FLAGS)[1] >> 3) & 0x1)

    def __set_edp_sel_rates(self, rates: List[float]):
        if self.__caps.edp:
            rates.sort()
            rates = rates[:8]
            to_write_list = [0 for i in range(8)]
            j = 0
            for i in range(len(rates)):
                if rates[i] in self.__get_edp_caps_rates():
                    to_write_list[j] = int(rates[i] / 0.0002)
                    j += 1

            rates.sort()
            self.__io.set(TSI_DPRX_EDP_SEL_RATES, to_write_list, data_type=c_uint32, data_count=8)

    def __get_edp_sel_rates(self) -> List[float]:
        if self.__caps.edp:
            rates = self.__io.get(TSI_DPRX_EDP_SEL_RATES, data_type=c_uint32, data_count=8)[1]
            res_rates = []
            for rate in rates:
                if rate == 0:
                    return res_rates
                res_rates.append(rate * 0.0002)
            return res_rates

    def __get_edp_caps_rates(self) -> List[float]:
        if self.__caps.edp:
            rates = self.__io.get(TSI_DPRX_EDP_CAPS_RATES_R, data_type=c_uint32, data_count=32)[1]
            new_rates = []
            for rate in rates:
                if rate > 0:
                    new_rates.append(rate * 0.0002)
            return new_rates
