# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#
import carb.audio
import omni.audiorecorder
import omni.kit.ui
import omni.ui
import threading
import time
import re
import asyncio
from typing import Callable
from omni.kit.window.filepicker import FilePickerDialog


class AudioRecorderWindowExtension(omni.ext.IExt):
    """Audio Recorder Window Extension"""

    class ComboModel(omni.ui.AbstractItemModel):
        class ComboItem(omni.ui.AbstractItem):
            def __init__(self, text):
                super().__init__()
                self.model = omni.ui.SimpleStringModel(text)

        def __init__(self):
            super().__init__()
            self._options = [
                ["16 bit PCM", carb.audio.SampleFormat.PCM16],
                ["24 bit PCM", carb.audio.SampleFormat.PCM24],
                ["32 bit PCM", carb.audio.SampleFormat.PCM32],
                ["float PCM", carb.audio.SampleFormat.PCM_FLOAT],
                ["Vorbis", carb.audio.SampleFormat.VORBIS],
                ["FLAC", carb.audio.SampleFormat.FLAC],
                ["Opus", carb.audio.SampleFormat.OPUS],
            ]

            self._current_index = omni.ui.SimpleIntModel()
            self._current_index.add_value_changed_fn(lambda a: self._item_changed(None))

            self._items = [AudioRecorderWindowExtension.ComboModel.ComboItem(text) for (text, value) in self._options]

        def get_item_children(self, item):
            return self._items

        def get_item_value_model(self, item, column_id):
            if item is None:
                return self._current_index
            return item.model

        def set_value(self, value):
            for i in range(0, len(self._options)):
                if self._options[i][1] == value:
                    self._current_index.as_int = i
                    break

        def get_value(self):
            return self._options[self._current_index.as_int][1]

    class FieldModel(omni.ui.AbstractValueModel):
        def __init__(self):
            super(AudioRecorderWindowExtension.FieldModel, self).__init__()
            self._value = ""

        def get_value_as_string(self):
            return self._value

        def begin_edit(self):
            pass

        def set_value(self, value):
            self._value = value
            self._value_changed()

        def end_edit(self):
            pass

        def get_value(self):
            return self._value

    def _choose_file_clicked(self):  # pragma: no cover
        dialog = FilePickerDialog(
            "Select File",
            apply_button_label="Select",
            click_apply_handler=lambda filename, dirname: self._on_file_pick(dialog, filename, dirname),
        )
        dialog.show()

    def _on_file_pick(self, dialog: FilePickerDialog, filename: str, dirname: str):  # pragma: no cover
        path = ""
        if dirname:
            path = f"{dirname}/{filename}"
        elif filename:
            path = filename

        dialog.hide()
        self._file_field.model.set_value(path)

    def _menu_callback(self, a, b):
        self._window.visible = not self._window.visible

    def _read_callback(self, data):  # pragma: no cover
        self._display_buffer[self._display_buffer_index] = data
        self._display_buffer_index = (self._display_buffer_index + 1) % self._display_len
        buf = []
        for i in range(self._display_len):
            buf += self._display_buffer[(self._display_buffer_index + i) % self._display_len]

        width = 512
        height = 128
        img = omni.audiorecorder.draw_waveform_from_blob_int16(
            input=buf,
            channels=1,
            width=width,
            height=height,
            fg_color=[0.89, 0.54, 0.14, 1.0],
            bg_color=[0.0, 0.0, 0.0, 0.0],
        )
        self._waveform_image_provider.set_bytes_data(img, [width, height])

    def _close_error_window(self):
        self._error_window.visible = False

    def _record_clicked(self):
        if self._recording:
            self._record_button.set_style({"image_url": "resources/glyphs/audio_record.svg"})
            self._recorder.stop_recording()
            self._recording = False
        else:
            result = self._recorder.begin_recording_int16(
                filename=self._file_field_model.get_value(),
                callback=self._read_callback,
                output_format=self._format_model.get_value(),
                buffer_length=200,
                period=25,
                length_type=carb.audio.UnitType.MILLISECONDS,
            )
            if result:
                self._record_button.set_style({"image_url": "resources/glyphs/timeline_stop.svg"})
                self._recording = True
            else:  # pragma: no cover
                self._error_window = omni.ui.Window(
                    "Audio Recorder Error", width=400, height=0, flags=omni.ui.WINDOW_FLAGS_NO_DOCKING
                )
                with self._error_window.frame:
                    with omni.ui.VStack():
                        with omni.ui.HStack():
                            omni.ui.Spacer()
                            self._error_window_label = omni.ui.Label(
                                "Failed to start recording. The file path may be incorrect or the device may be inaccessible.",
                                word_wrap=True,
                                width=380,
                                alignment=omni.ui.Alignment.CENTER,
                            )
                            omni.ui.Spacer()
                        with omni.ui.HStack():
                            omni.ui.Spacer()
                            self._error_window_ok_button = omni.ui.Button(
                                width=64, height=32, clicked_fn=self._close_error_window, text="ok"
                            )
                            omni.ui.Spacer()

    def _stop_clicked(self):
        pass

    def _create_tooltip(self, text):
        """Create a tooltip in a fixed style"""
        with omni.ui.VStack(width=400):
            omni.ui.Label(text, word_wrap=True)

    def on_startup(self):
        self._display_len = 8
        self._display_buffer = [[0] for i in range(self._display_len)]
        self._display_buffer_index = 0
        # self._ticker_pos = 0;
        self._recording = False
        self._recorder = omni.audiorecorder.create_audio_recorder()
        self._window = omni.ui.Window("Audio Recorder", width=600, height=240)
        with self._window.frame:
            with omni.ui.VStack(height=0, spacing=8):
                # file dialogue
                with omni.ui.HStack():
                    omni.ui.Button(
                        width=32,
                        height=32,
                        clicked_fn=self._choose_file_clicked,
                        style={"image_url": "resources/glyphs/folder.svg"},
                    )
                    self._file_field_model = AudioRecorderWindowExtension.FieldModel()
                    self._file_field = omni.ui.StringField(self._file_field_model, height=32)

                # waveform
                with omni.ui.HStack(height=128):
                    omni.ui.Spacer()
                    self._waveform_image_provider = omni.ui.ByteImageProvider()
                    self._waveform_image = omni.ui.ImageWithProvider(
                        self._waveform_image_provider,
                        width=omni.ui.Percent(95),
                        height=omni.ui.Percent(100),
                        fill_policy=omni.ui.IwpFillPolicy.IWP_STRETCH,
                    )
                    omni.ui.Spacer()

                # buttons
                with omni.ui.HStack():
                    with omni.ui.ZStack():
                        omni.ui.Spacer()
                        self._anim_label = omni.ui.Label("", alignment=omni.ui.Alignment.CENTER)

                    with omni.ui.VStack():
                        omni.ui.Spacer()
                        self._format_model = AudioRecorderWindowExtension.ComboModel()
                        self._format = omni.ui.ComboBox(
                            self._format_model,
                            height=0,
                            tooltip_fn=lambda: self._create_tooltip(
                                "The format for the output file."
                                + "The PCM formats will output as a WAVE file (.wav)."
                                + "FLAC will output as a FLAC file (.flac)."
                                + "Vorbis and Opus will output as an Ogg file (.ogg/.oga)."
                            ),
                        )
                        omni.ui.Spacer()

                    self._record_button = omni.ui.Button(
                        width=32,
                        height=32,
                        clicked_fn=self._record_clicked,
                        style={"image_url": "resources/glyphs/audio_record.svg"},
                    )

                    omni.ui.Spacer()

        # add a callback to open the window
        self._menuEntry = omni.kit.ui.get_editor_menu().add_item("Window/Audio Recorder", self._menu_callback)

        self._window.visible = False

    def on_shutdown(self):  # pragma: no cover
        self._recorder = None
        self._window = None
        self._menuEntry = None
