import pathlib
import time
import os
import math
from PIL import Image

import omni.kit.test_helpers_gfx.compare_utils

import carb.tokens
import carb.audio
import omni.audiorecorder
import omni.usd.audio
import omni.kit.audiodeviceenum
import omni.kit.test
import omni.log

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

# boilerplate to work around python lacking references
class IntReference:  # pragma: no cover
    def __init__(self, v):
        self._v = v;

    def _get_value(self):
        return self._v

    def _set_value(self, v):
        self._v = v

    value = property(_get_value, _set_value)

def wait_for_data(received: IntReference, ms: int, expected_data: int):  # pragma: no cover
    for i in range(ms):
        if received.value > expected_data:
            # log this so we can see timing margins on teamcity
            omni.log.warn("finished waiting after " + str(i / 1000.0) + " seconds")
            break
        if i > 0 and i % 1000 == 0:
            omni.log.warn("waited " + str(i // 1000) + " seconds")
        if i == ms - 1:
            self.assertTrue(not "timeout of " + str(ms) + "ms waiting for " + str(expected_data) + " frames")
        time.sleep(0.001)

class TestAudioRecorder(omni.kit.test.AsyncTestCase):  # pragma: no cover
    def setUp(self):
        self._recorder = omni.audiorecorder.create_audio_recorder()
        self.assertIsNotNone(self._recorder)

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


    def tearDown(self):
        self._recorder = None


    def test_callback_recording(self):
        valid = []
        for i in range(1, 16):
            valid.append(4800 * i)

        received = IntReference(0)
        def read_validation(data):
            nonlocal received
            received.value += len(data)
            self.assertTrue(len(data) in valid)

        def float_read_callback(data):
            read_validation(data)
            self.assertTrue(isinstance(data[0], float))

        def int_read_callback(data):
            read_validation(data)
            self.assertTrue(isinstance(data[0], int))

        # if there are no devices, this won't work
        if omni.kit.audiodeviceenum.acquire_audio_device_enum_interface().get_device_count(omni.kit.audiodeviceenum.Direction.CAPTURE) == 0:
            self.assertFalse(self._recorder.begin_recording_float(
                callback = float_read_callback,
            ))
            return

        # not open => should be some valid default format
        fmt = self._recorder.get_format()
        self.assertIsNotNone(fmt)
        self.assertGreaterEqual(fmt.channels, carb.audio.MIN_CHANNELS)
        self.assertLessEqual(fmt.channels, carb.audio.MAX_CHANNELS)
        self.assertGreaterEqual(fmt.frame_rate, carb.audio.MIN_FRAMERATE)
        self.assertLessEqual(fmt.frame_rate, carb.audio.MAX_FRAMERATE)

        self.assertTrue(self._recorder.begin_recording_float(
            callback = float_read_callback,
            buffer_length = 4800 * 16, # enormous buffer for test reliability
            period = 4800,
            length_type = carb.audio.UnitType.FRAMES,
            channels = 1
        ))

        # try again and it will fail because a recording is already in progress
        self.assertFalse(self._recorder.begin_recording_float(
            callback = float_read_callback,
            buffer_length = 4800 * 16, # enormous buffer for test reliability
            period = 4800,
            length_type = carb.audio.UnitType.FRAMES,
            channels = 1
        ))

        # not much we can test here aside from it being valid
        fmt = self._recorder.get_format()
        self.assertIsNotNone(fmt)
        self.assertEqual(fmt.channels, 1)
        self.assertGreaterEqual(fmt.frame_rate, carb.audio.MIN_FRAMERATE)
        self.assertLessEqual(fmt.frame_rate, carb.audio.MAX_FRAMERATE)
        self.assertEqual(fmt.format, carb.audio.SampleFormat.PCM_FLOAT)

        # this timeout seems absurd, but anything's possible with teamcity's timing
        wait_for_data(received, 8000, 4800)
        self._recorder.stop_recording()


        # try again with int16
        received.value = 0
        self.assertTrue(self._recorder.begin_recording_int16(
            callback = int_read_callback,
            buffer_length = 4800 * 16, # enormous buffer for test reliability
            period = 4800,
            length_type = carb.audio.UnitType.FRAMES,
            channels = 1
        ))

        # not much we can test here aside from it being valid
        fmt = self._recorder.get_format()
        self.assertIsNotNone(fmt)
        self.assertEqual(fmt.channels, 1)
        self.assertGreaterEqual(fmt.frame_rate, carb.audio.MIN_FRAMERATE)
        self.assertLessEqual(fmt.frame_rate, carb.audio.MAX_FRAMERATE)
        self.assertEqual(fmt.format, carb.audio.SampleFormat.PCM16)

        wait_for_data(received, 8000, 4800)
        self._recorder.stop_recording()


        # try again with int32
        received.value = 0
        self.assertTrue(self._recorder.begin_recording_int32(
            callback = int_read_callback,
            buffer_length = 4800 * 16, # enormous buffer for test reliability
            period = 4800,
            length_type = carb.audio.UnitType.FRAMES,
            channels = 1
        ))

        # not much we can test here aside from it being valid
        fmt = self._recorder.get_format()
        self.assertIsNotNone(fmt)
        self.assertEqual(fmt.channels, 1)
        self.assertGreaterEqual(fmt.frame_rate, carb.audio.MIN_FRAMERATE)
        self.assertLessEqual(fmt.frame_rate, carb.audio.MAX_FRAMERATE)
        self.assertEqual(fmt.format, carb.audio.SampleFormat.PCM32)

        wait_for_data(received, 8000, 4800)
        self._recorder.stop_recording()



        # test that the format is set as specified
        self.assertTrue(self._recorder.begin_recording_float(
            callback = float_read_callback,
            channels = 1, # 1 channel because that's all we can guarantee
                          # when we upgrade IAudioCapture, this should no longer
                          # be an issue
            frame_rate = 43217,
            buffer_length = 4800 * 16, # enormous buffer for test reliability
            period = 4800,
            length_type = carb.audio.UnitType.FRAMES
        ))

        fmt = self._recorder.get_format()
        self.assertEqual(fmt.channels, 1)
        self.assertEqual(fmt.frame_rate, 43217)
        self.assertEqual(fmt.format, carb.audio.SampleFormat.PCM_FLOAT)

        self._recorder.stop_recording()

    def _validate_sound(self, path):

        # to validate that this is a working sound, load it into a sound prim
        # and check if the sound asset loaded successfully
        context = omni.usd.get_context()
        self.assertIsNotNone(context)

        context.new_stage()
        stage = context.get_stage()
        self.assertIsNotNone(stage)

        prim = stage.DefinePrim("/sound", "OmniSound")
        self.assertIsNotNone(prim)
        prim.GetAttribute("filePath").Set(path)

        audio = omni.usd.audio.get_stage_audio_interface()
        self.assertIsNotNone(audio)


        i = 0
        while audio.get_sound_asset_status(prim) == omni.usd.audio.AssetLoadStatus.IN_PROGRESS:
            time.sleep(0.001)
            if i > 5000:
                raise Exception("asset load timed out")
            i += 1

        self.assertEqual(audio.get_sound_asset_status(prim), omni.usd.audio.AssetLoadStatus.DONE)


    def test_recording_to_file(self):

        received = IntReference(0)
        def read_callback(data):
            nonlocal received
            received.value += len(data)

        if omni.kit.audiodeviceenum.acquire_audio_device_enum_interface().get_device_count(omni.kit.audiodeviceenum.Direction.CAPTURE) == 0:
            self.assertFalse(self._recorder.begin_recording_float(
                callback = read_callback,
                filename = str(OUTPUTS_DIR.joinpath("test.oga"))
            ))
            self.assertFalse(self._recorder.begin_recording_to_file(
                filename = str(OUTPUTS_DIR.joinpath("test.oga"))
            ))
            return


        self.assertTrue(self._recorder.begin_recording_float(
            callback = read_callback,
            buffer_length = 4800 * 16, # enormous buffer for test reliability
            period = 4800,
            channels = 1,
            frame_rate = 48000,
            length_type = carb.audio.UnitType.FRAMES,
            output_format = carb.audio.SampleFormat.OPUS,
            filename = str(OUTPUTS_DIR.joinpath("test.opus.oga"))
        ))

        wait_for_data(received, 8000, 4800)
        self._recorder.stop_recording()

        self._validate_sound(str(OUTPUTS_DIR.joinpath("test.opus.oga")))

        # try again with a default output format
        self.assertTrue(self._recorder.begin_recording_float(
            callback = read_callback,
            buffer_length = 4800 * 16, # enormous buffer for test reliability
            period = 4800,
            channels = 1,
            frame_rate = 48000,
            length_type = carb.audio.UnitType.FRAMES,
            filename = str(OUTPUTS_DIR.joinpath("test.default.0.wav"))
        ))

        wait_for_data(received, 8000, 4800)
        self._recorder.stop_recording()

        self._validate_sound(str(OUTPUTS_DIR.joinpath("test.default.0.wav")))

        self.assertTrue(self._recorder.begin_recording_to_file(
            frame_rate = 73172,
            channels = 1,
            filename = str(OUTPUTS_DIR.joinpath("test.vorbis.oga")),
            output_format = carb.audio.SampleFormat.VORBIS
        ))

        time.sleep(2.0)

        self._recorder.stop_recording()

        self._validate_sound(str(OUTPUTS_DIR.joinpath("test.vorbis.oga")))

        # try again with a default output format
        self.assertTrue(self._recorder.begin_recording_to_file(
            frame_rate = 73172,
            channels = 1,
            filename = str(OUTPUTS_DIR.joinpath("test.default.1.wav")),
        ))

        time.sleep(2.0)

        self._recorder.stop_recording()

        self._validate_sound(str(OUTPUTS_DIR.joinpath("test.default.1.wav")))


    def test_bad_parameters(self):

        def read_callback(data):
            pass

        # MP3 is not supported
        self.assertFalse(self._recorder.begin_recording_float(
            callback = read_callback,
            output_format = carb.audio.SampleFormat.MP3,
            filename = str(OUTPUTS_DIR.joinpath("test.mp3"))
        ))

        # bad format
        self.assertFalse(self._recorder.begin_recording_float(
            callback = read_callback,
            output_format = carb.audio.SampleFormat.RAW,
            filename = str(OUTPUTS_DIR.joinpath("test.mp3"))
        ))

        # bad file name
        self.assertFalse(self._recorder.begin_recording_float(
            callback = read_callback,
            output_format = carb.audio.SampleFormat.OPUS,
            filename = "a/b/c/d/e/f/g/h/i/j.oga"
        ))

        self.assertFalse(self._recorder.begin_recording_float(
            callback = read_callback,
            channels = carb.audio.MAX_CHANNELS + 1
        ))

        self.assertFalse(self._recorder.begin_recording_float(
            callback = read_callback,
            frame_rate = carb.audio.MAX_FRAMERATE + 1
        ))

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

        samples_int16 = []
        samples_int32 = []
        samples_float = []
        for i in range(4800):
            samples_float.append(math.sin(i / 48.0))
            samples_int16.append(int(samples_float[-1] * (2 ** 15 - 1)))
            samples_int32.append(int(samples_float[-1] * (2 ** 31 - 1)))

        W = 256
        H = 256

        raw = omni.audiorecorder.draw_waveform_from_blob_float(samples_float, 1, 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(pathlib.Path(OUTPUTS_DIR).joinpath("waveform.float.png")))

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


        raw = omni.audiorecorder.draw_waveform_from_blob_int32(samples_int32, 1, 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(pathlib.Path(OUTPUTS_DIR).joinpath("waveform.int32.png")))

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


        raw = omni.audiorecorder.draw_waveform_from_blob_int16(samples_int16, 1, 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(pathlib.Path(OUTPUTS_DIR).joinpath("waveform.int16.png")))

        if not GENERATE_GOLDEN_IMAGES:
            # 1 pixel of difference was 756.0 difference
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath("waveform.int16.png"),
                    self._golden_path.joinpath("waveform.png"),
                    pathlib.Path(OUTPUTS_DIR).joinpath("waveform.int16.png.diff.png")),
                1024.0)
