# Copyright (c) 2023, 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 numpy as np
from omni.kit.test import AsyncTestCase
from omni.ujitso import *

kTokenTypeUSDPath = KeyToken(11)
kTokenTypeRTXPackedMeshLODs = KeyToken(12)
kTokenTypeRTXPackedMeshLOD = KeyToken(13)
kTokenTypeMeshLOD = KeyToken(14)
kTokenTypeTriangulatedMesh = KeyToken(15)
kTokenTypeMesh = KeyToken(16)
kTokenMeshReduceFactor = KeyToken(17)

class ProcessorImpl(Processor):
    def __init__(self, agent, key_token, get_info_impl, match_impl, gather_dependencies_impl, build_impl):
        super().__init__(get_info_impl, match_impl, gather_dependencies_impl, build_impl)

        self._agent = agent

        request_filter = RequestFilter([(key_token.value,)])
        agent.registry.register_processor(self, request_filter)

    def __del__(self):
        self._agent.registry.unregister_processor(self)

    @staticmethod
    def _match_impl(match_context, request_array):
        return OperationResult.SUCCESS, MatchResult.NORMAL_PRIORITY

    @staticmethod
    def _build_impl(build_context, build_job, build_handle):
        return OperationResult.SUCCESS

class ProcessorRTXPackedMeshLODs(ProcessorImpl):
    def __init__(self, agent, key_token):
        super().__init__(agent, key_token, ProcessorRTXPackedMeshLODs._get_info_impl, ProcessorRTXPackedMeshLODs._match_impl, ProcessorRTXPackedMeshLODs._gather_dependencies_impl, ProcessorRTXPackedMeshLODs._build_impl)

    @staticmethod
    def _get_info_impl(processor):
        return OperationResult.SUCCESS, ProcessorInformation("ProcessorRTXPackedMeshLODs", 1, 1)

    @staticmethod
    def _gather_dependencies_impl(dependency_context, dependency_job, dependency_handle):
        dynamic_request = DynamicRequest(dependency_job.request)
        dynamic_request.replace_key(KeyTokenEx(kTokenTypeRTXPackedMeshLODs), KeyTokenEx(kTokenTypeRTXPackedMeshLOD))

        dynamic_request.add_float(KeyTokenEx(kTokenMeshReduceFactor), 1.0)
        dependency_context.agent.service.add_dependency(dependency_context.agent, dependency_handle, dynamic_request.get_request())

        dynamic_request.replace_value_float(KeyTokenEx(kTokenMeshReduceFactor), 0.3)
        dependency_context.agent.service.add_dependency(dependency_context.agent, dependency_handle, dynamic_request.get_request())

        dynamic_request.replace_value_float(KeyTokenEx(kTokenMeshReduceFactor), 0.1)
        dependency_context.agent.service.add_dependency(dependency_context.agent, dependency_handle, dynamic_request.get_request())

        return OperationResult.SUCCESS

class ProcessorRTXPackedMeshLOD(ProcessorImpl):
    def __init__(self, agent, key_token):
        super().__init__(agent, key_token, ProcessorRTXPackedMeshLOD._get_info_impl, ProcessorRTXPackedMeshLOD._match_impl, ProcessorRTXPackedMeshLOD._gather_dependencies_impl, ProcessorRTXPackedMeshLOD._build_impl)

    @staticmethod
    def _get_info_impl(processor):
        return OperationResult.SUCCESS, ProcessorInformation("ProcessorRTXPackedMeshLOD", 1, 1)

    @staticmethod
    def _gather_dependencies_impl(dependency_context, dependency_job, dependency_handle):
        dynamic_request = DynamicRequest(dependency_job.request)
        dynamic_request.replace_key(KeyTokenEx(kTokenTypeRTXPackedMeshLOD), KeyTokenEx(kTokenTypeMeshLOD))
        dependency_context.agent.service.add_dependency(dependency_context.agent, dependency_handle, dynamic_request.get_request())
        return OperationResult.SUCCESS

class ProcessorMeshLOD(ProcessorImpl):
    def __init__(self, agent, key_token):
        super().__init__(agent, key_token, ProcessorMeshLOD._get_info_impl, ProcessorMeshLOD._match_impl, ProcessorMeshLOD._gather_dependencies_impl, ProcessorMeshLOD._build_impl)

    @staticmethod
    def _get_info_impl(processor):
        return OperationResult.SUCCESS, ProcessorInformation("ProcessorMeshLOD", 1, 1)

    @staticmethod
    def _gather_dependencies_impl(dependency_context, dependency_job, dependency_handle):
        dynamic_request = DynamicRequest(dependency_job.request)
        dynamic_request.replace_key(KeyTokenEx(kTokenTypeMeshLOD), KeyTokenEx(kTokenTypeTriangulatedMesh))
        dynamic_request.remove_key(KeyTokenEx(kTokenMeshReduceFactor))
        dependency_context.agent.service.add_dependency(dependency_context.agent, dependency_handle, dynamic_request.get_request())
        return OperationResult.SUCCESS

class ProcessorTriangulatedMesh(ProcessorImpl):
    def __init__(self, agent, key_token):
        super().__init__(agent, key_token, ProcessorTriangulatedMesh._get_info_impl, ProcessorTriangulatedMesh._match_impl, ProcessorTriangulatedMesh._gather_dependencies_impl, ProcessorTriangulatedMesh._build_impl)

    @staticmethod
    def _get_info_impl(processor):
        return OperationResult.SUCCESS, ProcessorInformation("ProcessorTriangulatedMesh", 1, 1)

    @staticmethod
    def _gather_dependencies_impl(dependency_context, dependency_job, dependency_handle):
        dynamic_request = DynamicRequest(dependency_job.request)
        dynamic_request.replace_key(KeyTokenEx(kTokenTypeTriangulatedMesh), KeyTokenEx(kTokenTypeMesh))
        dependency_context.agent.service.add_dependency(dependency_context.agent, dependency_handle, dynamic_request.get_request())
        return OperationResult.SUCCESS

class ProcessorMesh(ProcessorImpl):
    def __init__(self, agent, key_token):
        super().__init__(agent, key_token, ProcessorMesh._get_info_impl, ProcessorMesh._match_impl, ProcessorMesh._gather_dependencies_impl, ProcessorMesh._build_impl)

    @staticmethod
    def _get_info_impl(processor):
        return OperationResult.SUCCESS, ProcessorInformation("ProcessorMesh", 1, 1)

    @staticmethod
    def _gather_dependencies_impl(dependency_context, dependency_job, dependency_handle):
        return OperationResult.SUCCESS

class TestUJITSO(AsyncTestCase):
    """Test UJITSO (ported from TestUJITSO.cpp)"""

    async def test_ujitso_agent(self):
        """UJITSO agent test"""
        self._test_UJITSO()

    def _test_UJITSO(self, task_agent = None, task_service = None):
        factory_interface = acquire_factory_interface()
        self.assertIsNotNone(factory_interface)

        data_grid_interface = acquire_data_grid_interface()
        self.assertIsNotNone(data_grid_interface)
        data_grid = data_grid_interface.create_data_grid()
        self.assertIsNotNone(data_grid)
        local_data_store_interface = acquire_local_data_store_interface()
        self.assertIsNotNone(local_data_store_interface)
        data_store = local_data_store_interface.create(None, 0)
        self.assertIsNotNone(data_store)

        agent = factory_interface.create_agent(data_grid, data_store, task_agent, task_service, AgentConfigFlags.Default)
        self.assertIsNotNone(agent)
        self.assertIsNotNone(agent.agent)

        self._test_registry(agent)
        self._test_request(agent)
        self._test_lod_mockup(agent)
        self._test_build(agent, data_grid, data_store)

        factory_interface.destroy_agent(agent)
        data_grid_interface.destroy_data_grid(data_grid)
        local_data_store_interface.destroy(data_store)

        release_factory_interface(factory_interface)
        release_local_data_store_interface(local_data_store_interface)
        release_data_grid_interface(data_grid_interface)

    def _test_registry(self, agent):
        string1 = KeyTokenEx("Test1")
        string2 = KeyTokenEx("Test2")
        string3 = KeyTokenEx("Test1")

        self.assertTrue(string1 != string2)
        self.assertTrue(string2 != string3)
        self.assertTrue(string1 == string3)

    def _test_request(self, agent):
        agent_interface = agent.agent

        keys = [(123,)]
        key_token_rtype = UjitsoUtils.make_key_token_request_token_type()
        types = [(key_token_rtype.size, key_token_rtype.typeNameHash)]
        request_type = RequestType(keys, types)

        key_token_dtype = np.dtype([('value', np.uint32)])
        key_token_value = [1]
        request_values = np.array(key_token_value, key_token_dtype).tobytes()

        request = Request(request_type, request_values)

        request_callback_data = RequestCallbackData()

        operation_result, request_handle = agent_interface.request_build(agent, request, request_callback_data)
        self.assertEqual(operation_result, OperationResult.SUCCESS)
        self.assertIsNotNone(request_handle)
        self.assertNotEqual(request_handle.value, 0)

        operation_result = agent_interface.wait_all(agent, TIME_OUT_INFINITE)
        self.assertEqual(operation_result, OperationResult.SUCCESS)

        operation_result, result_handle = agent_interface.get_request_result(request_handle)
        self.assertEqual(operation_result, OperationResult.NOPROCESSOR_ERROR)

        operation_result = agent_interface.destroy_request(request_handle)
        self.assertEqual(operation_result, OperationResult.SUCCESS)

        operation_result = agent_interface.destroy_request(request_handle)
        self.assertEqual(operation_result, OperationResult.INVALIDHANDLE_ERROR)

        operation_result, result_handle = agent_interface.get_request_result(request_handle)
        self.assertEqual(operation_result, OperationResult.INVALIDHANDLE_ERROR)        

    def _test_lod_mockup(self, agent):
        agent_interface = agent.agent

        procPackedLODs = ProcessorRTXPackedMeshLODs(agent, kTokenTypeRTXPackedMeshLODs)
        procPackedLOD = ProcessorRTXPackedMeshLOD(agent, kTokenTypeRTXPackedMeshLOD)
        procMeshLOD = ProcessorMeshLOD(agent, kTokenTypeMeshLOD)
        procTriangulatedMesh = ProcessorTriangulatedMesh(agent, kTokenTypeTriangulatedMesh)
        procMesh = ProcessorMesh(agent, kTokenTypeMesh)

        dynamic_request = DynamicRequest()
        dynamic_request.add(KeyTokenEx(kTokenTypeRTXPackedMeshLODs))
        dynamic_request.add_string(KeyTokenEx(kTokenTypeUSDPath), "/Meshes/TestMesh")
        request = dynamic_request.get_request()
        self.assertIsNotNone(request)

        request_callback_data = RequestCallbackData()
        operation_result, request_handle = agent_interface.request_build(agent, request, request_callback_data)
        self.assertEqual(operation_result, OperationResult.SUCCESS)
        self.assertIsNotNone(request_handle)
        self.assertNotEqual(request_handle.value, 0)

        operation_result = agent_interface.wait_all(agent, TIME_OUT_INFINITE)
        self.assertEqual(operation_result, OperationResult.SUCCESS)

        agent_interface.destroy_request(request_handle)

    def _test_build(self, agent, data_grid, data_store):
        class ProcessorBuild(ProcessorImpl):
            def __init__(self, agent, key_token):
                super().__init__(agent, key_token, ProcessorBuild._get_info_impl, ProcessorBuild._match_impl, ProcessorBuild._gather_dependencies_impl, ProcessorBuild._build_impl)

            @staticmethod
            def _get_info_impl(processor):
                return OperationResult.SUCCESS, ProcessorInformation("ProcessorBuild", 1, 1)

            @staticmethod
            def _gather_dependencies_impl(dependency_context, dependency_job, dependency_handle):
                agent = dependency_context.agent
                service = agent.service
                dynamic_request = DynamicRequest(dependency_job.request)
                service.add_request_tuple_input(agent, dependency_handle, dynamic_request.get_request(), False, False)
                service.set_storage_context(agent, dependency_handle, "TestStorageContext")
                return OperationResult.SUCCESS

            @staticmethod
            def _build_impl(build_context, build_job, build_handle):
                agent = build_context.agent

                string_value = UjitsoUtils.get_request_value_string(agent, build_job.request, KeyTokenEx("String"), None)
                int_value = UjitsoUtils.get_request_value_int(agent, build_job.request, KeyTokenEx("IntParam"), 0)

                dtype = np.dtype('uint32')
                elem_size = dtype.itemsize

                operation_result, metadata = agent.service.allocate_meta_data_storage(agent, build_handle, elem_size * 255)
                self.assertEqual(operation_result, OperationResult.SUCCESS)

                # reinterpret data type
                metadata_array = np.frombuffer(metadata, dtype)

                for i in range(len(metadata_array)):
                    metadata_array[i] = int_value

                external_data = [np.frombuffer(string_value.encode(), dtype=np.uint8)]
                validation_data = [ValidationType.MANDATORY]
                operation_result = agent.service.store_external_data(agent, build_handle, external_data, validation_data)

                return OperationResult.SUCCESS

        kIntValue = 0x102
        kStringValue = "MyTestStringForFun"
        kTokenRequest = KeyTokenEx("String")
        test_processor = ProcessorBuild(agent, kTokenRequest)

        dynamic_request = DynamicRequest()
        dynamic_request.add_int(KeyTokenEx("IntParam"), kIntValue)
        dynamic_request.add_string(KeyTokenEx("String"), kStringValue)
        request = dynamic_request.get_request()

        request_callback_data = RequestCallbackData()

        agent_interface = agent.agent
        operation_result, request_handle = agent_interface.request_build(agent, request, request_callback_data)
        self.assertEqual(operation_result, OperationResult.SUCCESS)
        self.assertIsNotNone(request_handle)
        self.assertNotEqual(request_handle.value, 0)

        operation_result = agent_interface.wait_request(request_handle, TIME_OUT_INFINITE)
        self.assertEqual(operation_result, OperationResult.SUCCESS)
        
        operation_result, result_handle = agent_interface.get_request_result(request_handle)

        operation_result, metadata = agent_interface.get_request_meta_data(result_handle)
        self.assertEqual(operation_result, OperationResult.SUCCESS)
        
        # reinterpret data type
        metadata_array = np.frombuffer(metadata, dtype='uint32')

        for i in range(len(metadata_array)):
            self.assertEqual(metadata_array[i], kIntValue)

        operation_result, storages = agent_interface.get_request_external_data(result_handle)
        self.assertEqual(operation_result, OperationResult.SUCCESS)
        self.assertEqual(len(storages), 1)

        operation_result = agent_interface.validate_request_external_data(request_handle, result_handle, [True], True)
        self.assertEqual(operation_result, OperationResult.SUCCESS)

        storage = storages[0]
        operation_result, data_block = DataStoreUtils.copy_data_block(data_store, "", storage)
        string_value = ''.join(chr(v) for v in data_block)
        self.assertEqual(string_value, kStringValue)

        operation_result, storage_context = agent_interface.get_request_storage_context(request_handle)
        self.assertEqual(operation_result, OperationResult.SUCCESS)
        self.assertEqual(storage_context, "TestStorageContext")

        agent_interface.destroy_request(request_handle)
