communication_server

class BaseCommunicator:
    def __init__(self):
        self.process = None
        self.websocket_server = None
        self.websocket_rpc = None
        self.exit_code = None
        self._connected_client = None

    @property
    def server_is_running(self):
        if self.websocket_server is None:
            return False
        return self.websocket_server.server_is_running

    def _windows_file_process(self, src_dst_mapping, to_remove):
        """Windows specific file processing asking for admin permissions.

        It is required to have administration permissions to modify plugin
        files in TVPaint installation folder.

        Method requires `pywin32` python module.

        Args:
            src_dst_mapping (list, tuple, set): Mapping of source file to
                destination. Both must be full path. Each item must be iterable
                of size 2 `(C:/src/file.dll, C:/dst/file.dll)`.
            to_remove (list): Fullpath to files that should be removed.
        """

        import pythoncom
        from win32comext.shell import shell

        # Create temp folder where plugin files are temporary copied
        # - reason is that copy to TVPaint requires administartion permissions
        #   but admin may not have access to source folder
        tmp_dir = os.path.normpath(
            tempfile.mkdtemp(prefix="tvpaint_copy_")
        )

        # Copy source to temp folder and create new mapping
        dst_folders = collections.defaultdict(list)
        new_src_dst_mapping = []
        for old_src, dst in src_dst_mapping:
            new_src = os.path.join(tmp_dir, os.path.split(old_src)[1])
            shutil.copy(old_src, new_src)
            new_src_dst_mapping.append((new_src, dst))

        for src, dst in new_src_dst_mapping:
            src = os.path.normpath(src)
            dst = os.path.normpath(dst)
            dst_filename = os.path.basename(dst)
            dst_folder_path = os.path.dirname(dst)
            dst_folders[dst_folder_path].append((dst_filename, src))

        # create an instance of IFileOperation
        fo = pythoncom.CoCreateInstance(
            shell.CLSID_FileOperation,
            None,
            pythoncom.CLSCTX_ALL,
            shell.IID_IFileOperation
        )
        # Add delete command to file operation object
        for filepath in to_remove:
            item = shell.SHCreateItemFromParsingName(
                filepath, None, shell.IID_IShellItem
            )
            fo.DeleteItem(item)

        # here you can use SetOperationFlags, progress Sinks, etc.
        for folder_path, items in dst_folders.items():
            # create an instance of IShellItem for the target folder
            folder_item = shell.SHCreateItemFromParsingName(
                folder_path, None, shell.IID_IShellItem
            )
            for _dst_filename, source_file_path in items:
                # create an instance of IShellItem for the source item
                copy_item = shell.SHCreateItemFromParsingName(
                    source_file_path, None, shell.IID_IShellItem
                )
                # queue the copy operation
                fo.CopyItem(copy_item, folder_item, _dst_filename, None)

        # commit
        fo.PerformOperations()

        # Remove temp folder
        shutil.rmtree(tmp_dir)

    def _prepare_windows_plugin(self, launch_args):
        """Copy plugin to TVPaint plugins and set PATH to dependencies.

        Check if plugin in TVPaint's plugins exist and match to plugin
        version to current implementation version. Based on 64-bit or 32-bit
        version of the plugin. Path to libraries required for plugin is added
        to PATH variable.
        """

        host_executable = launch_args[0]
        executable_file = os.path.basename(host_executable)
        if "64bit" in executable_file:
            subfolder = "windows_x64"
        elif "32bit" in executable_file:
            subfolder = "windows_x86"
        else:
            raise ValueError(
                "Can't determine if executable "
                "leads to 32-bit or 64-bit TVPaint!"
            )

        plugin_files_path = get_plugin_files_path()
        # Folder for right windows plugin files
        source_plugins_dir = os.path.join(plugin_files_path, subfolder)

        # Path to libraries (.dll) required for plugin library
        # - additional libraries can be copied to TVPaint installation folder
        #   (next to executable) or added to PATH environment variable
        additional_libs_folder = os.path.join(
            source_plugins_dir,
            "additional_libraries"
        )
        additional_libs_folder = additional_libs_folder.replace("\\", "/")
        if (
            os.path.exists(additional_libs_folder)
            and additional_libs_folder not in os.environ["PATH"]
        ):
            os.environ["PATH"] += (os.pathsep + additional_libs_folder)

        # Path to TVPaint's plugins folder (where we want to add our plugin)
        host_plugins_path = os.path.join(
            os.path.dirname(host_executable),
            "plugins"
        )

        # Files that must be copied to TVPaint's plugin folder
        plugin_dir = os.path.join(source_plugins_dir, "plugin")

        to_copy = []
        to_remove = []
        # Remove old plugin name
        deprecated_filepath = os.path.join(
            host_plugins_path, "AvalonPlugin.dll"
        )
        if os.path.exists(deprecated_filepath):
            to_remove.append(deprecated_filepath)

        for filename in os.listdir(plugin_dir):
            src_full_path = os.path.join(plugin_dir, filename)
            dst_full_path = os.path.join(host_plugins_path, filename)
            if dst_full_path in to_remove:
                to_remove.remove(dst_full_path)

            if (
                not os.path.exists(dst_full_path)
                or not filecmp.cmp(src_full_path, dst_full_path)
            ):
                to_copy.append((src_full_path, dst_full_path))

        # Skip copy if everything is done
        if not to_copy and not to_remove:
            return

        # Try to copy
        try:
            self._windows_file_process(to_copy, to_remove)
        except Exception:
            log.error("Plugin copy failed", exc_info=True)

        # Validate copy was done
        invalid_copy = []
        for src, dst in to_copy:
            if not os.path.exists(dst) or not filecmp.cmp(src, dst):
                invalid_copy.append((src, dst))

        # Validate delete was dones
        invalid_remove = []
        for filepath in to_remove:
            if os.path.exists(filepath):
                invalid_remove.append(filepath)

        if not invalid_remove and not invalid_copy:
            return

        msg_parts = []
        if invalid_remove:
            msg_parts.append(
                "Failed to remove files: {}".format(", ".join(invalid_remove))
            )

        if invalid_copy:
            _invalid = [
                "\"{}\" -> \"{}\"".format(src, dst)
                for src, dst in invalid_copy
            ]
            msg_parts.append(
                "Failed to copy files: {}".format(", ".join(_invalid))
            )
        raise RuntimeError(" & ".join(msg_parts))

    def _launch_tv_paint(self, launch_args):
        flags = (
            subprocess.DETACHED_PROCESS
            | subprocess.CREATE_NEW_PROCESS_GROUP
        )
        env = os.environ.copy()
        # Remove QuickTime from PATH on windows
        # - quicktime overrides TVPaint's ffmpeg encode/decode which may
        #   cause issues on loading
        if platform.system().lower() == "windows":
            new_path = []
            for path in env["PATH"].split(os.pathsep):
                if path and "quicktime" not in path.lower():
                    new_path.append(path)
            env["PATH"] = os.pathsep.join(new_path)

        kwargs = {
            "env": env,
            "creationflags": flags
        }
        self.process = subprocess.Popen(launch_args, **kwargs)

    def _create_routes(self):
        self.websocket_rpc = BaseTVPaintRpc(
            self, loop=self.websocket_server.loop
        )
        self.websocket_server.add_route(
            "*", "/", self.websocket_rpc.handle_request
        )

    def _start_webserver(self):
        self.websocket_server.start()
        # Make sure RPC is using same loop as websocket server
        while not self.websocket_server.server_is_running:
            time.sleep(0.1)

    def _stop_webserver(self):
        self.websocket_server.stop()

    def _exit(self, exit_code=None):
        self._stop_webserver()
        if exit_code is not None:
            self.exit_code = exit_code

        if self.exit_code is None:
            self.exit_code = 0

    def stop(self):
        """Stop communication and currently running python process."""
        log.info("Stopping communication")
        self._exit()

    def launch(self, launch_args):
        """Prepare all required data and launch host.

        First is prepared websocket server as communication point for host,
        when server is ready to use host is launched as subprocess.
        """
        if platform.system().lower() == "windows":
            self._prepare_windows_plugin(launch_args)

        # Launch TVPaint and the websocket server.
        log.info("Launching TVPaint")
        self.websocket_server = WebSocketServer()

        self._create_routes()
        server_url = f"ws://localhost:{self.websocket_server.port}"
        os.environ["AYON_RPC_URL"] = server_url
        log.info(f"Added request handler for url: {server_url}")

        self._start_webserver()

        # Start TVPaint when server is running
        self._launch_tv_paint(launch_args)

        log.info("Waiting for client connection")
        while True:
            if self.process.poll() is not None:
                log.debug("Host process is not alive. Exiting")
                self._exit(1)
                return

            if self.websocket_rpc.client_connected():
                log.info("Client has connected")
                break
            time.sleep(0.5)

        self._on_client_connect()

        emit_event("application.launched")

    def _on_client_connect(self):
        self._initial_textfile_write()

    def _initial_textfile_write(self):
        """Show popup about Write to file at start of TVPaint."""
        tmp_file = tempfile.NamedTemporaryFile(
            mode="w", prefix="a_tvp_", suffix=".txt", delete=False
        )
        tmp_file.close()
        tmp_filepath = tmp_file.name.replace("\\", "/")
        george_script = (
            "tv_writetextfile \"strict\" \"append\" \"{}\" \"empty\""
        ).format(tmp_filepath)

        result = CommunicationWrapper.execute_george(george_script)

        # Remote the file
        os.remove(tmp_filepath)

        if result is None:
            log.warning(
                "Host was probably closed before plugin was initialized."
            )
        elif result.lower() == "forbidden":
            log.warning("User didn't confirm saving files.")

    def _client(self):
        if not self.websocket_rpc:
            log.warning("Communicator's server did not start yet.")
            return None

        for client in self.websocket_rpc.clients:
            if not client.ws.closed:
                return client
        log.warning("Client is not yet connected to Communicator.")
        return None

    def client(self):
        if not self._connected_client or self._connected_client.ws.closed:
            self._connected_client = self._client()
        return self._connected_client

    def send_request(self, method, params=None):
        client = self.client()
        if not client:
            return

        return self.websocket_rpc.send_request(
            client, method, params
        )

    def send_notification(self, method, params=None):
        client = self.client()
        if not client:
            return

        self.websocket_rpc.send_notification(
            client, method, params
        )

    def execute_george(self, george_script):
        """Execute passed goerge script in TVPaint."""
        return self.send_request(
            "execute_george", [george_script]
        )

    def execute_george_through_file(self, george_script):
        """Execute george script with temp file.

        Allows to execute multiline george script without stopping websocket
        client.

        On windows make sure script does not contain paths with backwards
        slashes in paths, TVPaint won't execute properly in that case.

        Args:
            george_script (str): George script to execute. May be multilined.
        """
        temporary_file = tempfile.NamedTemporaryFile(
            mode="w", prefix="a_tvp_", suffix=".grg", delete=False
        )
        temporary_file.write(george_script)
        temporary_file.close()
        temp_file_path = temporary_file.name.replace("\\", "/")
        self.execute_george("tv_runscript {}".format(temp_file_path))
        os.remove(temp_file_path)

class CommunicationWrapper:
    # TODO add logs and exceptions
    communicator = None

    log = logging.getLogger("CommunicationWrapper")

    @classmethod
    def create_qt_communicator(cls, *args, **kwargs):
        """Create communicator for Artist usage."""
        communicator = QtCommunicator(*args, **kwargs)
        cls.set_communicator(communicator)
        return communicator

    @classmethod
    def set_communicator(cls, communicator):
        if not cls.communicator:
            cls.communicator = communicator
        else:
            cls.log.warning("Communicator was set multiple times.")

    @classmethod
    def client(cls):
        if not cls.communicator:
            return None
        return cls.communicator.client()

    @classmethod
    def execute_george(cls, george_script):
        """Execute passed goerge script in TVPaint."""
        if not cls.communicator:
            return
        return cls.communicator.execute_george(george_script)

class MainThreadItem:
    """Structure to store information about callback in main thread.

    Item should be used to execute callback in main thread which may be needed
    for execution of Qt objects.

    Item store callback (callable variable), arguments and keyword arguments
    for the callback. Item hold information about it's process.
    """
    not_set = object()
    sleep_time = 0.1

    def __init__(self, callback, *args, **kwargs):
        self.done = False
        self.exception = self.not_set
        self.result = self.not_set
        self.callback = callback
        self.args = args
        self.kwargs = kwargs

    def execute(self):
        """Execute callback and store its result.

        Method must be called from main thread. Item is marked as `done`
        when callback execution finished. Store output of callback of exception
        information when callback raises one.
        """
        log.debug("Executing process in main thread")
        if self.done:
            log.warning("- item is already processed")
            return

        callback = self.callback
        args = self.args
        kwargs = self.kwargs
        log.info("Running callback: {}".format(str(callback)))
        try:
            result = callback(*args, **kwargs)
            self.result = result

        except Exception as exc:
            self.exception = exc

        finally:
            self.done = True

    def wait(self):
        """Wait for result from main thread.

        This method stops current thread until callback is executed.

        Returns:
            object: Output of callback. May be any type or object.

        Raises:
            Exception: Reraise any exception that happened during callback
                execution.
        """
        while not self.done:
            time.sleep(self.sleep_time)

        if self.exception is self.not_set:
            return self.result
        raise self.exception

    async def async_wait(self):
        """Wait for result from main thread.

        Returns:
            object: Output of callback. May be any type or object.

        Raises:
            Exception: Reraise any exception that happened during callback
                execution.
        """
        while not self.done:
            await asyncio.sleep(self.sleep_time)

        if self.exception is self.not_set:
            return self.result
        raise self.exception

class QtCommunicator(BaseCommunicator):
    label = os.getenv("AYON_MENU_LABEL") or "AYON"
    title = "{} Tools".format(label)
    menu_definitions = {
        "title": title,
        "menu_items": [
            {
                "callback": "workfiles_tool",
                "label": "Workfiles",
                "help": "Open workfiles tool"
            }, {
                "callback": "loader_tool",
                "label": "Load",
                "help": "Open loader tool"
            }, {
                "callback": "scene_inventory_tool",
                "label": "Scene inventory",
                "help": "Open scene inventory tool"
            }, {
                "callback": "publish_tool",
                "label": "Publish",
                "help": "Open publisher"
            }, {
                "callback": "library_loader_tool",
                "label": "Library",
                "help": "Open library loader tool"
            }, {
                "callback": "experimental_tools",
                "label": "Experimental tools",
                "help": "Open experimental tools dialog"
            }
        ]
    }

    def __init__(self, qt_app):
        super().__init__()
        self.callback_queue = collections.deque()
        self.qt_app = qt_app

    def _create_routes(self):
        self.websocket_rpc = QtTVPaintRpc(
            self, loop=self.websocket_server.loop
        )
        self.websocket_server.add_route(
            "*", "/", self.websocket_rpc.handle_request
        )

    def execute_in_main_thread(self, main_thread_item, wait=True):
        """Add `MainThreadItem` to callback queue and wait for result."""
        self.callback_queue.append(main_thread_item)
        if wait:
            return main_thread_item.wait()
        return

    async def async_execute_in_main_thread(self, main_thread_item, wait=True):
        """Add `MainThreadItem` to callback queue and wait for result."""
        self.callback_queue.append(main_thread_item)
        if wait:
            return await main_thread_item.async_wait()

    def main_thread_listen(self):
        """Get last `MainThreadItem` from queue.

        Must be called from main thread.

        Method checks if host process is still running as it may cause
        issues if not.
        """
        # check if host still running
        if self.process.poll() is not None:
            self._exit()
            return None

        if self.callback_queue:
            return self.callback_queue.popleft()
        return None

    def _on_client_connect(self):
        super()._on_client_connect()
        self._build_menu()

    def _build_menu(self):
        self.send_request(
            "define_menu", [self.menu_definitions]
        )

    def _exit(self, *args, **kwargs):
        super()._exit(*args, **kwargs)
        emit_event("application.exit")
        self.qt_app.exit(self.exit_code)

class QtTVPaintRpc(BaseTVPaintRpc):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        from ayon_core.tools.utils import host_tools
        self.tools_helper = host_tools.HostToolsHelper()

        route_name = self.route_name

        # Register methods
        self.add_methods(
            (route_name, self.workfiles_tool),
            (route_name, self.loader_tool),
            (route_name, self.publish_tool),
            (route_name, self.scene_inventory_tool),
            (route_name, self.library_loader_tool),
            (route_name, self.experimental_tools)
        )

    # Panel routes for tools
    async def workfiles_tool(self):
        log.info("Triggering Workfile tool")
        item = MainThreadItem(self.tools_helper.show_workfiles)
        self._execute_in_main_thread(item, wait=False)
        return

    async def loader_tool(self):
        log.info("Triggering Loader tool")
        item = MainThreadItem(self.tools_helper.show_loader)
        self._execute_in_main_thread(item, wait=False)
        return

    async def publish_tool(self):
        log.info("Triggering Publish tool")
        item = MainThreadItem(self.tools_helper.show_publisher_tool)
        self._execute_in_main_thread(item, wait=False)
        return

    async def scene_inventory_tool(self):
        """Open Scene Inventory tool.

        Function can't confirm if tool was opened becauise one part of
        SceneInventory initialization is calling websocket request to host but
        host can't response because is waiting for response from this call.
        """
        log.info("Triggering Scene inventory tool")
        item = MainThreadItem(self.tools_helper.show_scene_inventory)
        # Do not wait for result of callback
        self._execute_in_main_thread(item, wait=False)
        return

    async def library_loader_tool(self):
        log.info("Triggering Library loader tool")
        item = MainThreadItem(self.tools_helper.show_library_loader)
        self._execute_in_main_thread(item, wait=False)
        return

    async def experimental_tools(self):
        log.info("Triggering Library loader tool")
        item = MainThreadItem(self.tools_helper.show_experimental_tools_dialog)
        self._execute_in_main_thread(item, wait=False)
        return

    async def _async_execute_in_main_thread(self, item, **kwargs):
        await self.communication_obj.async_execute_in_main_thread(
            item, **kwargs
        )

    def _execute_in_main_thread(self, item, **kwargs):
        return self.communication_obj.execute_in_main_thread(item, **kwargs)

class WebsocketServerThread(threading.Thread):
    """ Listener for websocket rpc requests.

        It would be probably better to "attach" this to main thread (as for
        example Harmony needs to run something on main thread), but currently
        it creates separate thread and separate asyncio event loop
    """
    def __init__(self, module, port, loop):
        super(WebsocketServerThread, self).__init__()
        self.is_running = False
        self.server_is_running = False
        self.port = port
        self.module = module
        self.loop = loop
        self.runner = None
        self.site = None
        self.tasks = []

    def run(self):
        self.is_running = True

        try:
            log.debug("Starting websocket server")

            self.loop.run_until_complete(self.start_server())

            log.info(
                "Running Websocket server on URL:"
                " \"ws://localhost:{}\"".format(self.port)
            )

            asyncio.ensure_future(self.check_shutdown(), loop=self.loop)

            self.server_is_running = True
            self.loop.run_forever()

        except Exception:
            log.warning(
                "Websocket Server service has failed", exc_info=True
            )
        finally:
            self.server_is_running = False
            # optional
            self.loop.close()

        self.is_running = False
        log.info("Websocket server stopped")

    async def start_server(self):
        """ Starts runner and TCPsite """
        self.runner = web.AppRunner(self.module.app)
        await self.runner.setup()
        self.site = web.TCPSite(self.runner, "localhost", self.port)
        await self.site.start()

    def stop(self):
        """Sets is_running flag to false, 'check_shutdown' shuts server down"""
        self.is_running = False

    async def check_shutdown(self):
        """ Future that is running and checks if server should be running
            periodically.
        """
        while self.is_running:
            while self.tasks:
                task = self.tasks.pop(0)
                log.debug("waiting for task {}".format(task))
                await task
                log.debug("returned value {}".format(task.result))

            await asyncio.sleep(0.5)

        log.debug("## Server shutdown started")

        await self.site.stop()
        log.debug("# Site stopped")
        await self.runner.cleanup()
        log.debug("# Server runner stopped")
        tasks = [
            task for task in asyncio.all_tasks()
            if task is not asyncio.current_task()
        ]
        list(map(lambda task: task.cancel(), tasks))  # cancel all the tasks
        results = await asyncio.gather(*tasks, return_exceptions=True)
        log.debug(f"Finished awaiting cancelled tasks, results: {results}...")
        await self.loop.shutdown_asyncgens()
        # to really make sure everything else has time to stop
        await asyncio.sleep(0.07)
        self.loop.stop()