rpc_server

class AyonLocalHost(Route):
    """Handle Menu calls from JS part of ComfyUI plugin."""

    qt_thread: QThread_interface = None

    def init(self, **kwargs):
        """Override of init method. Can return anything."""
        return kwargs

    @classmethod
    def register_qrpc_manager(cls, qrpc: QThread_interface):
        cls.qt_thread = qrpc

    @decorators.proxy
    async def pingAyonMenu(self, message: str) -> str:  # noqa: PLR6301, N802
        """Returns message sent from server."""
        return message

    @decorators.proxy
    async def requestToolByName(self, tool_name: str) -> None:  # noqa: N802
        """Schedule tool in thread."""
        log.debug(f"origin {self.socket.request.headers}")  # noqa: G004
        if self.qt_thread:
            self.qt_thread.schedule(show_tool_by_name, tool_name)
            return f"{tool_name} scheduled in qt_thread"
        return tool_name

    @decorators.proxy
    async def requestSaveByName(  # noqa: N802
        self, file_name: str, workfile_contents: str
    ) -> None:
        """Schedule saving workfile in thread.

        File name is without extension (as represented in ComfyUI tab)
        """
        log.debug(f"origin {self.socket.request.headers}")  # noqa: G004
        log.debug("Saving workfile in place")
        if self.qt_thread:
            self.qt_thread.schedule(
                overwrite_workfile, file_name, workfile_contents
            )
            return f"{file_name} save scheduled in qt_thread"
        return file_name

class RPCServerThread(Thread):
    """Manages event loop for the server that recieves messages.

    In particular, deal with the JavaScript part of the plugin.
    """

    _instance: ClassVar[RPCServerThread] = None

    def __init__(
        self,
        port: int = 55056,
        qthread: QThread_interface = None,
        *,
        https: bool = False,
    ):
        self._app = None
        self._port = port
        self._setup = False
        self._is_running = False
        self._is_https = https
        self._loop: asyncio.AbstractEventLoop = None
        self._shutdown_event = None
        self._had_first_client = False

        self._origin = pull_origin_from_settings()
        log.info(self._origin)
        AyonLocalHost.register_qrpc_manager(qthread)
        self._qt_thread = qthread
        self.__class__._instance = self  # noqa: SLF001

        super().__init__()

    @classmethod
    def get_thread(cls) -> RPCServerThread:
        """Return stored thread."""
        return cls._instance

    def setup_server(self) -> None:
        """Set up server, do not start it yet."""
        self._app = aiohttp.web.Application()
        self._app.router.add_route("*", "/ws/", WebSocketAsync)

        WebSocketAsync.add_route("ayonComfyUI", AyonLocalHost)
        self._setup = True

    def run(self) -> None:
        """Do server run."""
        if not self.is_set_up:
            return
        self._loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._loop)
        self._shutdown_event = asyncio.Event()
        self._loop.run_until_complete(self.async_run())

    async def async_run(self) -> None:
        """Run server async."""
        runner = aiohttp.web.AppRunner(self._app)
        await runner.setup()
        site = aiohttp.web.TCPSite(runner, "localhost", self._port)
        await site.start()
        log.info(
            f"Websocket Server running on ws://localhost:{self._port}/ws/"  # noqa: G004
        )

        # Shutdown conditional
        # await self._shutdown_event.wait()
        await self.check_any_origin_client()

        try:
            # Close all websocket conns
            await WebSocketAsync.close()
        except TypeError as e:
            log.info(f"Error {e} ignored. Already closed Websocket.")  # noqa: G004

        # stopping server and clean up
        await site.stop()
        await runner.cleanup()
        log.info("RPC webocket thread stopped and cleaned up.")

    async def check_any_origin_client(self) -> None:
        """If shutdown flag isn't set, check for valid frontend conns."""
        while not self._shutdown_event.is_set():
            if (
                get_client_from_origin(self._origin) is None
                and self._had_first_client
            ):
                self._qt_thread.sig_onfrontendcon_fail.emit()
                break
            await asyncio.sleep(0.5)

    def flag_valid_client(self) -> None:
        """Notify server a verified origin client connection has been made."""
        self._had_first_client = True

    @property
    def is_set_up(self) -> bool:
        """Return whether the server is set up.

        This doesn't mean that the server is running,
        but it means that the server is ready to run.
        """
        return self._setup

    @property
    def loop(self) -> asyncio.AbstractEventLoop:
        """Return contained loop."""
        return self._loop

    def stop(self) -> None:
        """Stop server from running."""
        self._loop.call_soon_threadsafe(self._shutdown_event.set)