from UniTAP.libs.lib_tsi.tsi_types import *
from UniTAP.utils import function_scheduler
from .ag_utils import *
from UniTAP.libs.lib_tsi.tsi_io import PortIO, PortProtocol
from UniTAP.dev.modules import MemoryManager
from .types import *
from .private_types import *
from typing import Union


class AudioGenerator:
    """
    Class `AudioGenerator` allows working with generating audio from Source (TX - transmitter). You can configure
    audio generator `setup`, apply settings and start generate audio `apply`, stop generate audio `stop_generate`,
    read audio generator `status` and get current `audio_mode`.
    """
    def __init__(self, port_io: PortIO, memory_manager: MemoryManager):
        self.__io = port_io
        self.__memory_manager = memory_manager
        self.__status = AGStatus.Unknown
        self.__audio_mode = AudioMode()

    def setup(self, audio_mode: AudioMode = AudioMode(),
              audio_pattern: Union[AudioPattern, str] = AudioPattern.SignalSine,
              signal_frequency: int = 1000, amplitude: int = 60):
        """
        Configure audio generator. Possible two variants of configuration:
        - From 'wav' or 'bin' file.
        - From `AudioPattern` parameters.

        Args:
            audio_mode (AudioMode) - object of `AudioMode`
            audio_pattern (Union[AudioPattern, str]) - object of `AudioPattern` or path to audio file ('bin' or 'wave')
            signal_frequency (int)
            amplitude (int)
        """
        self.stop_generate()
        if isinstance(audio_pattern, str) and check_file_format(audio_pattern) != AudioFileFormat.UNKNOWN:
            if check_file_format(audio_pattern) == AudioFileFormat.BIN:
                data, _, size = load_from_bin_file(path=audio_pattern)
            else:
                data, audio_mode, size = load_from_wave_file(path=audio_pattern)

            self.__memory_manager.make_default()
            self.__memory_manager.set_memory_block_index(MemoryManager.MemoryOwner.MO_AudioGenerator)
            self.__io.set(TSI_AUDGEN_SIGNAL_BLOCK, MemoryManager.MemoryOwner.MO_AudioGenerator.value)
            self.__memory_manager.memory_write(data, size)
            self.__io.set(TSI_AUDGEN_AUDIO_SIZE, size)

            ag_struct = self.__io.get(TSI_AUDGEN_CONFIG, AudioConfigStructure)[1]
            ag_struct.non_lpcm = 0
            ag_struct.loop = 1

            audio_sts = create_audio_sts(audio_mode=audio_mode)
            self.__io.set(TSI_AUDGEN_CHANNELS_STS, audio_sts, data_type=c_uint8, data_count=len(audio_sts))

            if self.__io.protocol() == PortProtocol.HDMI:
                tp = 0

                if audio_mode.channel_count == 8 and audio_mode.sample_rate >= 64000:
                    tp = 3

                ag_struct.n_selector = 1
                ag_struct.cts_selector = 1

                self.__io.set(TSI_AUDGEN_CONFIG, ag_struct.value())

                self.__io.set(TSI_AUDGEN_PACKET_TYPE, tp)
            else:
                self.__io.set(TSI_AUDGEN_CONFIG, ag_struct.value())

            self.__io.set(TSI_AUDGEN_SIGNAL_TYPE, AudioPattern.CustomAudio.value)
            self.__io.set(TSI_AUDGEN_CHANNEL_COUNT, audio_mode.channel_count)
            self.__io.set(TSI_AUDGEN_SAMPLE_RATE, audio_mode.sample_rate)
            self.__io.set(TSI_AUDGEN_SAMPLE_SIZE, audio_mode.bits)
        elif isinstance(audio_pattern, AudioPattern):
            list_sample_rate = [22050, 44100, 88200, 176400, 24000, 48000, 96000, 192000, 32000, 768000]
            list_bits = [16, 20, 24]
            audio_pattern = AudioPattern(audio_pattern.value)

            if not (1 <= audio_mode.channel_count <= 8):
                raise ValueError(f"'Channel count' must be in range: 1 - 8")

            if audio_mode.sample_rate not in list_sample_rate:
                raise ValueError(f"'Sample rate' must be in list {list_sample_rate}")

            if audio_mode.bits not in list_bits:
                raise ValueError(f"'Bits' must be in list {list_bits}")

            if signal_frequency <= 0:
                raise ValueError(f"'Signal frequency' must be more than 0")

            if amplitude not in [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]:
                raise ValueError(f"'Amplitude' must be more than 0")

            memory_block_no = 2

            self.__memory_manager.make_default()
            self.__io.set(TSI_AUDGEN_SIGNAL_BLOCK, memory_block_no)
            self.__io.set(TSI_AUDGEN_SIGNAL_TYPE, audio_pattern.value)
            self.__io.set(TSI_AUDGEN_CHANNEL_COUNT, audio_mode.channel_count)
            self.__io.set(TSI_AUDGEN_SAMPLE_RATE, audio_mode.sample_rate)
            self.__io.set(TSI_AUDGEN_SAMPLE_SIZE, audio_mode.bits)
            self.__io.set(TSI_AUDGEN_SIGNAL_FREQ, signal_frequency)
            self.__io.set(TSI_AUDGEN_SIGNAL_VOLUME, amplitude)
        else:
            available_variants = '\n'.join([e.name for e in AudioPattern])
            raise ValueError(f"Incorrect value of audio pattern - {audio_pattern}\n"
                             f"Available audio patterns: {available_variants}")

    def apply(self) -> bool:
        """
        Apply settings and start generate audio.

        Returns:
            object of `bool` type - generation was enabled successfully or not.
        """
        self.__io.set(TSI_W_AUDGEN_CONTROL, 1)

        def is_apply_ag_success(ag: AudioGenerator):
            return ag.status == AGStatus.Running

        return function_scheduler(is_apply_ag_success, self, interval=1, timeout=10)

    def stop_generate(self) -> bool:
        """
        Stop generate audio.

        Returns:
            object of `bool` type - generation was disabled successfully or not.
        """
        return True if self.__io.set(TSI_W_AUDGEN_CONTROL, 0) >= TSI_SUCCESS else False

    @property
    def status(self) -> AGStatus:
        """
        Return audio generator status.

        Returns:
            object of `AGStatus` type
        """
        self.__status = AGStatus(self.__io.get(TSI_AUDGEN_STATUS_R, c_int)[1] & 0x1)
        return self.__status

    @property
    def audio_mode(self) -> AudioMode:
        """
        Return current audio mode.

        Returns:
            object of `AudioMode` type
        """
        self.__audio_mode.sample_rate = self.__io.get(TSI_AUDGEN_SAMPLE_RATE, c_int32)[1]
        self.__audio_mode.bits = self.__io.get(TSI_AUDGEN_SAMPLE_SIZE, c_int32)[1]
        self.__audio_mode.channel_count = self.__io.get(TSI_AUDGEN_CHANNEL_COUNT, c_int32)[1]
        return self.__audio_mode
