APIClient-Agent/app/ui/websocket_panel.py

"""APIClient - Agent — WebSocket client panel."""
import asyncio
import queue
import time

from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QLineEdit,
    QPushButton, QTextEdit, QLabel, QSplitter
)
from PyQt6.QtCore import QThread, pyqtSignal, Qt
from PyQt6.QtGui import QFont

from app.ui.theme import Colors, restyle


class WsWorker(QThread):
    """Runs the asyncio WebSocket loop in a background thread.

    Messages to send are passed via a thread-safe queue.
    Received messages and connection events are emitted as Qt signals.
    """

    message_received = pyqtSignal(str)
    connected        = pyqtSignal()
    disconnected     = pyqtSignal(str)

    def __init__(self, url: str):
        super().__init__()
        self.url       = url
        self._send_q:  queue.Queue = queue.Queue()
        self._running: bool = False
        self._loop:    asyncio.AbstractEventLoop | None = None

    # ── Called from UI thread ─────────────────────────────────────────────────

    def send(self, message: str):
        """Thread-safe: enqueue a message to be sent."""
        self._send_q.put(message)

    def stop(self):
        self._running = False
        if self._loop and self._loop.is_running():
            self._loop.call_soon_threadsafe(self._loop.stop)

    # ── Worker thread ─────────────────────────────────────────────────────────

    def run(self):
        try:
            import websockets
        except ImportError:
            self.disconnected.emit("WebSocket support not installed.\nRun: pip install websockets")
            return

        self._loop    = asyncio.new_event_loop()
        self._running = True
        try:
            self._loop.run_until_complete(self._connect(websockets))
        except Exception as e:
            self.disconnected.emit(str(e))
        finally:
            self._loop.close()
            self._loop = None

    async def _connect(self, websockets):
        try:
            async with websockets.connect(self.url) as ws:
                self.connected.emit()
                while self._running:
                    # Drain outbound queue
                    while not self._send_q.empty():
                        try:
                            msg = self._send_q.get_nowait()
                            await ws.send(msg)
                        except queue.Empty:
                            break

                    # Non-blocking receive with timeout
                    try:
                        msg = await asyncio.wait_for(ws.recv(), timeout=0.1)
                        self.message_received.emit(f"← {msg}")
                    except asyncio.TimeoutError:
                        continue
                    except Exception:
                        break

        except Exception as e:
            self.disconnected.emit(str(e))
            return

        self.disconnected.emit("Connection closed")


class WebSocketPanel(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        layout = QVBoxLayout(self)
        layout.setContentsMargins(12, 12, 12, 12)
        layout.setSpacing(8)

        # ── Connection bar ────────────────────────────────────────────────────
        conn_row = QHBoxLayout()
        url_label = QLabel("URL:")
        url_label.setObjectName("fieldLabel")
        self.url_input = QLineEdit()
        self.url_input.setObjectName("urlBar")
        self.url_input.setPlaceholderText("ws://localhost:8080/socket")
        self.connect_btn = QPushButton("Connect")
        self.connect_btn.setObjectName("accent")
        self.connect_btn.setFixedWidth(100)
        self.connect_btn.clicked.connect(self._toggle_connection)
        self.status_label = QLabel("● Disconnected")
        self.status_label.setObjectName("statusErr")
        conn_row.addWidget(url_label)
        conn_row.addWidget(self.url_input, 1)
        conn_row.addWidget(self.connect_btn)
        conn_row.addWidget(self.status_label)
        layout.addLayout(conn_row)

        splitter = QSplitter(Qt.Orientation.Vertical)

        # Message log
        self.log = QTextEdit()
        self.log.setReadOnly(True)
        self.log.setFont(QFont("JetBrains Mono, Fira Code, Consolas, monospace", 10))
        splitter.addWidget(self.log)

        # Send area
        send_w = QWidget()
        send_layout = QVBoxLayout(send_w)
        send_layout.setContentsMargins(0, 4, 0, 0)
        send_layout.setSpacing(6)

        self.send_input = QTextEdit()
        self.send_input.setPlaceholderText("Type message to send…")
        self.send_input.setMaximumHeight(80)
        self.send_input.setFont(QFont("JetBrains Mono, Fira Code, Consolas, monospace", 10))

        send_btn_row = QHBoxLayout()
        self.send_btn = QPushButton("Send")
        self.send_btn.setObjectName("sendBtn")
        self.send_btn.setEnabled(False)
        self.send_btn.setFixedWidth(90)
        self.send_btn.clicked.connect(self._send_message)

        clear_btn = QPushButton("Clear Log")
        clear_btn.setObjectName("ghost")
        clear_btn.clicked.connect(self.log.clear)

        send_btn_row.addWidget(self.send_btn)
        send_btn_row.addWidget(clear_btn)
        send_btn_row.addStretch()

        send_layout.addWidget(self.send_input)
        send_layout.addLayout(send_btn_row)
        splitter.addWidget(send_w)

        splitter.setSizes([350, 150])
        layout.addWidget(splitter)

        self._worker: WsWorker | None = None

    # ── Slots ─────────────────────────────────────────────────────────────────

    def _toggle_connection(self):
        if self._worker and self._worker.isRunning():
            self._worker.stop()
            self._worker.wait(1000)
            self._worker = None
            self._set_disconnected("Disconnected by user")
        else:
            url = self.url_input.text().strip()
            if not url:
                return
            self._worker = WsWorker(url)
            self._worker.connected.connect(self._on_connected)
            self._worker.disconnected.connect(self._on_disconnected)
            self._worker.message_received.connect(self._on_message)
            self._worker.start()
            self.connect_btn.setText("Connecting…")
            self.connect_btn.setEnabled(False)
            restyle(self.status_label, "statusWarn")
            self.status_label.setText("● Connecting…")

    def _on_connected(self):
        restyle(self.status_label, "statusOk")
        self.status_label.setText("● Connected")
        self.connect_btn.setText("Disconnect")
        self.connect_btn.setEnabled(True)
        self.send_btn.setEnabled(True)
        self._log("── Connected ──", Colors.SUCCESS)

    def _on_disconnected(self, reason: str):
        self._set_disconnected(reason)
        self._log(f"── {reason} ──", Colors.ERROR)

    def _set_disconnected(self, reason: str = "Disconnected"):
        restyle(self.status_label, "statusErr")
        self.status_label.setText(f"● {reason}")
        self.connect_btn.setText("Connect")
        self.connect_btn.setEnabled(True)
        self.send_btn.setEnabled(False)

    def _on_message(self, msg: str):
        self._log(msg, Colors.INFO)

    def _send_message(self):
        msg = self.send_input.toPlainText().strip()
        if msg and self._worker:
            self._worker.send(msg)
            self._log(f"→ {msg}", Colors.WARNING)
            self.send_input.clear()

    def _log(self, text: str, color: str = "#D4D4D4"):
        ts = time.strftime("%H:%M:%S")
        safe_text = text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
        self.log.append(f'<span style="color:{color}">[{ts}] {safe_text}</span>')