# 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.settings
import carb.dictionary
import omni.audioplayer
import omni.kit.ui
import omni.ui
import threading
import time
import re
import asyncio
import enum
from typing import Callable
from omni.kit.window.filepicker import FilePickerDialog

PERSISTENT_SETTINGS_PREFIX = "/persistent"

class EndReason(enum.Enum):

    # sound finished naturally
    FINISHED = 0,

    # sound was explicitly stopped
    STOPPED = 1,

    # seeked to a new location in the sound (causes an end callback)
    SEEK = 2,

    # the previous sound ended because another one is being played
    QUEUED_NEW_SOUND = 3,


class AudioPlayerWindowExtension(omni.ext.IExt):
    """Audio Player Window Extension"""

    class FieldModel(omni.ui.AbstractValueModel):
        def __init__(self, end_edit_callback):
            super(AudioPlayerWindowExtension.FieldModel, self).__init__()
            self._end_edit_callback = end_edit_callback
            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):
            self._end_edit_callback(self._value)

    class SliderModel(omni.ui.AbstractValueModel):
        def __init__(self, update_callback, end_edit_callback):
            super(AudioPlayerWindowExtension.SliderModel, self).__init__()
            self._update_callback = update_callback
            self._end_edit_callback = end_edit_callback
            self._value = 0

        def get_value_as_int(self):  # pragma: no cover
            return int(self._value)

        def get_value_as_float(self):  # pragma: no cover
            return float(self._value)

        def begin_edit(self):  # pragma: no cover
            pass

        def set_value(self, value):  # pragma: no cover
            self._value = value
            self._value_changed()
            self._update_callback(self._value)

        def end_edit(self):  # pragma: no cover
            self._end_edit_callback(self._value)

    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)

        # this has to be called manually because set_value doesn't do it
        self._file_field_end_edit(path)

    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 _set_pause_button(self):  # pragma: no cover
        self._play_button.set_style({"image_url": "resources/glyphs/timeline_pause.svg"})

    def _set_play_button(self):  # pragma: no cover
        self._play_button.set_style({"image_url": "resources/glyphs/timeline_play.svg"})

    def _timeline_str(self, time):  # pragma: no cover
        sec = ":{:02.0f}".format(time % 60)
        if time > 60.0 * 60.0:
            return "{:1.0f}".format(time // (60 * 60)) + ":{:02.0f}".format((time // 60) % 60) + sec
        else:
            return "{:1.0f}".format(time // 60) + sec

    def _timeline_ticker(self):  # pragma: no cover
        if not self._playing:
            return

        time = self._player.get_play_cursor()
        self._timeline_cursor_label.text = self._timeline_str(time)
        self._timeline_slider.model.set_value(time * self._timeline_slider_scale)

        # if the window was closed, stop the player
        if not self._window.visible:
            self._end_reason = EndReason.STOPPED
            self._player.stop_sound()

        self._ticker = threading.Timer(0.25, self._timeline_ticker).start()

    def _loading_ticker(self):
        labels = {0: "Loading", 1: "Loading.", 2: "Loading..", 3: "Loading..."}

        if not self._loading:
            self._loading_label.text = ""
            return

        self._loading_label.text = labels[self._loading_counter % 4]
        self._loading_counter += 1
        self._loading_timer = threading.Timer(0.25, self._loading_ticker).start()

    def _play_sound(self, time):
        self._loading = True
        self._player.play_sound(
            self._file_field.model.get_value_as_string(), time
        )

    def _close_error_window(self):  # pragma: no cover
        self._error_window.visible = False

    def _set_play_cursor(self, time):  # pragma: no cover
        self._end_reason = EndReason.SEEK
        self._player.set_play_cursor(time)

    def _file_loaded(self, success):  # pragma: no cover
        self._loading = False
        if not success:
            self._playing = False
            self._set_play_button()
            error_text = "Loading failed"
            file_name = self._file_field.model.get_value_as_string()
            if re.search("^.*.(m4a|aac)$", file_name):
                error_text = (
                    f"Failed to load file '{file_name}' codec not supported - only Vorbis, FLAC and WAVE are supported"
                )
            else:
                error_text = f"Failed to load file '{file_name}' codec not supported (only Vorbis, FLAC, MP3 and WAVE are supported), file does not exist or the file is corrupted"

            self._error_window = omni.ui.Window(
                "Audio Player 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(
                            error_text, 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()

            self._waveform_image_provider.set_bytes_data([0, 0, 0, 0], [1, 1])
            return

        if self._new_file:
            width = 2048
            height = 64
            raw_image = self._player.draw_waveform(width, height, [0.89, 0.54, 0.14, 1.0], [0.0, 0.0, 0.0, 0.0])
            self._waveform_image_provider.set_bytes_data(raw_image, [width, height])
            self._new_file = False

        self._timeline_end_label.text = self._timeline_str(self._player.get_sound_length())
        self._sound_length = self._player.get_sound_length()
        self._timeline_slider_scale = 1.0 / self._sound_length

        # set the timeline ticker going
        self._timeline_ticker()

        # set this back to default
        self._end_reason = EndReason.FINISHED


    def _play_finished(self):  # pragma: no cover
        if self._end_reason != EndReason.SEEK and self._end_reason != EndReason.QUEUED_NEW_SOUND:
            self._playing = False

            # set the slider to finished
            self._timeline_cursor_label.text = self._timeline_str(0)
            self._timeline_slider.model.set_value(0.0)

        if self._end_reason == EndReason.FINISHED or self._end_reason == EndReason.STOPPED:
            self._set_play_button()

        if self._end_reason == EndReason.FINISHED and self._settings.get_as_bool(PERSISTENT_SETTINGS_PREFIX + "/audio/context/closeAudioPlayerOnStop"):
            self._window.visible = False

    def _play_clicked(self):  # pragma: no cover
        if self._loading:
            return

        if self._playing:
            if self._paused:
                self._player.unpause_sound()
                self._set_pause_button()
                self._paused = False
            else:
                self._player.pause_sound()
                self._set_play_button()
                self._paused = True
            return

        self._playing = True
        self._paused = False
        self._load_result_label.text = ""
        self._loading_ticker()
        self._set_pause_button()
        self._play_sound(self._timeline_slider.model.get_value_as_float() * self._sound_length)

    def _file_field_end_edit(self, value):
        self._loading = True
        self._new_file = True
        self._load_result_label.text = ""
        self._loading_ticker()
        self._stop_clicked()
        self._player.load_sound(self._file_field.model.get_value_as_string())

    def _stop_clicked(self):  # pragma: no cover
        if self._loading:
            return

        self._end_reason = EndReason.STOPPED
        self._player.stop_sound()
        self._playing = False
        self._paused = False

    def _slider_end_edit(self, value):  # pragma: no cover
        if self._loading:
            return

        if not self._playing:
            return

        self._set_play_cursor(value * self._sound_length)

    def _slider_changed(self, value):  # pragma: no cover
        if not self._playing and not self._loading:
            self._timeline_cursor_label.text = self._timeline_str(value * self._sound_length)

    def open_window(self):
        """
        Make the window become visible

        Args:
            No arguments

        Returns:
            No return value
        """
        self._window.visible = True

    def open_window_and_play(self, path):  # pragma: no cover
        """
        Make the window become visible then begin playing a file

        Args:
            path: The file to begin playing

        Returns:
            No return value
        """

        self._playing = True
        self._loading = True;
        self._paused = False
        self._new_file = True
        self._window.visible = True
        self._load_result_label.text = ""
        self._loading_ticker()
        self._set_pause_button()
        self._end_reason = EndReason.QUEUED_NEW_SOUND
        self._file_field.model.set_value(path)
        self._play_sound(0.0)

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

    def _on_menu_click(self, menu, value):  # pragma: no cover
        if self._content_window is None:
            return

        protocol = self._content_window.get_selected_icon_protocol()
        path = self._content_window.get_selected_icon_path()
        if not path.startswith(protocol):
            path = protocol + path

        self.open_window_and_play(path)

    def _on_menu_check(self, url):
        return not not re.search("^.*\\.(wav|wave|ogg|oga|flac|fla|mp3|m4a|spx|opus|adpcm)$", url)

    def _on_browser_click(self, menu, value):  # pragma: no cover
        if self._content_browser is None:
            return

        # protocol = self._content_browser.get_selected_icon_protocol()
        # path = self._content_browser.get_selected_icon_path()
        # if not path.startswith(protocol):
        #    path = protocol + path

        self.open_window_and_play(value)

    def _on_content_browser_load(self):  # pragma: no cover
        import omni.kit.window.content_browser

        self._content_browser = omni.kit.window.content_browser.get_content_window()
        if self._content_browser is not None:
            self._content_browser_entry = self._content_browser.add_context_menu(
                "Play Audio", "audio_play.svg", self._on_browser_click, self._on_menu_check
            )

    def _on_content_browser_unload(self):  # pragma: no cover
        if self._content_browser is not None:
            self._content_browser.delete_context_menu("Play Audio")

        self._content_browser_entry = None
        self._content_browser = None

    def _on_player_event(self, event):
        if event.type == int(omni.audioplayer.CallbackType.LOADED):
            success = event.payload["success"]
            self._file_loaded(success)
        elif event.type == int(omni.audioplayer.CallbackType.ENDED):
            self._play_finished()
        else:
            print("unrecognized type " + str(event.type))

    def on_startup(self):
        self._content_browser = None
        self._hooks = []
        manager = omni.kit.app.get_app().get_extension_manager()
        # current content window
        self._hooks.append(
            manager.subscribe_to_extension_enable(
                on_enable_fn=lambda _: self._on_content_browser_load(),
                on_disable_fn=lambda _: self._on_content_browser_unload(),
                ext_name="omni.kit.window.content_browser",
                hook_name="omni.kit.window.audioplayer omni.kit.window.content_browser listener",
            )
        )

        self._loading_counter = 0
        self._ticker = None
        self._loading = False
        self._end_reason = EndReason.FINISHED
        self._new_file = True
        self._sound_length = 0
        self._timeline_slider_scale = 0
        self._file = ""
        self._playing = False
        self._paused = False
        self._player = omni.audioplayer.create_audio_player()
        self._sub = self._player.get_event_stream().create_subscription_to_pop(self._on_player_event)
        self._window = omni.ui.Window("Audio Player", width=600, height=200)
        self._settings = carb.settings.get_settings()
        self._settings.set_default_bool(PERSISTENT_SETTINGS_PREFIX + "/audio/context/closeAudioPlayerOnStop", False)
        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 = AudioPlayerWindowExtension.FieldModel(self._file_field_end_edit)
                    self._file_field = omni.ui.StringField(self._file_field_model, height=32)

                # timeline slider
                with omni.ui.HStack(height=64):
                    self._timeline_cursor_label = omni.ui.Label("0:00", width=25)
                    omni.ui.Label(" / ", width=10)
                    self._timeline_end_label = omni.ui.Label("0:00", width=25)
                    self._timeline_slider_model = AudioPlayerWindowExtension.SliderModel(
                        self._slider_changed, self._slider_end_edit
                    )
                    with omni.ui.ZStack():
                        self._waveform_image_provider = omni.ui.ByteImageProvider()
                        self._waveform_image = omni.ui.ImageWithProvider(
                            self._waveform_image_provider,
                            width=omni.ui.Percent(100),
                            height=omni.ui.Percent(100),
                            fill_policy=omni.ui.IwpFillPolicy.IWP_STRETCH,
                        )
                        with omni.ui.VStack():
                            omni.ui.Spacer()
                            self._timeline_slider = omni.ui.FloatSlider(
                                self._timeline_slider_model,
                                height=0,
                                style={
                                    "color": 0x00FFFFFF,
                                    "background_color": 0x00000000,
                                    "draw_mode": omni.ui.SliderDrawMode.HANDLE,
                                    "font_size": 22,
                                },
                            )
                            omni.ui.Spacer()

                # buttons
                with omni.ui.HStack():
                    with omni.ui.ZStack():
                        omni.ui.Spacer()
                        self._load_result_label = omni.ui.Label(
                            "", alignment=omni.ui.Alignment.CENTER, style={"color": 0xFF0000FF}
                        )

                    self._play_button = omni.ui.Button(
                        width=32,
                        height=32,
                        clicked_fn=self._play_clicked,
                        style={"image_url": "resources/glyphs/timeline_play.svg"},
                    )
                    omni.ui.Button(
                        width=32,
                        height=32,
                        clicked_fn=self._stop_clicked,
                        style={"image_url": "resources/glyphs/timeline_stop.svg"},
                    )

                    with omni.ui.ZStack():
                        omni.ui.Spacer()
                        self._loading_label = omni.ui.Label("", alignment=omni.ui.Alignment.CENTER)

                with omni.ui.HStack(alignment=omni.ui.Alignment.LEFT, width=100):
                    omni.ui.Label("Close on Stop", alignment=omni.ui.Alignment.LEFT)
                    omni.ui.Spacer()
                    self._auto_close_on_stop = omni.ui.CheckBox(alignment=omni.ui.Alignment.LEFT)
                    omni.ui.Spacer()

        self._auto_close_on_stop.model.set_value(
            self._settings.get_as_bool(PERSISTENT_SETTINGS_PREFIX + "/audio/context/closeAudioPlayerOnStop")
        )
        self._dict = carb.dictionary.get_dictionary()
        self._auto_close_on_stop.model.add_value_changed_fn(
            lambda a, b=self._settings: b.set_bool(
                PERSISTENT_SETTINGS_PREFIX + "/audio/context/closeAudioPlayerOnStop", a.get_value_as_bool()
            )
        )

        def on_change(item, event_type):  # pragma: no cover
            self._auto_close_on_stop.model.set_value(self._dict.get(item))

        self._subscription = self._settings.subscribe_to_node_change_events(
            PERSISTENT_SETTINGS_PREFIX + "/audio/context/closeAudioPlayerOnStop", on_change
        )
        # add a callback to open the window
        # FIXME: disabled until the bugs are worked out
        self._menuEntry = omni.kit.ui.get_editor_menu().add_item("Window/Audio Player", self._menu_callback)

        self._window.visible = False

    def on_shutdown(self):  # pragma: no cover
        self._end_reason = EndReason.STOPPED
        self._player.stop_sound()

        if self._ticker != None:
            self._ticker.cancel()

        self._settings.unsubscribe_to_change_events(self._subscription)
        self._subscription = None

        # run the unload function to avoid breaking the extension when it reloads
        self._on_content_browser_unload()

        # remove the subscription before the player to avoid events with a dead player
        self._sub = None
        self._player = None
        self._window = None
        self._menuEntry = None
        self._content_window_entry = None
