import carb
import carb.events
import asyncio
import omni.client
import omni.kit.tagging
from idl.connection.transport import TransportError

# rename this to omni.tagging.client after omni.tagging is deprecated
import omni.tagging_client as tagging_client


class TaggingClientContextAsync:
    """Create a tag client context."""

    def __init__(self, ov_server, **kwargs):
        self.ov_server = ov_server

    async def init_client(self):
        """Create tagging client."""
        self._helper = tagging_client.TaggingClientHelper(self.ov_server)
        # connect to the tagging service
        self.service = await self._helper.get_tagging_service()
        # check that tagging server is responding
        res = await self.service.generate_client_id()
        # get the ID of the client (this will likely be deprecated)
        self.client_id = res.client_id

    async def close_client(self, client_msg: str = None, *args, **kwargs):
        """Close tagging client."""
        await self._helper.__aexit__(*args, **kwargs)
        # log the given message on success
        if client_msg is not None:
            tag_utils_logger.debug(client_msg)


class OmniKitTaggingDelegate:
    """Delegate that handles when async tagging events are finished.
    """

    def updated_tags_for_url(self, url: str):
        pass

    def finished_getting_tags(self, urls: str):
        pass


class OmniKitTagging(omni.ext.IExt):
    """The Tagging plugin. This plugin holds tagging client contexts."""

    static_instance = None

    def on_startup(self):
        OmniKitTagging.static_instance = self
        self._connections = {}
        self._auth_tokens = {}
        self._path_results = {}
        self._tagging_clients = {}
        self._conn_subscription = omni.client.register_connection_status_callback(self._server_status_changed)
        self._delegate = None
        self._initializing = False

    def on_shutdown(self):
        for client in self._tagging_clients:
            asyncio.ensure_future(self._tagging_clients[client].close_client())
        self._tagging_clients = None
        self._conn_subscription = None
        self._auth_subscription = None
        self._delegate = None
        OmniKitTagging.static_instance = None

    def set_delegate(self, delegate: OmniKitTaggingDelegate):
        self._delegate = delegate

    def _server_status_changed(self, url, status):
        def _server_info_callback(result, server_info):
            nonlocal url
            broken_url = omni.client.break_url(url)
            if result == omni.client.Result.OK:
                self._auth_tokens[broken_url.host] = server_info.auth_token
            else:
                carb.log_error(f"Server info error: {str(result)}")

        broken_url = omni.client.break_url(url)
        if status == omni.client.ConnectionStatus.CONNECTED:
            self._connections[broken_url.host] = True
            omni.client.get_server_info_with_callback(url, _server_info_callback)
        else:
            self._connections[broken_url.host] = False

    async def _verify_connection(self, host):
        if host is None:
            return False

        if host not in self._connections:
            carb.log_warn(f"Unable to get tags. Not connected to host: {host}")
            return False

        if host not in self._auth_tokens:
            carb.log_error(f"Error getting auth token for host: {host}")
            return False

        if host not in self._tagging_clients:
            self._initializing = True
            self._tagging_clients[host] = TaggingClientContextAsync(ov_server=host)
            await self._tagging_clients[host].init_client()
            self._initializing = False

            if self._tagging_clients[host].service is None:
                carb.log_error(f"Error creating tagging client service on {host}")
                return False

        return True

    def tags_action(self, url, new_tag, action="add", old_tag=None):
        if action.lower() == "add":
            modify_type = tagging_client.TagModifyType.Add
        elif action.lower() == "set":
            modify_type = tagging_client.TagModifyType.Set
        elif action.lower() == "reset":
            modify_type = tagging_client.TagModifyType.Reset
        elif action.lower() == "remove":
            modify_type = tagging_client.TagModifyType.Remove
        elif action.lower() == "update":
            modify_type = tagging_client.TagModifyType.Remove
            asyncio.ensure_future(self.tags_action_async(url, old_tag, modify_type))
            modify_type = tagging_client.TagModifyType.Add
        else:
            carb.log_error(f"provided action is currently not supported: ({action})")

        asyncio.ensure_future(self.tags_action_async(url, new_tag, modify_type))

    async def tags_action_async(self, url, tag, modify_type):
        broken_url = omni.client.break_url(url)
        host = broken_url.host
        if not await self._verify_connection(host):
            return

        tag_dict = self.validate_tag(tag)
        if tag_dict:

            # 5 second timout for service initialization.
            timeout = 5.0
            while self._initializing and timeout > 0:
                await asyncio.sleep(0.5)
                timeout -= 0.5

            constructed_tag = tagging_client.Tag(
                name=tag_dict["key"], tag_namespace=tag_dict["full_namespace"], value=tag_dict["value"]
            )
            try:
                # get the tags
                results = await self._tagging_clients[host].service.modify_tags(
                    auth_token=self._auth_tokens[host],
                    client_id=self._tagging_clients[host].client_id,
                    paths=[broken_url.path],
                    tags=[constructed_tag],
                    modify_type=modify_type,
                )
            except AttributeError:
                # service is still being initialized
                carb.log_warn("Timeout waiting for tagging service initialization")
                return
            except TransportError as msg:
                # lost connection to tagging service
                carb.log_warn(f"TransportError: {msg}")
                self._tagging_clients = {}
                return

        if self._delegate:
            self._delegate.updated_tags_for_url(url)

    def get_tags(self, urls, callback):
        asyncio.ensure_future(self.get_tags_async(urls, callback))

    async def get_tags_async(self, urls, callback=None):

        paths = {}
        for url in urls:
            broken_url = omni.client.break_url(url)
            host = broken_url.host

            if not await self._verify_connection(host):
                return

            # 5 second timout for service initialization.
            timeout = 5.0
            while self._initializing and timeout > 0:
                await asyncio.sleep(0.5)
                timeout -= 0.5

            if host not in paths:
                paths[host] = [broken_url.path]
            else:
                paths[host].append(broken_url.path)

        for host in paths:
            try:
                # get the tags
                results = await self._tagging_clients[host].service.get_tags(
                    auth_token=self._auth_tokens[host], paths=paths[host]
                )
            except AttributeError as msg:
                # service is still being initialized
                carb.log_warn(f"Timeout waiting for tagging service initialization: {msg}")
                return
            except TransportError as msg:
                # lost connection to tagging service
                self._tagging_clients = {}
                return

            for i, result in enumerate(results["path_result"]):
                if result["connection_status"] != 1: # OmniTaggingStatusCodeOk
                    carb.log_warn(f"Tagging Connection Error Code: {result['connection_status']}")
                    self._tagging_clients = {}
                    return

                # TODO check connection status
                if host not in self._path_results:
                    self._path_results[host] = {}

                self._path_results[host][paths[host][i]] = result["tags"]

        if callback is not None:
            callback()
        if self._delegate:
            self._delegate.finished_getting_tags(urls)

        return self.tags_for_url(urls[0])

    def tags_for_url(self, url):

        broken_url = omni.client.break_url(url)
        host = broken_url.host
        path = broken_url.path

        if host in self._path_results and path in self._path_results[host]:
            return self._path_results[host][path]

    def pack_tag(self, namespace, hidden_label, kv_pair=None, key=None, value=None):
        """ Changes a keyvalue pair tag into a string like "namespace.key=value"

            Adds a leading . when the namespace has one, and ommits the =value when no value is present.
            Format of tags:
                [namespace].[key]{=value}
                .[namespace].[label].[key]{=value}
            For Example:
                appearance.car
                .appearance.generated.truck=60
                .appearance.excluded.banana
                tag_without_namespace
                .excluded.tag_without_namespace
        """
        if kv_pair is None:
            if not value:
                kv_pair = key
            else:
                kv_pair = key + "=" + value
        if hidden_label and len(hidden_label) > 0:
            if namespace == "":
                return ".".join(["", hidden_label, kv_pair])
            return ".".join(["", namespace, hidden_label, kv_pair])
        if namespace == "":
            return kv_pair
        return ".".join([namespace, kv_pair])

    def validate_tag(self, tag, allow_hidden_tags=True):
        """ Validate the tag, make sure there is a category, namespace, and key,
            and return the dict if valid.
        """
        is_hidden = False
        namespace = None
        hidden_label = None
        kv_split = tag.split("=")
        full_key = kv_split[0]
        val = ""
        if len(kv_split) > 2:
            carb.log_warn(f"Too many '=' characters in tag: {tag}")
            return False
        if ".." in full_key:
            carb.log_warn(f"Invalid tag (too many consecutive .'s): {tag}")
            return False

        if len(kv_split) > 1:
            val = kv_split[1]

        dot_split = full_key.split(".")

        # the full_namespace combines the hidden/system namespace with the normal one.
        if len(dot_split) <= 1:
            key = full_key
            full_namespace = ""
        else:
            key = dot_split[-1]
            full_namespace = ".".join(dot_split[:-1])

        # remove the key
        dot_split = dot_split[:-1]

        if len(dot_split) >= 1 and dot_split[0] == "":
            # a hidden/system namespace begins with a .
            is_hidden = True
            # remove the empty string
            dot_split = dot_split[1:]

            if len(dot_split) == 0:
                carb.log_warn(f"Tag has an empty system namespace: {tag}")
                return False
            elif len(dot_split) == 1:
                namespace = ""
                hidden_label = "." + dot_split[0]
            else:
                # a namespace may contain multiple .'s, but the system namespace may not
                hidden_label = "." + dot_split[-1]
                namespace = ".".join(dot_split[:-1])
        else:
            if full_namespace == "":
                namespace = ""
            else:
                namespace = ".".join(dot_split)

        if not key:
            return False

        if is_hidden and not allow_hidden_tags:
            return False

        return {
            "hidden_label": hidden_label,
            "namespace": namespace,
            "key": key,
            "value": val,
            "full_namespace": full_namespace,
        }

    def ordered_tag_list(self, tag_dict_list):
        """ Returns a sorted (user tags, then generated in order of value) list of unique tags.
        """
        tag_tuples = []
        namespaces = list(set([d["tag_namespace"] for d in tag_dict_list]))

        user_namespaces = []
        generated_namespaces = []
        excluded_namespaces = []
        for n in namespaces:
            if n[0] != ".":
                user_namespaces.append(n)
            elif n[0] == "." and n.split(".")[-1] == "generated":
                generated_namespaces.append(n)
            elif n[0] == "." and n.split(".")[-1] == "excluded":
                excluded_namespaces.append(n)

        user_namespaces.sort()
        generated_namespaces.sort()

        for gen_namespace in generated_namespaces:
            stripped = gen_namespace[1 : -(1 + len("generated"))]
            if stripped not in user_namespaces:
                user_namespaces.append(stripped)

        def safe_int(value):
            try:
                return int(value)
            except ValueError:
                return 0

        # for each namespace, add user tags
        # then add generated tags (if they are not already added or excluded) sorted by value
        for n in user_namespaces:
            g = "." + n + ".generated"
            e = "." + n + ".excluded"

            # tags will contain tuples (name, value)
            tags = []
            excluded = []
            for tag_dict in tag_dict_list:
                if tag_dict["tag_namespace"] == n:
                    tags.append((tag_dict["name"], 100))
                elif tag_dict["tag_namespace"] == e and tag_dict["name"] not in tags:
                    excluded.append(tag_dict["name"])

            for tag_dict in tag_dict_list:
                if tag_dict["tag_namespace"] == g:
                    if tag_dict["name"] not in tags and tag_dict["name"] not in excluded:
                        tags.append((tag_dict["name"], safe_int(tag_dict["value"])))

            if len(tags) > 0:
                # sort by name, then by value descending
                tags.sort(key=lambda x: x[0])
                tag_tuples += sorted(tags, key=lambda x: -x[1])

        tags = []
        for t in tag_tuples:
            # prevent duplicates
            if t[0] not in tags:
                tags.append(t[0])

        return tags

    def unpack_raw_tags(self, tag_dict_list):
        """ Changes a list of strings into a dictionary.

            Example tag: {'name': 'notebook', 'tag_namespace': 'appearance', 'value': ''}
            Example tag: {'name': 'book', 'tag_namespace': 'appearance.generated', 'value': '50'}
            Returns a dictionary like:
            {
                "appearance": {
                    "car": "",
                    ".generated": { "truck": 99 },
                    ".excluded": { "banana": "" }
                }
            }
        """
        tags = {}
        namespace = None

        for tag_dict in tag_dict_list:
            # join the namespace to the name and value of the tag.
            # Examples: appearance.notebook, appearance.generated.book=10
            if tag_dict["tag_namespace"] == "":
                tag = tag_dict["name"]
            else:
                tag = ".".join([tag_dict["tag_namespace"], tag_dict["name"]])

            if "value" in tag_dict and tag_dict["value"]:
                tag = f"{tag}={tag_dict['value']}"

            # recover namespace, key, value from single string
            tag = self.validate_tag(tag)
            if tag:
                hidden_label = tag["hidden_label"]
                namespace = tag["namespace"]
                key = tag["key"]
                val = tag["value"]
                if len(val) > 50:
                    val = "(value too long to display)"

                if namespace not in tags:
                    tags[namespace] = {}
                if hidden_label and hidden_label not in tags[namespace]:
                    tags[namespace][hidden_label] = {}
                if hidden_label:
                    tags[namespace][hidden_label][key] = val
                else:
                    tags[namespace][key] = val

        return tags


def get_tagging_instance():
    return OmniKitTagging.static_instance
