from typing import Optional, Union, Type
from UniTAP.libs.lib_tsi import PortIO
from UniTAP.libs.lib_tsi.tsi_private_types import *
from .link_tx_types import *
from .private_link_tx_types import *
from .link_status_common import *


class LinkDisplayPortConfig:
    """
    Class `LinkDisplayPortConfig` allows settings link configuration on Source (TX - transmitter) side.
    - Set configuration `set`.
    - Get current configuration on link `get`.
    """
    def __init__(self, port_io: PortIO, caps: DPTXHWCaps):
        self.__io = port_io
        self.__caps = caps
        self.__dp2_custom_rate = self.__get_dp2_custom_rates()

    def set(self, config: DisplayPortLinkConfig):
        """
        Write new configuration on DP link.

        Args:
            config (`DisplayPortLinkConfig`) - `LinkConfig.DP8b10b`, `LinkConfig.DP128b132b` or `LinkConfig.eDP`.
        """
        if isinstance(config, LinkConfig.DP8b10b):
            self.__set_lane_count_dp14(config.lane_count)
            self.__set_bit_rate_14(config.bit_rate)
            self.__set_lt_features(config)
            self.__set_config_ssc(config.ssc)
            self.__enable_mst(config.mst)
            self.__set_mst_channel_count(config.mst_stream_count)
            self.__try_fec_after_lt(config.fec)
            self.__set_post_lt_features(config)
            self.__sdp_control(config)
        elif isinstance(config, LinkConfig.DP128b132b):
            self.__set_lane_count_dp21(config.lane_count)
            self.__set_bit_rate_21(config.bit_rate)
            self.__set_lt_features(config)
            self.__set_lttpr(config.lttpr)
            self.__set_config_ssc(config.ssc)
            self.__set_post_lt_features(config)
            self.__enable_mst(config.mst)
            self.__set_mst_channel_count(config.mst_stream_count)
            self.__sdp_control(config)
        elif isinstance(config, LinkConfig.eDP):
            self.__set_lane_count_dp14(config.lane_count)
            self.__set_lt_features(config)
            self.__set_edp_sel_rate(config.eDp_cur_rate)
        else:
            raise TypeError("Incorrect DisplayPortLinkConfig format!")

    def get(self, config_type: Union[Type[DisplayPortLinkConfig], None] = None) -> DisplayPortLinkConfig:
        """
        Returns current DP source link configuration.

        Returns:
            object of `DisplayPortLinkConfig` type
        """
        if config_type is not None and issubclass(config_type, LinkConfig.eDP):
            lt_features = self._read_lt_features()
            link_status = LinkConfig.eDP()
            link_status.lane_count = self.__get_lane_count_dp14()
            link_status.force_edp = lt_features.force_edp
            link_status.eDp_cur_rate = self.__get_edp_sel_rate()
            link_status.eDp_supported_rates = self.__get_edp_caps_rates()
            link_status.eDp_aux_preamble = lt_features.edp_aux_preamble
            return link_status
        elif config_type is not None and issubclass(config_type, LinkConfig.DP8b10b):
            return self.__fill_dp14_config()
        elif config_type is not None and issubclass(config_type, LinkConfig.DP128b132b) or \
                self.__link_encoding() == DpLinkEncoding.LE_128b132b:
            return self.__fill_dp21_config()
        elif config_type is None and self.__link_encoding() == DpLinkEncoding.LE_8b10b:
            return self.__fill_dp14_config()
        elif config_type is None and \
                (self.__caps.support_10gbps or self.__caps.support_20gbps or self.__caps.support_13_5gbps) and \
                self.__link_encoding() == DpLinkEncoding.LE_128b132b:
            return self.__fill_dp21_config()
        else:
            raise ValueError('Incorrect link encoding type')

    def __fill_dp14_config(self) -> LinkConfig.DP8b10b:
        lt_features = self._read_lt_features()
        post_lt_features = self.__get_post_lt_features()
        link_status = LinkConfig.DP8b10b()
        link_status.lane_count = self.__get_lane_count_dp14()
        link_status.bit_rate = self.__get_bit_rate_14()
        link_status.mst = self.__mst_enabled()
        link_status.mst_stream_count = self.__mst_stream_count()
        link_status.enhanced_framing_mode = lt_features.framing
        link_status.auto_seed = lt_features.auto_seed
        link_status.ssc = self.__get_config_ssc()
        link_status.fec = self.__fec_after_lt_state()
        link_status.force_edid_timings_after_lt = post_lt_features.force_edid_timings
        link_status.adaptive_sync_auto_enable = post_lt_features.as_auto_enable
        link_status.split_sdp = self.__read_sdp_control().dp_split_sdp != 0 if self.__caps.sdp_split else False
        return link_status

    def __fill_dp21_config(self) -> LinkConfig.DP128b132b:
        link_status = LinkConfig.DP128b132b()
        link_status.lane_count = self.__get_lane_count_dp21()
        if self.__check_custom_bit_rate() != 0:
            link_status.bit_rate = self.__check_custom_bit_rate()
        else:
            link_status.bit_rate = self.__get_bit_rate_21()
        lt_features = self._read_lt_features()
        post_lt_features = self.__get_post_lt_features()
        link_status.force_dp_128_132 = lt_features.force_dp2
        link_status.enhanced_framing_mode = lt_features.framing
        link_status.max_link_bandwidth_supported = lt_features.max_link_bw_policy
        link_status.old_dp2_lt = lt_features.old_dp2_lt
        link_status.lttpr = self.__lttpr_active()
        link_status.try_dp_128_132 = lt_features.try_dp2
        link_status.auto_seed = lt_features.auto_seed
        link_status.ssc = self.__get_config_ssc()
        link_status.force_edid_timings_after_lt = post_lt_features.force_edid_timings
        link_status.adaptive_sync_auto_enable = post_lt_features.as_auto_enable
        link_status.mst = self.__mst_enabled()
        link_status.mst_stream_count = self.__mst_stream_count()
        link_status.crc_16 = self.__read_sdp_control().dp2_crc_sdp != 0 if self.__caps.dp2_sdp_crc else False
        return link_status

    def _read_lt_features(self) -> LTFeatures:
        return self.__io.get(TSI_DPTX_LT_FEATURES, LTFeatures)[1]

    def __get_post_lt_features(self) -> PostLtFeatures:
        return self.__io.get(TSI_DPTX_POST_LT_FEATURES, PostLtFeatures)[1]

    def __set_post_lt_features(self, config: DisplayPortLinkConfig):
        post_lt_features = self.__get_post_lt_features()
        if config.force_edid_timings_after_lt is not None and self.__caps.edid_parser:
            post_lt_features.force_edid_timings = config.force_edid_timings_after_lt
        if config.adaptive_sync_auto_enable is not None and self.__caps.adaptive_sync:
            post_lt_features.as_auto_enable = config.adaptive_sync_auto_enable
        self.__io.set(TSI_DPTX_POST_LT_FEATURES, post_lt_features.value(), c_uint32)

    def __get_dp2_custom_rates(self) -> list:
        rates = self.__io.get(TSI_DP2TX_CUSTOM_RATE_CAPS_R, c_uint32, 32)[1]
        custom_rates = []
        for i in range(len(rates)):
            custom_rates.append(float(rates[i]) * 200000 / 1000000000)

        return custom_rates

    def __set_lane_count_dp14(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_DPTX_LINK_CFG_LANES, lane_count, c_uint32)

    def __get_lane_count_dp14(self) -> int:
        return self.__io.get(TSI_DPTX_LINK_CFG_LANES)[1]

    def __set_lane_count_dp21(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_DP2TX_LT_SLC, lane_count, c_uint32)

    def __get_lane_count_dp21(self) -> int:
        return self.__io.get(TSI_DP2TX_LT_SLC)[1]

    def __set_bit_rate_14(self, bit_rate: Optional[float]):
        if bit_rate is None:
            return
        if int(round(bit_rate / 0.27)) > self.__caps.max_link_rate:
            raise ValueError(f"Following 8b/10b bit rate {bit_rate} is not supported on device. "
                             f"Max supported link rate {round(self.__caps.max_link_rate * 0.27, 2)}")
        self.__io.set(TSI_DPTX_LINK_CFG_BIT_RATE, int(round(bit_rate / 0.27)), c_uint32)

    def __get_bit_rate_14(self) -> float:
        return round(self.__io.get(TSI_DPTX_LINK_CFG_BIT_RATE)[1] * 0.27, 2)

    def __set_bit_rate_21(self, bit_rate: Optional[Union[float, int]]):
        if bit_rate is None:
            return
        if bit_rate in self.__dp2_custom_rate:
            if not self.__caps.support_10gbps:
                raise ValueError(f"Following bit rate {bit_rate} is not supported on device.")
            else:
                res = self.__io.get(TSI_DP2TX_CUSTOM_RATE_MAP, c_uint32, 3)[1]
                self.__io.set(TSI_DP2TX_CUSTOM_RATE_MAP, (int(bit_rate * 1000000000 / 200000), res[1], res[2]),
                              c_uint32, 3)
        elif (bit_rate == 10.0 and not self.__caps.support_10gbps) or \
                (bit_rate == 13.5 and not self.__caps.support_13_5gbps) or \
                (bit_rate == 20.0 and not self.__caps.support_20gbps):
            raise ValueError(f"Following bit rate {bit_rate} is not supported on device.")
        self.__io.set(TSI_DP2TX_LT_SBR, DP21_LinkRateRev.get(bit_rate), c_uint32)

    def __get_bit_rate_21(self) -> float:
        return DP21_LinkRate.get(self.__io.get(TSI_DP2TX_LT_SBR)[1])

    def __support_ssc(self) -> bool:
        return (self.__io.get(TSI_DPTX_DOWNSPREAD_STATUS_R, c_uint32)[1] >> 1) & 1 != 0

    def __set_config_ssc(self, ssc_conf: Optional[SSCConfig]):
        if ssc_conf is None:
            return
        if not self.__support_ssc():
            raise ValueError(f"SSC is not supported on the device.")
        self.__io.set(TSI_DPTX_DOWNSPREAD_CONTROL, int(ssc_conf.enabled), c_uint32)
        if ssc_conf.enabled:
            if not (1 <= ssc_conf.amplitude <= 50):
                raise ValueError(f"Incorrect amplitude value {ssc_conf.amplitude}. Must be from range: 0.1 - 0.5")
            self.__io.set(TSI_DPTX_DOWNSPREAD_AMP, ssc_conf.amplitude, c_uint32)
            if not (30000 <= ssc_conf.frequency <= 63000):
                raise ValueError(f"Incorrect frequency value {ssc_conf.frequency}. Must be from range: 30000 - 63000")
            self.__io.set(TSI_DPTX_DOWNSPREAD_FREQ, ssc_conf.frequency, c_uint32)

    def __get_config_ssc(self) -> SSCConfig:
        if not self.__support_ssc():
            return None
        ssc_conf = SSCConfig()
        ssc_conf.enabled = (self.__io.get(TSI_DPTX_DOWNSPREAD_STATUS_R)[1] & 1) != 0
        ssc_conf.frequency = self.__io.get(TSI_DPTX_DOWNSPREAD_FREQ)[1]
        ssc_conf.amplitude = self.__io.get(TSI_DPTX_DOWNSPREAD_AMP)[1]
        return ssc_conf

    def __dp2_support_rates(self) -> bool:
        return self.__caps.support_10gbps or self.__caps.support_13_5gbps or self.__caps.support_20gbps

    def __set_lttpr(self, enable: Optional[bool]):
        if enable is None:
            return
        if not self.__dp2_support_rates():
            raise ValueError(f"LTTPR is not supported on the device.")
        self.__io.set(TSI_DPTX_LTTPR_CONTROL, int(enable))

    def __set_lt_features(self, config: Optional[DisplayPortLinkConfig]):
        if config is None:
            return
        res = self._read_lt_features()
        if isinstance(config, LinkConfig.DP8b10b):
            if config.enhanced_framing_mode is not None:
                res.framing = int(config.enhanced_framing_mode)
            if config.auto_seed is not None:
                res.auto_seed = int(config.auto_seed)
            if self.__caps.support_10gbps or self.__caps.support_20gbps or self.__caps.support_13_5gbps:
                res.try_dp2 = 0
            res.try_edp = 0
        elif isinstance(config, LinkConfig.DP128b132b):
            if not self.__dp2_support_rates():
                raise ValueError(f"DP 2.1 is not supported on the device.")
            if config.try_dp_128_132 is not None:
                res.try_dp2 = int(config.try_dp_128_132)
            if config.force_dp_128_132 is not None:
                res.force_dp2 = int(config.force_dp_128_132)
            if config.old_dp2_lt is not None:
                res.old_dp2_lt = int(config.old_dp2_lt)
            if config.enhanced_framing_mode is not None:
                res.framing = int(config.enhanced_framing_mode)
            if config.auto_seed is not None:
                res.auto_seed = int(config.auto_seed)
            if config.max_link_bandwidth_supported is not None:
                res.max_link_bw_policy = int(config.max_link_bandwidth_supported)
            res.try_edp = 0
        elif isinstance(config, LinkConfig.eDP):
            if self.__caps.edp:
                if config.force_edp is not None:
                    res.force_edp = config.force_edp
                if config.eDp_aux_preamble is not None:
                    res.edp_aux_preamble = config.eDp_aux_preamble
                res.try_edp = 1
            elif not self.__caps.edp:
                raise ValueError("EDP is not supported.")
        else:
            raise TypeError("Incorrect DisplayPortLinkConfig format!")

        self.__io.set(TSI_DPTX_LT_FEATURES, res.value(), c_uint32)

    def __enable_mst(self, enable: Optional[bool]):
        if enable is None:
            return
        self.__io.set(TSI_DPTX_COMMAND_W, 3 if enable else 4, c_uint32)

    def __set_mst_channel_count(self, count: Optional[int]):
        if count is None:
            return
        if not (0 <= count <= self.__caps.mst_stream_count):
            raise ValueError(f"Incorrect stream count {count}. Available stream count {self.__caps.mst_stream_count}")
        self.__io.set(TSI_PG_ENABLED_STREAM_COUNT, count, c_uint32)

    def __link_encoding(self) -> DpLinkEncoding:
        result, value = self.__io.get(TSI_DPTX_LINK_MODE_R, c_uint)

        if value == 0:
            return DpLinkEncoding.LE_8b10b
        else:
            return DpLinkEncoding.LE_128b132b

    def __mst_enabled(self) -> bool:
        return (self.__io.get(TSI_DPTX_LINK_STATUS_BITS_R)[1] & (1 << 30)) != 0

    def __mst_stream_count(self) -> int:
        result, mst_status = self.__io.get(TSI_DPTX_MST_STATUS_R, c_uint)
        return (mst_status >> 8) & 0xFF

    def __lttpr_active(self) -> Optional[bool]:
        if self.__link_encoding == DpLinkEncoding.LE_128b132b:
            return (self.__io.get(TSI_DPTX_LTTPR_CONTROL, c_uint32)[1] & 1) != 0
        else:
            return None

    def __check_custom_bit_rate(self) -> float:
        res = self.__io.get(TSI_DP2TX_CUSTOM_RATE_MAP, c_uint32, 3)[1]
        for i in range(3):
            if res[i] != 0:
                return res[i] * 200000 / 1000000000
        return 0

    def __try_fec_after_lt(self, enable: bool):
        result = self.__io.get(TSI_DPTX_FEC_CTRL, c_int)
        val = result[1]
        if enable:
            val |= 0x2
        else:
            val &= ~0x2
        self.__io.set(TSI_DPTX_FEC_CTRL, val)

    def __fec_after_lt_state(self) -> bool:
        return (self.__io.get(TSI_DPTX_FEC_CTRL, c_int)[1] & 0x2) != 0

    def __sdp_control(self, config: Optional[DisplayPortLinkConfig]):
        if config is None:
            return
        sdp_control = self.__read_sdp_control()
        if self.__caps.dp2_sdp_crc and isinstance(config, LinkConfig.DP128b132b) and config.crc_16 is not None:
            sdp_control.dp2_crc_sdp = config.crc_16
        if self.__caps.sdp_split and isinstance(config, LinkConfig.DP8b10b) and config.split_sdp is not None:
            sdp_control.dp_split_sdp = config.split_sdp
        self.__write_sdp_control(sdp_control)

    def __write_sdp_control(self, sdp_control: SdpControl):
        self.__io.set(TSI_DPTX_SDP_CTRL, sdp_control.value())

    def __read_sdp_control(self) -> SdpControl:
        return self.__io.get(TSI_DPTX_SDP_CTRL, SdpControl)[1]

    def __set_edp_sel_rate(self, rate: float):
        if self.__caps.edp and rate in self.__get_edp_caps_rates():
            self.__io.set(TSI_DPTX_EDP_LT_SBR, int(rate / 0.0002), c_uint32)
        elif not self.__caps.edp:
            raise ValueError("EDP is not supported.")

    def __get_edp_sel_rate(self) -> int:
        if self.__caps.edp:
            return self.__io.get(TSI_DPTX_EDP_LT_SBR)[1] * 0.0002

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