import os
import asyncio
import weakref
import omni
import omni.usd
import carb

from typing import Callable
from omni.kit.widget.prompt import PromptButtonInfo, PromptManager
from .omni_client_wrapper import OmniClientWrapper
from .main_window import CollectMainWindow
from .collector import Collector, CollectorException, CollectorFailureOptions, FlatCollectionTextureOptions
from .progress_popup import ProgressPopup
from .utils import Utils
from omni.kit.menu.utils import MenuItemDescription

g_singleton = None


def get_instance():
    global g_singleton
    return g_singleton


class PublicExtension(omni.ext.IExt):
    def on_startup(self):
        self._main_window = None
        self._context_menu_name = None
        self._register_menus()

        global g_singleton
        g_singleton = self

    def on_shutdown(self):
        global g_singleton
        g_singleton = None

        content_window = self.get_content_window()
        if content_window and self._context_menu_name:
            content_window.delete_context_menu(self._context_menu_name)
            self._context_menu_name = None

        if self._main_window:
            self._main_window.destroy()
            self._main_window = None

    def collect(self, filepath: str, finish_callback: Callable[[], None] = None) -> None:
        """
        Collect a usd file.
        Args:
            filepath: Path to usd file to be collected.
        """
        if not omni.usd.is_usd_writable_filetype(filepath):
            self._show_file_not_supported_popup()
        else:
            self._show_main_window(filepath, finish_callback)

    def _register_menus(self):
        content_window = self.get_content_window()
        if content_window:
            self._context_menu_name = content_window.add_context_menu(
                "Collect Asset", "upload.svg",
                lambda b, c: self._on_menu_click(b, c),
                omni.usd.is_usd_writable_filetype
            )

        def _collect():
            stage = omni.usd.get_context().get_stage()
            self.collect(stage.GetRootLayer().identifier)

        def _enable_collect_menu():
            stage = omni.usd.get_context().get_stage()
            return stage is not None and not stage.GetRootLayer().anonymous

        self._file_menu_list = [
            MenuItemDescription(
                name="Collect As...",
                glyph="none.svg",
                appear_after="Save Flattened As...",
                enable_fn=_enable_collect_menu,
                onclick_fn=_collect,
            )
        ]

        omni.kit.menu.utils.add_menu_items(self._file_menu_list, "File")

    def _get_collection_dir(self, folder, usd_file_name):
        stage_name = os.path.splitext(usd_file_name)[0]
        if not folder.endswith("/"):
            folder += "/"

        folder = Utils.normalize_path(folder)
        return f"{folder}Collected_{stage_name}"

    def _show_folder_exist_popup(
        self,
        usd_path,
        collect_dir,
        usd_only,
        flat_collection,
        material_only,
        texture_option,
        finish_callback: Callable[[], None] = None,
    ):
        def on_confirm():
            self._start_collecting(
                usd_path, collect_dir, usd_only, flat_collection, material_only, texture_option, finish_callback
            )

        def on_cancel():
            self._show_main_window(usd_path)

        PromptManager.post_simple_prompt(
            "Overwrite",
            "The target directory already exists, do you want to overwrite it?",
            PromptButtonInfo("Confirm", on_confirm),
            PromptButtonInfo("Cancel", on_cancel),
            modal=False,
        )

    def _start_collecting(
        self,
        usd_path,
        collect_folder,
        usd_only,
        flat_collection,
        material_only,
        texture_option,
        finish_callback: Callable[[], None] = None,
    ):
        progress_popup = self._show_progress_popup()
        progress_popup.status_text = "Collecting dependencies..."
        collector = Collector(
            usd_path,
            collect_folder,
            usd_only,
            flat_collection,
            material_only,
            texture_option=FlatCollectionTextureOptions(texture_option),
        )

        collector_weakref = weakref.ref(collector)

        def on_cancel():
            carb.log_info("Cancelling collector...")
            if not collector_weakref():
                return

            collector_weakref().cancel()

        progress_popup.set_cancel_fn(on_cancel)

        def on_progress(step, total):
            progress_popup.status_text = f"Collecting USD {os.path.basename(usd_path)}..."
            if total != 0:
                progress_popup.progress = float(step) / total
            else:
                progress_popup.progress = 0.0

        def on_finish():
            if finish_callback:
                finish_callback()

            progress_popup.hide()

            self._refresh_current_directory()

            if self._main_window:
                self._main_window.set_collect_fn(None)

            if not collector_weakref():
                return
            collector_weakref().destroy()

        asyncio.ensure_future(collector.collect(on_progress, on_finish))

    def _show_main_window(self, usd_path, finish_callback: Callable[[], None] = None):
        if not self._main_window:
            self._main_window = CollectMainWindow(None, None)

        def on_collect(collect_folder, usd_only, flat_collection, material_only, texture_option):
            existed = OmniClientWrapper.exists_sync(collect_folder)
            if existed:
                self._show_folder_exist_popup(
                    usd_path, collect_folder, usd_only, flat_collection, material_only, texture_option, finish_callback
                )
            else:
                self._start_collecting(
                    usd_path, collect_folder, usd_only, flat_collection, material_only, texture_option, finish_callback
                )

        self._main_window.set_collect_fn(on_collect)
        default_folder = self._get_collection_dir(os.path.dirname(usd_path), os.path.basename(usd_path))
        self._main_window.show(default_folder)

    def _on_menu_click(self, menu, value):
        self.collect(value)

    def _refresh_current_directory(self):
        content_window = self.get_content_window()
        if content_window:
            content_window.refresh_current_directory()

    def _show_progress_popup(self):
        progress_popup = ProgressPopup("Collecting")
        progress_popup.progress = 0
        progress_popup.show()

        return progress_popup

    def _show_file_not_supported_popup(self):
        PromptManager.post_simple_prompt("Warning", "Only USD file can be collected")
    
    def get_content_window(self):
        try:
            import omni.kit.window.content_browser as content
            
            return content.get_content_window()
        except Exception as e:
            pass
        
        return None

