"""Unit tests for the timeline playback nodes in this extension"""

import carb.events
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.kit.commands
import omni.kit.stage_templates
import omni.kit.test
import omni.timeline
import omni.usd


# ======================================================================
class TestTimelineNodes(ogts.OmniGraphTestCase):
    TEST_GRAPH_PATH = "/World/TestGraph"

    FRAME_RATE = 64  # power of two so that FRAME_DURATION is an exact binary number
    FRAME_DURATION = 1.0 / FRAME_RATE
    START_FRAME = 10
    END_FRAME = 20
    START_TIME = START_FRAME * FRAME_DURATION
    END_TIME = END_FRAME * FRAME_DURATION

    TEST_MESSAGE_NAME = "omni.graph.nodes/triggerTimelineNode"
    TEST_MESSAGE_EVENT = carb.events.type_from_string(TEST_MESSAGE_NAME)

    def flush(self):
        # The timeline is async, we need to commit any pending changes before we can read its last state.
        self._timeline.commit()

    def current_time(self):
        self.flush()
        return self._timeline.get_current_time()

    def current_frame(self):
        return self._timeline.time_to_time_code(self.current_time())

    def assert_from_start(self, frame: int, message: str):
        self.assertEqual(self.current_frame(), self.START_FRAME + frame, message)

    def is_playing(self):
        self.flush()
        return self._timeline.is_playing()

    def is_looping(self):
        self.flush()
        return self._timeline.is_looping()

    def dump(self, header=""):
        self.flush()
        tl = self._timeline
        print(
            f"""
Timeline {header}
tick:{tl.get_current_tick()}
time:{tl.get_current_time()}
start:{tl.get_start_time()}
end:{tl.get_end_time()}
playing:{tl.is_playing()}
looping:{tl.is_looping()}
"""
        )

    async def tick(self):
        await omni.kit.app.get_app().next_update_async()

    async def fire(self):
        app = omni.kit.app.get_app()
        message_bus = app.get_message_bus_event_stream()
        message_bus.push(self.TEST_MESSAGE_EVENT)
        await self.tick()

    # ----------------------------------------------------------------------
    async def setUp(self):
        """Set up test environment, to be torn down when done"""
        await super().setUp()

        fps = self.FRAME_RATE

        # make sure we put the global timeline in the same initial state
        self._timeline = omni.timeline.get_timeline_interface()
        self._timeline.stop()
        self._timeline.set_current_time(0)
        self._timeline.set_end_time(self.END_TIME)
        self._timeline.set_start_time(self.START_TIME)
        self._timeline.set_time_codes_per_second(fps)
        self._timeline.set_looping(False)
        self._timeline.set_fast_mode(False)
        self._timeline.set_prerolling(True)
        self._timeline.set_auto_update(True)
        self._timeline.set_play_every_frame(True)
        self._timeline.set_ticks_per_frame(1)
        self._timeline.set_target_framerate(fps)
        self._timeline.commit()

    # ----------------------------------------------------------------------
    async def test_timeline_start_node(self):
        """Test OgnTimelineStart node"""
        # convert to action graph
        og.Controller.edit({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})

        controller = og.Controller()
        keys = og.Controller.Keys

        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnMessage", "omni.graph.action.OnMessageBusEvent"),
                    ("StartTimeline", "omni.graph.nodes.StartTimeline"),
                ],
                keys.CONNECT: [
                    ("OnMessage.outputs:execOut", "StartTimeline.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnMessage.inputs:onlyPlayback", False),
                    ("OnMessage.inputs:eventName", self.TEST_MESSAGE_NAME),
                ],
            },
        )

        # Initialize
        await self.tick()

        self.assertFalse(self.is_playing(), "Should not have started yet")

        # Trigger the node
        await self.fire()

        self.assertTrue(self.is_playing(), "Should have started")

        self.assert_from_start(0, "Should have started at frame 0 from start")

        await self.tick()

        self.assert_from_start(1, "Should have advanced one frame after starting")

        # Trigger the node again
        await self.fire()

        self.assert_from_start(2, "Should not have restarted after starting twice")

        # Pause the timeline
        self._timeline.pause()
        self._timeline.commit()

        await self.tick()
        self.assert_from_start(2, "Should have been paused at frame 2")

        # Start again
        await self.fire()
        self.assert_from_start(2, "Should be at frame 2 after paused and starting again")

        await self.tick()
        self.assert_from_start(3, "Should be at frame 3 after paused and playing for one frame")

    # ----------------------------------------------------------------------

    async def test_timeline_stop_node(self):
        """Test OgnTimelineStop node"""
        # convert to action graph
        og.Controller.edit({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})

        controller = og.Controller()
        keys = og.Controller.Keys

        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnMessage", "omni.graph.action.OnMessageBusEvent"),
                    ("StopTimeline", "omni.graph.nodes.StopTimeline"),
                ],
                keys.CONNECT: [
                    ("OnMessage.outputs:execOut", "StopTimeline.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnMessage.inputs:onlyPlayback", False),
                    ("OnMessage.inputs:eventName", self.TEST_MESSAGE_NAME),
                ],
            },
        )

        # Initialize
        await self.tick()

        self.assertFalse(self.is_playing(), "Should not have started")

        # Trigger the node
        await self.fire()

        await self.tick()

        # Stopping a stopped timeline has no effect
        self.assertTrue(self.current_frame() == 0, "Should not have started when stopping")

        # Start the timeline
        self._timeline.play()

        self.assertTrue(self.is_playing(), "Should have started")

        play_frames = 5

        # Advance the timeline N frames
        for _ in range(play_frames):
            await self.tick()

        self.assert_from_start(play_frames, f"Should have played {play_frames} frames")

        # Trigger the node
        await self.fire()

        stop_frame = play_frames + 1

        self.assertFalse(self.is_playing(), "Should have stopped")
        self.assert_from_start(stop_frame, f"Should have stopped at frame {stop_frame}")

        await self.tick()

        self.assert_from_start(
            stop_frame,
            f"Should still be at frame {stop_frame} after stopped one frame",
        )

        # Trigger the node again
        await self.fire()

        # Should still be at the current frame
        self.assert_from_start(
            stop_frame,
            f"Should still be at frame {stop_frame} after stopping twice",
        )

    # ----------------------------------------------------------------------

    async def test_timeline_loop_node(self):
        """Test OgnTimelineLoop node"""
        # convert to action graph
        og.Controller.edit({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "execution"})

        controller = og.Controller()
        keys = og.Controller.Keys

        controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("OnMessage", "omni.graph.action.OnMessageBusEvent"),
                    ("LoopTimeline", "omni.graph.nodes.LoopTimeline"),
                ],
                keys.CONNECT: [
                    ("OnMessage.outputs:execOut", "LoopTimeline.inputs:execIn"),
                ],
                keys.SET_VALUES: [
                    ("OnMessage.inputs:onlyPlayback", False),
                    ("OnMessage.inputs:eventName", self.TEST_MESSAGE_NAME),
                ],
            },
        )

        def set_looping(value):
            controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.SET_VALUES: [
                        ("LoopTimeline.inputs:loop", value),
                    ],
                },
            )

        # Initialize
        await self.tick()

        self.assertFalse(self.is_looping(), "Should not be looping yet")

        # Enable looping
        set_looping(True)
        await self.fire()

        self.assertTrue(self.is_looping(), "Should be looping")

        # Disable looping
        set_looping(False)
        await self.fire()

        self.assertFalse(self.is_looping(), "Should not be looping")

    # ----------------------------------------------------------------------

    async def test_timeline_get_node(self):
        """Test OgnTimelineGet node"""
        # convert to push graph
        og.Controller.edit({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "push"})

        controller = og.Controller()
        keys = og.Controller.Keys

        (_, [node], _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("GetTimeline", "omni.graph.nodes.GetTimeline"),
                ],
            },
        )

        def get_attr_value(name):
            attr = og.Controller.attribute(f"outputs:{name}", node)
            return og.Controller(attr).get()

        def assert_attr_value(name, value):
            self.assertEqual(get_attr_value(name), value, name)

        async def assert_node_values(playing, looping):
            await self.tick()
            assert_attr_value("time", self._timeline.get_current_time())
            assert_attr_value("frame", self.current_frame())
            assert_attr_value("startTime", self.START_TIME)
            assert_attr_value("startFrame", self.START_FRAME)
            assert_attr_value("endTime", self.END_TIME)
            assert_attr_value("endFrame", self.END_FRAME)
            assert_attr_value("framesPerSecond", self.FRAME_RATE)
            assert_attr_value("isPlaying", playing)
            assert_attr_value("isLooping", looping)

        await assert_node_values(playing=False, looping=False)

        self._timeline.pause()

        await assert_node_values(playing=False, looping=False)

        self._timeline.play()

        await assert_node_values(playing=True, looping=False)
        await assert_node_values(playing=True, looping=False)

        self._timeline.set_looping(True)

        await assert_node_values(playing=True, looping=True)
        await assert_node_values(playing=True, looping=True)

    # ----------------------------------------------------------------------

    async def test_timeline_set_node(self):
        """Test OgnTimelineSet node"""
        # convert to push graph
        og.Controller.edit({"graph_path": self.TEST_GRAPH_PATH, "evaluator_name": "push"})

        controller = og.Controller()
        keys = og.Controller.Keys

        (_, [node], _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("SetTimeline", "omni.graph.nodes.SetTimeline"),
                ],
            },
        )

        def get_attr_value(name):
            attr = og.Controller.attribute(f"outputs:{name}", node)
            return og.Controller(attr).get()

        async def set_node_property(property_name, property_value):
            attr = og.Controller.attribute("inputs:propName", node)
            attr.set(property_name)
            attr = og.Controller.attribute("inputs:propValue", node)
            attr.set(property_value)
            await self.tick()
            self.flush()

        async def assert_set_position(target_frame, expected_frame, clamped):
            await set_node_property("Frame", target_frame)
            self.assertEqual(self.current_frame(), expected_frame, "Wrong Frame")
            self.assertEqual(get_attr_value("clamped"), clamped, "Wrong clamping")

            target_time = target_frame * self.FRAME_DURATION
            expected_time = expected_frame * self.FRAME_DURATION
            await set_node_property("Time", target_time)
            self.assertEqual(self.current_time(), expected_time, "Wrong Time")
            self.assertEqual(get_attr_value("clamped"), clamped, "Wrong clamping")

        for f in range(self.START_FRAME, self.END_FRAME):
            await assert_set_position(f, f, False)

        await assert_set_position(self.START_FRAME - 1, self.START_FRAME, True)
        await assert_set_position(self.END_FRAME + 1, self.END_FRAME, True)

        await set_node_property("StartTime", 0)
        self.assertEqual(self._timeline.get_start_time(), 0, "Wrong StartTime")

        await set_node_property("StartFrame", 1)
        self.assertEqual(self._timeline.get_start_time(), self.FRAME_DURATION, "Wrong StartFrame")

        await set_node_property("EndTime", 10)
        self.assertEqual(self._timeline.get_end_time(), 10, "Wrong EndTime")

        await set_node_property("EndFrame", 2)
        self.assertEqual(self._timeline.get_end_time(), 2 * self.FRAME_DURATION, "Wrong EndFrame")

        await set_node_property("FramesPerSecond", 30)
        self.assertEqual(self._timeline.get_time_codes_per_seconds(), 30, "Wrong FramesPerSecond")

    # ----------------------------------------------------------------------
