from typing import Optional
from UniTAP.libs.lib_tsi import PortIO
from UniTAP.libs.lib_tsi.tsi_private_types import *
from UniTAP.libs.lib_tsi.tsi_types import *
from .private_link_rx_types import RoutedLTStatusPrivate
from .link_rx_types import RoutedLTStatus, RoutedLTConfig
from UniTAP.dev.ports.modules.dpcd import DPCDRegisters


class DisplayPortAUXController:
    """
    Class `DisplayPortAUXController` describes information about DP AUX Controller. Contains following info:
    - configure routed LT `exec_routed_lt`
    - read routed T status `status`
    - request PHY test pattern `request_phy_test_pattern`
    """
    def __init__(self, port_io: PortIO, dpcd: DPCDRegisters):
        self.__io = port_io
        self.__dpcd = dpcd

    def enable(self, enable: bool):
        """
        Enable or disable Routed LT.

        Args:
            enable ('bool')

        """
        status = self.status()

        ctrl = status.dp20_old_lt << 1
        if enable:
            ctrl |= 0x01
        else:
            ctrl &= ~0x01

        res = self.__io.set(TSI_DPRX_LT_ROUTE_CONTROL_W, ctrl, c_uint32)
        if not res:
            raise Exception("The feature requires a license key and the required license key is not installed.")

    def exec_routed_lt(self, config: RoutedLTConfig, use_ta_request: bool = False):
        """

        Execute routed link training with transferred configuration.

        Args:
            config ('RoutedLTConfig')
            use_ta_request (`bool`)
        """
        max_rate_dp14 = 0
        max_rate_dp20 = 0

        if config.is128b132b:
            if config.link_bw == 1:
                max_rate_dp20 = 1
            elif config.link_bw == 2:
                max_rate_dp20 = 3
            elif config.link_bw == 4:
                max_rate_dp20 = 5

            max_rate_dp14 = 0x1e
        else:
            max_rate_dp14 = config.link_bw

        res = self.__set_routed_lt_config(config.int_value())
        if not res:
            raise Exception("The feature requires a license key and the required license key is not installed.")
        self.__set_force_cable_status_to_plugged(True)

        dpcd_rev = 0x14
        dpcd_01_val = max_rate_dp14
        dpcd_02_val = config.lane_count | 0xc0
        dpcd_03_val = 0x81
        dpcd_06_val = 0x03 if config.is128b132b else 0x01

        self.__dpcd.write(0x00, dpcd_rev)
        self.__dpcd.write(0x01, dpcd_01_val)
        self.__dpcd.write(0x02, dpcd_02_val)
        self.__dpcd.write(0x03, dpcd_03_val)
        self.__dpcd.write(0x06, dpcd_06_val)

        self.__dpcd.write(0x0e, 0x80)
        self.__dpcd.write(0x2200, dpcd_rev)
        self.__dpcd.write(0x2201, dpcd_01_val)
        self.__dpcd.write(0x2202, dpcd_02_val)
        self.__dpcd.write(0x2203, dpcd_03_val)
        self.__dpcd.write(0x2206, dpcd_06_val)

        self.__dpcd.write(0x2215, max_rate_dp20)

        if config.is128b132b:
            self.__dpcd.write(0x21, 0x01)
            self.__dpcd.write(0x90, 0xff)
            dsc_data = [0x0f, 0x21, 0x03, 0x03, 0xeb, 0x07, 0x01, 0x00, 0x00, 0x1f, 0x0e, 0x11, 0x08, 0x07, 0x00, 0x00]
            self.__dpcd.write(0x60, dsc_data)
            self.__dpcd.write(0x2217, 0x06)
        else:
            self.__dpcd.write(0x21, 0x00)
            self.__dpcd.write(0x90, 0x00)
            dsc_data = [0x00] * 16
            self.__dpcd.write(0x60, dsc_data)
            self.__dpcd.write(0x2217, 0x00)

        self.__dpcd.write(0x100, 0)
        self.__dpcd.write(0x101, 0)
        self.__dpcd.write(0x103, 0)
        self.__dpcd.write(0x104, 0)
        self.__dpcd.write(0x105, 0)
        self.__dpcd.write(0x106, 0)
        self.__dpcd.write(0x107, 0)
        self.__dpcd.write(0x108, 1)
        self.__dpcd.write(0x206, 0)
        self.__dpcd.write(0x207, 0)

        if use_ta_request:
            self.__dpcd.write(0x219, config.link_bw)
            self.__dpcd.write(0x220, config.lane_count)

            dpcd_218 = self.__dpcd.read(0x218, 1)
            dpcd_218_data = int.from_bytes(dpcd_218.data, byteorder='big', signed=True)
            dpcd_218_data |= 0x01
            dpcd_218_data &= ~0x30
            if config.is128b132b:
                dpcd_218_data |= 0x10
            self.__dpcd.write(dpcd_218.base, dpcd_218_data)

            dpcd_201 = self.__dpcd.read(0x201, 1)
            dpcd_201_data = int.from_bytes(dpcd_201.data, byteorder='big', signed=True)
            dpcd_201_data |= 0x02
            self.__dpcd.write(dpcd_201.base, dpcd_201_data)

            self.__generate_short_hpd_pulse(750)
        else:
            dpcd_201 = self.__dpcd.read(0x201, 1)
            dpcd_201_data = int.from_bytes(dpcd_201.data, byteorder='big', signed=True)
            dpcd_201_data &= ~2
            self.__dpcd.write(dpcd_201.base, dpcd_201_data)

            dpcd_218 = self.__dpcd.read(0x218, 1)
            dpcd_218_data = int.from_bytes(dpcd_218.data, byteorder='big', signed=True)
            dpcd_218_data &= ~0xf
            self.__dpcd.write(dpcd_218.base, dpcd_218_data)

            self.__generate_short_hpd_pulse(500 * 1000)

    def request_phy_test_pattern(self, pattern: int, sq_num: int, config: RoutedLTConfig):
        """

        Request PHY test pattern.

        Args:
            config ('RoutedLTConfig')
            pattern (`int`)
            sq_num (`int`)
        """

        self.__set_force_cable_status_to_plugged(True)

        self.__dpcd.write(0x00, 0x14)
        self.__dpcd.write(0x218, 0x18 if config.is128b132b else 0x08)
        self.__dpcd.write(0x219, config.link_bw)
        self.__dpcd.write(0x220, config.lane_count)

        self.__dpcd.write(0x248, pattern)

        if config.is128b132b and 0x48 <= pattern <= 0x4b:
            self.__dpcd.write(0x249, sq_num)

        if config.is128b132b:
            value = config.ffe | (config.ffe << 4)
        else:
            value = (config.vs & 3) | ((config.pe & 3) << 2)
            value |= value << 4

        self.__dpcd.write(0x206, value)
        self.__dpcd.write(0x207, value)
        self.__dpcd.write(0x260, 0x0)

        dpcd_201 = self.__dpcd.read(0x201, 1)
        dpcd_201_data = int.from_bytes(dpcd_201.data, byteorder='big', signed=True)
        dpcd_201_data |= 0x02
        self.__dpcd.write(dpcd_201.base, dpcd_201_data)

        self.__generate_short_hpd_pulse(750)

    def status(self) -> RoutedLTStatus:
        """
        Get Routed link training status.

        Returns:
            object of `RoutedLTStatus`
        """
        routed_lt_status = self.__read_status()
        if routed_lt_status is None:
            raise Exception("The feature requires a license key and the required license key is not installed.")
        return RoutedLTStatus(enabled=routed_lt_status.enabled,
                              dp20_old_lt=routed_lt_status.dp2OldLt,
                              state=routed_lt_status.state,
                              success=routed_lt_status.success,
                              step=routed_lt_status.step)

    def __set_routed_lt_config(self, value: int) -> bool:
        res = self.__io.set(TSI_DPRX_LT_ROUTE_CREATE, value, c_uint32)
        return res >= 0

    def __set_force_cable_status_to_plugged(self, enable: bool):
        val = 0x3C if enable else 0
        self.__io.set(TSI_DPRX_HPD_FORCE, val, c_int)

    def __read_status(self) -> Optional[RoutedLTStatusPrivate]:
        res, status = self.__io.get(TSI_DPRX_LT_ROUTE_STATUS_R, RoutedLTStatusPrivate)

        if res >= 0:
            return status
        return None

    def __generate_short_hpd_pulse(self, value: int):
        self.__io.set(TSI_DPRX_HPD_PULSE_W, value, c_int)
