import omni.kit.commands
import omni.kit.tagging
import omni.client


class AddTagCommand(omni.kit.commands.Command):
    """
    Add a tag to an asset **Command**.
    Fails if we are not connected to the tagging server.

    Args:
        asset_path (str): Path to the asset. Should start with omniverse: or it will just return
        tag (str): packed string tag, such as appearance.Car, or .appearance.generated.Truck=99
        excluded (str): excluded version of the tag (or empty string) to remove if present
    """

    def __init__(self, asset_path: str, tag: str, excluded: str = ""):
        self._asset_path = asset_path
        self._tag = tag
        self._excluded = excluded

    def do(self):
        broken_url = omni.client.break_url(self._asset_path)
        if broken_url.scheme != "omniverse":
            return
        tagging = omni.kit.tagging.get_tagging_instance()
        tagging.tags_action(self._asset_path, self._tag, action="add")
        if self._excluded:
            tagging.tags_action(self._asset_path, self._excluded, action="remove")

    def undo(self):
        broken_url = omni.client.break_url(self._asset_path)
        if broken_url.scheme != "omniverse":
            return
        tagging = omni.kit.tagging.get_tagging_instance()
        tagging.tags_action(self._asset_path, self._tag, action="remove")
        if self._excluded:
            tagging.tags_action(self._asset_path, self._excluded, action="add")


class RemoveTagCommand(omni.kit.commands.Command):
    """
    Remove a tag to an asset **Command**.
    Fails if we are not connected to the tagging server.

    Args:
        asset_path (str): Path to the asset. Should start with omniverse: or it will just return
        tag (str): packed string tag, such as appearance.Car, or .appearance.generated.Truck=99
    """

    def __init__(self, asset_path: str, tag: str):
        self._asset_path = asset_path
        self._tag = tag

    def do(self):
        broken_url = omni.client.break_url(self._asset_path)
        if broken_url.scheme != "omniverse":
            return
        tagging = omni.kit.tagging.get_tagging_instance()
        tagging.tags_action(self._asset_path, self._tag, action="remove")

    def undo(self):
        broken_url = omni.client.break_url(self._asset_path)
        if broken_url.scheme != "omniverse":
            return
        tagging = omni.kit.tagging.get_tagging_instance()
        tagging.tags_action(self._asset_path, self._tag, action="add")


class UpdateTagCommand(omni.kit.commands.Command):
    """
    Remove a tag to an asset **Command**.
    Fails if we are not connected to the tagging server.

    Args:
        asset_path (str): Path to the asset. Should start with omniverse: or it will just return
        old_tag (str): packed string tag, such as appearance.Car, or .appearance.generated.Truck=99
        new_tag (str): packed string tag, such as appearance.Car, or .appearance.generated.Truck=99
    """

    def __init__(self, asset_path: str, old_tag: str, new_tag):
        self._asset_path = asset_path
        self._old_tag = old_tag
        self._new_tag = new_tag

    def do(self):
        broken_url = omni.client.break_url(self._asset_path)
        if broken_url.scheme != "omniverse":
            return
        tagging = omni.kit.tagging.get_tagging_instance()
        tagging.tags_action(self._asset_path, old_tag=self._old_tag, action="update", new_tag=self._new_tag)

    def undo(self):
        broken_url = omni.client.break_url(self._asset_path)
        if broken_url.scheme != "omniverse":
            return
        tagging = omni.kit.tagging.get_tagging_instance()
        tagging.tags_action(self._asset_path, old_tag=self._new_tag, action="update", new_tag=self._old_tag)


omni.kit.commands.register_all_commands_in_module(__name__)
