import os
import importlib

import carb
import carb.settings

try:
    import omni.renderer_capture
    omni_renderer_capture_present = True
except ImportError:
    omni_renderer_capture_present = False

import omni.ext
import omni.kit.app

from omni.hydra.engine.stats import HydraEngineStats


class Extension(omni.ext.IExt):
    def __init__(self):
        super().__init__()
        pass

    def _set_default_settings(self):
        self._settings.set_default_int("/app/captureFrame/startFrame", -1)
        self._settings.set_default("/app/captureFrame/startMultipleFrame/0", -1)
        self._settings.set_default_bool("/app/captureFrame/closeApplication", False)
        self._settings.set_default_string("/app/captureFrame/fileName", "no-filename-specified")
        self._settings.set_default_string("/app/captureFrame/outputPath", "")
        self._settings.set_default_bool("/app/captureFrame/setAlphaTo1", True)
        self._settings.set_default_bool("/app/captureFrame/saveFps", False)
        self._settings.set_default_bool("/app/captureFrame/hdr", False)
        self._settings.set_default_int("/app/captureFrame/asyncBufferSizeMB", 2048)
        self._settings.set_default_bool("/renderer/gpuProfiler/record", False)
        self._settings.set_default_int("/renderer/gpuProfiler/maxIndent", 1)

    def on_startup(self):
        self._settings = carb.settings.get_settings()
        self._set_default_settings()

        self._app = omni.kit.app.get_app()

        try:
            module_omni_usd = importlib.import_module('omni.usd')
            self._usd_context = module_omni_usd.get_context()
            self._opened_state = module_omni_usd.StageState.OPENED
        except ImportError:
            self._usd_context = None
            self._opened_state = None

        if omni_renderer_capture_present:
            self._renderer_capture = omni.renderer_capture.acquire_renderer_capture_interface()
            self._renderer_capture.start_frame_updates()
        else:
            self._renderer_capture = None
            carb.log_error("Autocapture initialization failed: renderer.capture extension should be present!")
            return

        # Initial configuration
        self._frame_no = 0
        # App is exiting before last image has been saved, exit after _quitFrameCounter frames
        self._quitFrameCounter = 10
        self._multiple_frame_no = 0
        self._start_frame = self._settings.get("/app/captureFrame/startFrame")
        self._start_multiple_frame = self._settings.get("/app/captureFrame/startMultipleFrame")
        self._close_app = self._settings.get("/app/captureFrame/closeApplication")
        self._file_name = self._settings.get("/app/captureFrame/fileName")
        self._output_path = self._settings.get("/app/captureFrame/outputPath")

        if len(self._output_path) == 0:
            module_carb_tokens = importlib.import_module('carb.tokens')
            self._output_path = module_carb_tokens.get_tokens_interface().resolve("${kit}") + "/../../../outputs/"

        self._record_gpu_performance = self._settings.get("/renderer/gpuProfiler/record")
        self._recording_max_indent = self._settings.get("/renderer/gpuProfiler/maxIndent")
        self._gpu_perf = []

        # viewport_api = get_active_viewport()
        # self.__stats = HydraEngineStats(viewport_api.usd_context_name, viewport_api.hydra_engine)
        self.__stats = HydraEngineStats()

        self._count_loading_frames = False
        self._next_frame_exit = False

        if self._start_frame > 0 or self._start_multiple_frame[0] > 0:

            def on_post_update(e: carb.events.IEvent):
                if not self._app.is_app_ready():
                    return

                if self._next_frame_exit:
                    if self._quitFrameCounter <= 0:
                         self._app.post_quit()
                    self._quitFrameCounter = self._quitFrameCounter - 1
                    return None

                count_frame = True
                if not self._count_loading_frames and self._usd_context is not None:
                    if self._usd_context.get_stage_state() != self._opened_state:
                        count_frame = True

                if count_frame:
                    if self._record_gpu_performance and self.__stats:
                        frame_perf = self.__stats.get_nested_gpu_profiler_result(self._recording_max_indent)
                        dev_count = len(frame_perf)
                        has_data = False
                        for dev_idx in range(dev_count):
                            if len(frame_perf[dev_idx]) > 0:
                                has_data = True
                                break
                        if has_data:
                            if len(self._gpu_perf) == 0:
                                self._gpu_perf.extend(frame_perf)
                                for dev_idx in range(dev_count):
                                    self._gpu_perf[dev_idx] = {}

                            for dev_idx in range(dev_count):
                                self._gpu_perf[dev_idx]["frame %d" % (self._frame_no)] = frame_perf[dev_idx]

                    self._frame_no += 1
                    if self._start_frame > 0:
                        if self._frame_no >= self._start_frame:
                            self._renderer_capture.capture_next_frame_swapchain(self._output_path + self._file_name)
                            self._next_frame_exit = self._close_app

                    if self._start_multiple_frame[0] > 0 and self._multiple_frame_no < len(self._start_multiple_frame):                        
                        if self._frame_no >= self._start_multiple_frame[self._multiple_frame_no] :
                            self._renderer_capture.capture_next_frame_swapchain(self._output_path + self._file_name + "_" +str(self._start_multiple_frame[self._multiple_frame_no]))
                            self._multiple_frame_no += 1
                            if self._multiple_frame_no >= len(self._start_multiple_frame):
                                self._next_frame_exit = self._close_app

            self._post_update_subs = self._app.get_post_update_event_stream().create_subscription_to_pop(on_post_update, name="Autocapture post-update")


    def on_shutdown(self):
        if self._record_gpu_performance:
            json_filename = self._output_path + self._file_name + ".json"

            dump_json = {}
            dev_count = len(self._gpu_perf)
            for dev_idx in range(dev_count):
                dump_json["GPU-%d" % (dev_idx)] = self._gpu_perf[dev_idx]

            import json
            with open(json_filename, 'w', encoding='utf-8') as json_file:
                json.dump(dump_json, json_file, ensure_ascii=False, indent=4)

        self._gpu_perf = None
        self._settings = None
        self._usd_context = None
        self._opened_state = None
        self._renderer_capture = None
        self._post_update_subs = None
