import asyncio
import aiohttp
import carb
import os
import toml
import time
import omni.client
import webbrowser

from omni.kit.menu.utils import MenuItemDescription, MenuAlignment
from omni import ui
from typing import Union
from .style import Styles


class CacheStateDelegate(ui.MenuDelegate):
    def __init__(self, cache_enabled, **kwargs):
        super().__init__(**kwargs)
        
        self._cache_enabled = cache_enabled
        self._cache_widget = None
    
    def destroy(self):
        self._cache_widget = None

    def build_item(self, item: ui.MenuHelper):
        with ui.HStack(width=0, style={"margin" : 0}):
            ui.Spacer(width=5)
            self._cache_widget = ui.HStack(content_clipping=1, width=0, style=Styles.CACHE_STATE_ITEM_STYLE)
            ui.Spacer(width=10)
        
        self.update_live_state(self._cache_enabled)

    def get_menu_alignment(self):
        return MenuAlignment.RIGHT
    
    def update_menu_item(self, menu_item: Union[ui.Menu, ui.MenuItem], menu_refresh: bool):
        if isinstance(menu_item, ui.MenuItem):
            menu_item.visible = False
        
    def update_live_state(self, cache_enabled):
        if not self._cache_widget:
            return

        margin = 2
        self._cache_enabled = cache_enabled
        self._cache_widget.clear()
        with self._cache_widget:
            ui.Label("CACHE: ")
            if cache_enabled:
                ui.Label("ON", style={"color": 0xff00b86b})
            else:
                with ui.VStack():
                    ui.Spacer(height=margin)
                    with ui.ZStack(width=0):
                        with ui.HStack(width=20):
                            ui.Spacer()
                            ui.Label("OFF", name="offline", width=0)
                            ui.Spacer(width=margin)
                            with ui.VStack(width=0):
                                ui.Spacer()
                                ui.Image(width=14, height=14, name="doc")
                                ui.Spacer()
                            ui.Spacer()
                        button = ui.InvisibleButton(width=20)
                        button.set_clicked_fn(lambda: webbrowser.open('https://docs.omniverse.nvidia.com/nucleus/cache_troubleshoot.html', new=2))
                    ui.Spacer(height=margin)


class CacheStateMenu:
    def __init__(self):
        self._live_menu_name = "Cache State Widget"
        self._menu_list = [MenuItemDescription(name="placeholder", show_fn=lambda: False)]

        global_config_path = carb.tokens.get_tokens_interface().resolve("${omni_global_config}")
        self._omniverse_config_path = os.path.join(global_config_path, "omniverse.toml").replace("\\", "/")
        self._all_cache_apis = []
        self._cache_enabled = False
        self._last_time_check = 0
        self._ping_cache_future = None
        self._update_subscription = None

    def _initialize(self):
        if os.path.exists(self._omniverse_config_path):
            try:
                contents = toml.load(self._omniverse_config_path)
            except Exception as e:
                carb.log_error(f"Unable to parse {self._omniverse_config_path}. File corrupted?")
                contents = None

            if contents:
                self._all_cache_apis = self._load_all_cache_server_apis(contents)
                if self._all_cache_apis:
                    self._cache_enabled = True
                    self._update_subscription = omni.kit.app.get_app().get_update_event_stream().create_subscription_to_pop(
                        self._on_update, name="omni.kit.widget.live update"
                    )
                else:
                    carb.log_warn("Unable to detect Omniverse Cache Server. Consider installing it for better IO performance.")
                    self._cache_enabled = False
        else:
            carb.log_warn(f"Unable to detect Omniverse Cache Server. File {self._omniverse_config_path} is not found."
                          f" Consider installing it for better IO performance.")


    def register_menu_widgets(self):
        self._initialize()
        self._cache_state_delegate = CacheStateDelegate(self._cache_enabled)
        omni.kit.menu.utils.add_menu_items(self._menu_list, name=self._live_menu_name, delegate=self._cache_state_delegate)

    def unregister_menu_widgets(self):
        omni.kit.menu.utils.remove_menu_items(self._menu_list, self._live_menu_name)
        if self._cache_state_delegate:
            self._cache_state_delegate.destroy()
        self._cache_state_delegate = None
        self._menu_list = None
        self._update_subscription = None
        self._all_cache_apis = []

        try:
            if self._ping_cache_future and not self._ping_cache_future.done():
                self._ping_cache_future.cancel()
                self._ping_cache_future = None
        except Exception:
            self._ping_cache_future = None
    
    def _on_update(self, dt):
        if not self._cache_state_delegate or not self._all_cache_apis:
            return
        
        if not self._ping_cache_future or self._ping_cache_future.done():
            now = time.time()
            duration = now - self._last_time_check

            # 30s
            if duration < 30:
                return

            self._last_time_check = now

            async def _ping_cache():
                session = aiohttp.ClientSession()

                cache_enabled = None
                for cache_api in self._all_cache_apis:
                    try:
                        async with session.head(cache_api):
                            ''' If we're here the service port is alive '''
                            cache_enabled = True
                    except Exception as e:
                        cache_enabled = False
                        break
                
                await session.close()
            
                if cache_enabled is not None and self._cache_enabled != cache_enabled:
                    self._cache_enabled = cache_enabled
                    if self._cache_state_delegate:
                        self._cache_state_delegate.update_live_state(self._cache_enabled)

            self._ping_cache_future = asyncio.ensure_future(_ping_cache())
    
    def _load_all_cache_server_apis(self, config_contents):
        mapping = os.environ.get("OMNI_CONN_CACHE", None)
        if mapping:
            mapping = f"*#{mapping},f"
        else:
            mapping = os.environ.get("OMNI_CONN_REDIRECTION_DICT", None)
            if not mapping:
                mapping = os.environ.get("OM_REDIRECTION_DICT", None)

            if not mapping:
                connection_library_dict = config_contents.get("connection_library", None)
                if connection_library_dict:
                    mapping = connection_library_dict.get("proxy_dict", None)
        
        all_proxy_apis = set([])
        if mapping:
            mapping = mapping.lstrip("\"")
            mapping = mapping.rstrip("\"")
            mapping = mapping.lstrip("'")
            mapping = mapping.rstrip("'")

            redirections = mapping.split(";")
            for redirection in redirections:
                parts = redirection.split("#")
                if not parts or len(parts) < 2:
                    continue
                
                source, target = parts[0], parts[1]
                targets = target.split(",")
                if not targets:
                    continue

                if len(targets) > 1:
                    proxy_address = targets[0]
                else:
                    proxy_address = target
                
                if not proxy_address.startswith("http://") and not proxy_address.startswith("https://"):
                    proxy_address_api = f"http://{proxy_address}/ping"
                else:
                    if proxy_address.endswith("/"):
                        proxy_address_api = f"{proxy_address}ping"
                    else:
                        proxy_address_api = f"{proxy_address}/ping"
                all_proxy_apis.add(proxy_address_api)

        return all_proxy_apis


