import pathlib
import time
import pathlib
from PIL import Image

import carb.tokens
import omni.audioplayer
import os

import omni.kit.test
import omni.kit.test_helpers_gfx.compare_utils

OUTPUTS_DIR = pathlib.Path(omni.kit.test.get_test_output_path())


# This test needs to come first alphabetically since the bug this tests didn't
# trigger if the plugin was already loaded and carb plugins don't ever unload
# during normal kit operations.
class Test0AudioPlayerHang(omni.kit.test.AsyncTestCase):  # pragma: no cover
    def setUp(self):
        extension_path = carb.tokens.get_tokens_interface().resolve("${omni.audioplayer}")
        self._test_path = pathlib.Path(extension_path).joinpath("data").joinpath("tests").absolute()

        self._player = omni.audioplayer.create_audio_player()
        self.assertIsNotNone(self._player)

    def tearDown(self):
        # if you don't do this, the player never gets destroyed until plugin unload
        self._player = None

    def test_unload_during_load(self):
        self._player.load_sound(str(self._test_path.joinpath("long.oga")))

        # Finish the test so the player will close and the plugin will unload.
        # At one point, this caused a GIL deadlock in ~AudioPlayer().
        # This test should verify that this deadlock no longer happens.

class TestAudioPlayer(omni.kit.test.AsyncTestCase):  # pragma: no cover
    def setUp(self):
        self._player = omni.audioplayer.create_audio_player()
        self.assertIsNotNone(self._player)

        extension_path = carb.tokens.get_tokens_interface().resolve("${omni.audioplayer}")
        self._test_path = pathlib.Path(extension_path).joinpath("data").joinpath("tests").absolute()
        self._golden_path = self._test_path.joinpath("golden")

    def tearDown(self):
        self._player = None

    def test_basic(self):
        events = []
        expected_load_result = True
        def event_callback(event):
            nonlocal events
            nonlocal expected_load_result
            events.append(event.type)
            if event.type == int(omni.audioplayer.CallbackType.LOADED):
                self.assertEqual(event.payload["success"], expected_load_result)

        sub = self._player.get_event_stream().create_subscription_to_pop(event_callback)

        # load the sound and verify that a loaded event shows up
        last = len(events)
        self._player.load_sound(str(self._test_path.joinpath("long.oga")))

        i = 0
        while len(events) == last:
            time.sleep(0.001)
            if i > 30000:
                raise Exception("hang while loading a test asset")
            i += 1

        self.assertEqual(len(events), last + 1)
        self.assertEqual(events[-1], int(omni.audioplayer.CallbackType.LOADED))

        # check that get_sound_length() works
        self.assertLess(abs(self._player.get_sound_length() - 120.0), 0.001)

        # play the sound and verify that a load event is sent
        last = len(events)
        self._player.play_sound(str(self._test_path.joinpath("long.oga")))

        # it should load quickly since it's already been loaded once
        time.sleep(1)
        self.assertEqual(len(events), last + 1)
        self.assertEqual(events[-1], int(omni.audioplayer.CallbackType.LOADED))

        # pause and unpause it a few times
        last = len(events)
        self._player.pause_sound()
        time.sleep(1)
        self._player.unpause_sound()
        time.sleep(1)
        self._player.pause_sound()
        time.sleep(1)

        # nothing should have been sent
        self.assertEqual(len(events), last)

        self.assertGreater(self._player.get_play_cursor(), 0.0)

        # modify the play cursor, this should send an ended then loaded event
        last = len(events)
        self._player.set_play_cursor(60.0)

        i = 0
        while len(events) <= last + 1:
            time.sleep(0.001)
            if i > 30000:
                raise Exception("hang while loading a test asset")
            i += 1

        self.assertEqual(len(events), last + 2)
        self.assertEqual(events[-2], int(omni.audioplayer.CallbackType.ENDED))
        self.assertEqual(events[-1], int(omni.audioplayer.CallbackType.LOADED))

        self.assertGreaterEqual(self._player.get_play_cursor(), 60.0)

        # stop it and verify that the ended event is sent
        last = len(events)
        self._player.stop_sound()

        i = 0
        while len(events) == last:
            time.sleep(0.001)
            if i > 30000:
                raise Exception("hang while loading a test asset")
            i += 1

        self.assertEqual(len(events), last + 1)
        self.assertEqual(events[-1], int(omni.audioplayer.CallbackType.ENDED))

        # play the same sound again and see if events are sent
        last = len(events)
        self._player.play_sound(str(self._test_path.joinpath("long.oga")))

        i = 0
        while len(events) == last:
            time.sleep(0.001)
            if i > 30000:
                raise Exception("hang while loading a test asset")
            i += 1

        self.assertEqual(len(events), last + 1)
        self.assertEqual(events[-1], int(omni.audioplayer.CallbackType.LOADED))

        # play again to see the ended and loaded event
        last = len(events)
        self._player.play_sound(str(self._test_path.joinpath("long.oga")))

        i = 0
        while len(events) <= last + 1:
            time.sleep(0.001)
            if i > 30000:
                raise Exception("hang while loading a test asset")
            i += 1

        self.assertEqual(len(events), last + 2)
        self.assertEqual(events[-2], int(omni.audioplayer.CallbackType.ENDED))
        self.assertEqual(events[-1], int(omni.audioplayer.CallbackType.LOADED))

        # Note: there is no test for callbacks with play_sound() 2 times back to back
        # because whether callbacks for both calls are sent is dependent on timing.
        # The less recent sound playback may just be cancelled.


    def test_draw_waveform(self):
        # toggle in case you want to regenerate waveforms
        GENERATE_GOLDEN_IMAGES = False

        W = 256
        H = 256

        loaded = False
        def event_callback(event):
            if event.type == int(omni.audioplayer.CallbackType.LOADED):
                nonlocal loaded
                loaded = True

        sub = self._player.get_event_stream().create_subscription_to_pop(event_callback)

        self._player.load_sound(str(self._test_path.joinpath("short-2ch.oga")))

        i = 0
        while not loaded:
            time.sleep(0.001)
            if i > 30000:
                raise Exception("hang while loading a test asset")
            i += 1

        self._player.draw_waveform_to_file(str(OUTPUTS_DIR.joinpath("waveform.png")), W, H, [1.0, 1.0, 1.0, 1.0], [0.0, 0.0, 0.0, 1.0])

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath("waveform.png"),
                    self._golden_path.joinpath("waveform.png"),
                    pathlib.Path(OUTPUTS_DIR).joinpath("waveform.png.diff.png")),
                0.1)

        raw = self._player.draw_waveform(W, H, [1.0, 1.0, 1.0, 1.0], [0.0, 0.0, 0.0, 1.0])
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(OUTPUTS_DIR.joinpath("waveform.raw.png")))

        if not GENERATE_GOLDEN_IMAGES:
            # this image is slightly different than the original and this compare
            # helper function cannot handle even a slight difference, so we need to
            # have a separate golden image for this
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath("waveform.raw.png"),
                    self._golden_path.joinpath("waveform.raw.png"),
                    pathlib.Path(OUTPUTS_DIR).joinpath("waveform.raw.png.diff.png")),
                0.1)

    def _load_and_play_data_in_memory(self, filename, frame_rate, channels, format, flags = 0, expected_length = 2.0):
        events = []
        expected_load_result = True
        def event_callback(event):
            nonlocal events
            nonlocal expected_load_result
            events.append(event.type)
            if event.type == int(omni.audioplayer.CallbackType.LOADED):
                self.assertEqual(event.payload["success"], expected_load_result)

        # subscribe to events from the player.
        sub = self._player.get_event_stream().create_subscription_to_pop(event_callback)

        # load the test file into a bytes object.
        file_data = pathlib.Path(self._test_path.joinpath(filename)).read_bytes()

        # load the sound and verify that a 'loaded' event shows up.
        last = len(events)
        self._player.load_sound_in_memory(filename, file_data, frame_rate, channels, format, flags)

        i = 0
        while len(events) == last:
            time.sleep(0.001)
            if i > 30000:
                raise Exception("hang while loading a test asset")
            i += 1

        self.assertEqual(len(events), last + 1)
        self.assertEqual(events[-1], int(omni.audioplayer.CallbackType.LOADED))

        # check that get_sound_length() works.
        self.assertLess(abs(self._player.get_sound_length() - expected_length), 0.001)

        # play the sound and verify that a load event is sent.
        last = len(events)
        self._player.play_sound_in_memory(filename, file_data, frame_rate, channels, format, 0.0, flags)

        # it should load quickly since it's already been loaded once.
        time.sleep(1)
        self.assertEqual(len(events), last + 1)
        self.assertEqual(events[-1], int(omni.audioplayer.CallbackType.LOADED))

        # stop it and verify that the ended event is sent.
        last = len(events)
        self._player.stop_sound()

        i = 0
        while len(events) == last:
            time.sleep(0.001)
            if i > 30000:
                raise Exception("hang while loading a test asset")
            i += 1

        self.assertEqual(len(events), last + 1)
        self.assertEqual(events[-1], int(omni.audioplayer.CallbackType.ENDED))

        # clean up.
        sub = None

    def test_data_in_memory(self):
        self._load_and_play_data_in_memory("tone-440-8.wav", 44100, 1, omni.audioplayer.RawPcmFormat.PCM_8)
        self._load_and_play_data_in_memory("tone-440-16.wav", 44100, 1, omni.audioplayer.RawPcmFormat.PCM_16)
        self._load_and_play_data_in_memory("tone-440-32.wav", 44100, 1, omni.audioplayer.RawPcmFormat.PCM_32)
        self._load_and_play_data_in_memory("tone-440-float.wav", 44100, 1, omni.audioplayer.RawPcmFormat.PCM_FLOAT)

    def test_raw_data_in_memory(self):
        self._load_and_play_data_in_memory("tone-440-8.raw", 44100, 1, omni.audioplayer.RawPcmFormat.PCM_8, omni.audioplayer.FLAG_RAW_DATA)
        self._load_and_play_data_in_memory("tone-440-16.raw", 44100, 1, omni.audioplayer.RawPcmFormat.PCM_16, omni.audioplayer.FLAG_RAW_DATA)
        self._load_and_play_data_in_memory("tone-440-32.raw", 44100, 1, omni.audioplayer.RawPcmFormat.PCM_32, omni.audioplayer.FLAG_RAW_DATA)
        self._load_and_play_data_in_memory("tone-440-float.raw", 44100, 1, omni.audioplayer.RawPcmFormat.PCM_FLOAT, omni.audioplayer.FLAG_RAW_DATA)
