import carb
import asyncio
import omni.ui as ui
import omni.kit.app
import omni.kit.usd.layers as layers
import omni.kit.notification_manager as nm

from omni.kit.widget.prompt import PromptManager
from pxr import Sdf
from .file_picker import FilePicker, FileBrowserMode, FileBrowserSelectionType


class LiveSessionEndWindow:

    def __init__(self, layers_interface: layers.Layers, layer_identifier):
        self._window = None
        self._live_syncing = layers_interface.get_live_syncing()
        self._base_layer_identifier = layer_identifier
        self._buttons = []
        self._options_combo = None
        self._file_picker = None

        self._key_functions = {
            int(carb.input.KeyboardInput.ENTER): self._on_ok_button_fn,
            int(carb.input.KeyboardInput.ESCAPE): self._on_cancel_button_fn
        }

        self._build_ui()

    def destroy(self):
        for button in self._buttons:
            button.set_clicked_fn(None)
        self._buttons.clear()
        if self._window:
            self._window.visible = False
        self._window = None
        self._options_combo = None
        if self._file_picker:
            self._file_picker.destroy()
        self._file_picker = None
    
    async def __aenter__(self):
        await self._live_syncing.broadcast_merge_started_message_async(self._base_layer_identifier)
        self.visible = True
        # Wait until dialog disappears
        while self.visible:
            await asyncio.sleep(0.1)

    async def __aexit__(self, exc_type, exc, tb):
        self.visible = False

    @property
    def visible(self):
        return self._window and self._window.visible
    
    @visible.setter
    def visible(self, value):
        if self._window:
            self._window.visible = value

    def _on_ok_button_fn(self):
        self.visible = False
        if not self._live_syncing.is_layer_in_live_session(self._base_layer_identifier):
            return

        prompt = None

        async def pre_merge(current_session):
            nonlocal prompt
            app = omni.kit.app.get_app()
            await app.next_update_async()
            await app.next_update_async()
            prompt = PromptManager.post_simple_prompt(
                "Merging",
                "Merging live layers into base layers...",
                None,
                shortcut_keys=False
            )
        
        async def post_merge(success):
            nonlocal prompt
            if prompt:
                prompt.destroy()
                prompt = None
        
            if success:
                await self._live_syncing.broadcast_merge_done_message_async(
                    layer_identifier=self._base_layer_identifier
                )
        
        current_session = self._live_syncing.get_current_live_session(self._base_layer_identifier)
        comment = self._description_field.model.get_value_as_string()
        option = self._options_combo.model.get_item_value_model().as_int
        if option == 1:
            def save_model(file_path: str, overwrite_existing: bool):
                layer = Sdf.Layer.FindOrOpen(file_path)
                if not layer:
                    layer = Sdf.Layer.CreateNew(file_path)
                
                if not layer:
                    error = f"Failed to save live changes to layer {file_path} as it's not writable."
                    carb.log_error(error)
                    nm.post_notification(error, status=nm.NotificationStatus.WARNING)

                    return

                asyncio.ensure_future(
                    self._live_syncing.merge_and_stop_live_session_async(
                        file_path, comment, pre_merge=pre_merge, post_merge=post_merge,
                        layer_identifier=self._base_layer_identifier
                    )
                )

            usd_context = self._live_syncing.usd_context
            stage = usd_context.get_stage()
            root_layer_identifier = stage.GetRootLayer().identifier
            root_layer_name = current_session.name
            self._show_file_picker(save_model, root_layer_identifier, root_layer_name)
        elif option == 0:
            asyncio.ensure_future(
                self._live_syncing.merge_and_stop_live_session_async(
                    comment=comment, pre_merge=pre_merge, post_merge=post_merge,
                    layer_identifier=self._base_layer_identifier
                )
            )

    def _on_cancel_button_fn(self):
        self.visible = False

    def _on_key_pressed_fn(self, key, mod, pressed):
        if not pressed:
            return

        func = self._key_functions.get(key)
        if func:
            func()
    
    def _create_file_picker(self):
        filter_options = [
            (r"^(?=.*.usd$)((?!.*\.(sublayer)\.usd).)*$", "USD File (*.usd)"),
            (r"^(?=.*.usda$)((?!.*\.(sublayer)\.usda).)*$", "USDA File (*.usda)"),
            (r"^(?=.*.usdc$)((?!.*\.(sublayer)\.usdc).)*$", "USDC File (*.usdc)"),
            ("(.*?)", "All Files (*.*)"),
        ]
        layer_file_picker = FilePicker(
            "Save Live Changes",
            FileBrowserMode.SAVE,
            FileBrowserSelectionType.FILE_ONLY,
            filter_options,
            [".usd", ".usda", ".usdc", ".usd"],
        )

        return layer_file_picker

    def _show_file_picker(self, file_handler, default_location=None, default_filename=None):
        if not self._file_picker:
            self._file_picker = self._create_file_picker()

        self._file_picker.set_file_selected_fn(file_handler)
        self._file_picker.show(default_location, default_filename)

    def _on_description_begin_edit(self, model):
        self._description_field_hint_label.visible = False

    def _on_description_end_edit(self, model):
        if len(model.get_value_as_string()) == 0:
            self._description_field_hint_label.visible = True

    def _build_ui(self):
        """Construct the window based on the current parameters"""
        self._window = ui.Window("Merge Options", visible=False, height=0, dockPreference=ui.DockPreference.DISABLED)
        self._window.flags = (
            ui.WINDOW_FLAGS_NO_COLLAPSE | ui.WINDOW_FLAGS_NO_SCROLLBAR |
            ui.WINDOW_FLAGS_NO_RESIZE | ui.WINDOW_FLAGS_NO_MOVE
        )
        self._window.set_key_pressed_fn(self._on_key_pressed_fn)
        self._window.flags = self._window.flags | ui.WINDOW_FLAGS_MODAL

        current_session = self._live_syncing.get_current_live_session(self._base_layer_identifier)
        with self._window.frame:
            with ui.VStack(height=0):
                ui.Spacer(width=0, height=10)
                with ui.HStack(height=0):
                    ui.Spacer()
                    ui.Label(
                        "Live Session is ending.", word_wrap=True, alignment=ui.Alignment.CENTER, width=self._window.width - 80, height=0
                    )
                    ui.Spacer()
                with ui.HStack(height=0):
                    ui.Spacer()
                    ui.Label(
                        f"How do you want to merge changes from '{current_session.name}' session?",
                        alignment=ui.Alignment.CENTER, word_wrap=True, width=self._window.width - 80, height=0
                    )
                    ui.Spacer()
                ui.Spacer(width=0, height=10)
                with ui.HStack(height=0):
                    ui.Spacer()
                    self._options_combo = ui.ComboBox(
                        0, "Merge to corresponding layers", "Merge to a new layer",
                        word_wrap=True, width=self._window.width - 80, height=0
                    )
                    ui.Spacer()

                ui.Spacer(width=0, height=10)

                with ui.HStack(height=0):
                    ui.Spacer(width=40)
                    self._checkpoint_comment_frame = ui.Frame()
                    with self._checkpoint_comment_frame:
                        with ui.VStack(height=0, spacing=5):
                            ui.Label("Checkpoint Description")
                            with ui.ZStack():
                                self._description_field = ui.StringField(multiline=True, height=80)
                                self._description_field_hint_label = ui.Label(
                                    " Description", alignment=ui.Alignment.LEFT_TOP, style={"color": 0xFF3F3F3F}
                                )
                                self._description_begin_edit_sub = self._description_field.model.subscribe_begin_edit_fn(
                                    self._on_description_begin_edit
                                )
                                self._description_end_edit_sub = self._description_field.model.subscribe_end_edit_fn(
                                    self._on_description_end_edit
                                )
                    ui.Spacer(width=40)

                ui.Spacer(width=0, height=30)
                with ui.HStack(height=0):
                    ui.Spacer(height=0)
                    ok_button = ui.Button("CONTINUE", name="confirm_button", width=120, height=0)
                    ok_button.set_clicked_fn(self._on_ok_button_fn)
                    cancel_button = ui.Button("CANCEL", name="cancel_button", width=120, height=0)
                    cancel_button.set_clicked_fn(self._on_cancel_button_fn)
                    self._buttons.append(ok_button)
                    self._buttons.append(cancel_button)
                    ui.Spacer(height=0)
                ui.Spacer(width=0, height=10)
