APIClient-Agent/app/ui/main_window.py

"""APIClient - Agent — Main Window."""
from PyQt6.QtWidgets import (
    QMainWindow, QSplitter, QWidget, QVBoxLayout, QHBoxLayout,
    QLabel, QComboBox, QPushButton, QStatusBar, QTabWidget,
    QInputDialog, QMessageBox, QFileDialog, QApplication
)
from PyQt6.QtCore import Qt, QThread, pyqtSignal
from PyQt6.QtGui import QKeySequence, QShortcut

from app.ui.tabs_manager import TabsManager
from app.ui.response_panel import ResponsePanel
from app.ui.sidebar import CollectionsSidebar
from app.ui.theme import Colors, toggle as toggle_theme, is_dark
from app.core import http_client, storage
from app.core.test_runner import run_tests
from app.models import HttpRequest

APP_VERSION = "2.0.0"
APP_NAME    = "APIClient - Agent"


class RequestWorker(QThread):
    finished = pyqtSignal(object, list)

    def __init__(self, req: HttpRequest, variables: dict):
        super().__init__()
        self.req        = req
        self.variables  = variables
        self._cancelled = False

    def run(self):
        if self.req.pre_request_script.strip():
            try:
                exec(self.req.pre_request_script, {"__builtins__": {}})  # noqa: S102
            except Exception:
                pass
        if self._cancelled:
            return
        resp  = http_client.send_request(self.req, self.variables)
        tests = run_tests(self.req.test_script, resp)
        self.finished.emit(resp, tests)

    def cancel(self):
        self._cancelled = True


class EnvBar(QWidget):
    """Top branding + environment selector bar."""

    def __init__(self, parent=None):
        super().__init__(parent)
        self.setObjectName("envBar")
        self.setFixedHeight(46)

        layout = QHBoxLayout(self)
        layout.setContentsMargins(16, 0, 16, 0)
        layout.setSpacing(6)

        brand = QLabel("APIClient")
        brand.setObjectName("brandName")

        sub = QLabel("Agent")
        sub.setObjectName("brandSub")

        layout.addWidget(brand)
        layout.addWidget(sub)
        layout.addStretch()

        env_label = QLabel("ENV")
        env_label.setObjectName("envChip")
        layout.addWidget(env_label)

        self.env_combo = QComboBox()
        self.env_combo.setObjectName("methodCombo")
        self.env_combo.setMinimumWidth(180)
        layout.addWidget(self.env_combo)

        self.manage_btn = QPushButton("Manage")
        self.manage_btn.setObjectName("ghost")
        self.manage_btn.setToolTip("Manage Environments  (Ctrl+E)")
        layout.addWidget(self.manage_btn)

        self.ai_btn = QPushButton("✦ AI")
        self.ai_btn.setObjectName("accent")
        self.ai_btn.setFixedWidth(60)
        self.ai_btn.setToolTip("Toggle AI Chat  (Ctrl+Shift+A)")
        layout.addWidget(self.ai_btn)

        self.theme_btn = QPushButton("◑")
        self.theme_btn.setObjectName("ghost")
        self.theme_btn.setFixedWidth(32)
        self.theme_btn.setToolTip("Toggle Light / Dark Theme")
        layout.addWidget(self.theme_btn)


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle(f"{APP_NAME}  v{APP_VERSION}")
        self.setMinimumSize(1280, 800)
        self._worker: RequestWorker | None = None

        storage.init_db()
        self._build_ui()
        self._build_menu()
        self._build_shortcuts()
        self._update_env_selector()

    # ── UI Construction ───────────────────────────────────────────────────────

    def _build_ui(self):
        root = QWidget()
        root_layout = QVBoxLayout(root)
        root_layout.setContentsMargins(0, 0, 0, 0)
        root_layout.setSpacing(0)

        self.env_bar = EnvBar()
        self.env_bar.env_combo.currentIndexChanged.connect(self._on_env_changed)
        self.env_bar.manage_btn.clicked.connect(self._open_env_dialog)
        self.env_bar.theme_btn.clicked.connect(self._toggle_theme)
        self.env_bar.ai_btn.clicked.connect(self._toggle_ai_chat)
        root_layout.addWidget(self.env_bar)

        splitter = QSplitter(Qt.Orientation.Horizontal)
        splitter.setHandleWidth(1)

        self.sidebar = CollectionsSidebar()
        self.sidebar.request_selected.connect(self._load_request_in_tab)
        splitter.addWidget(self.sidebar)

        self.workspace = QTabWidget()
        self.workspace.setObjectName("workspaceTabs")

        self.workspace.addTab(self._build_http_workspace(), "  HTTP  ")

        from app.ui.websocket_panel import WebSocketPanel
        self.ws_panel = WebSocketPanel()
        self.workspace.addTab(self.ws_panel, "  WebSocket  ")

        from app.ui.mock_server_panel import MockServerPanel
        self.mock_panel = MockServerPanel()
        self.workspace.addTab(self.mock_panel, "  Mock Server  ")

        splitter.addWidget(self.workspace)

        from app.ui.ai_chat_panel import AIChatPanel
        self.chat_panel = AIChatPanel()
        splitter.addWidget(self.chat_panel)
        splitter.setSizes([260, 940, 360])   # give chat panel real size first
        self.chat_panel.hide()               # THEN hide — splitter remembers 360
        self._main_splitter = splitter

        # Wire apply signals
        self.chat_panel.apply_body.connect(lambda c: self._ai_apply("body", c))
        self.chat_panel.apply_params.connect(lambda c: self._ai_apply("params", c))
        self.chat_panel.apply_headers.connect(lambda c: self._ai_apply("headers", c))
        self.chat_panel.apply_test.connect(lambda c: self._ai_apply("test", c))

        root_layout.addWidget(splitter)

        self.setCentralWidget(root)

        self._status_bar = QStatusBar()
        self._status_bar.setFixedHeight(26)
        self.setStatusBar(self._status_bar)
        self._status_bar.showMessage(f"Ready  —  {APP_NAME}  v{APP_VERSION}")

    def _build_http_workspace(self) -> QWidget:
        w = QWidget()
        layout = QVBoxLayout(w)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)

        self.tabs_manager = TabsManager()
        self.tabs_manager.send_requested.connect(self._send)

        self.response_panel = ResponsePanel()

        splitter = QSplitter(Qt.Orientation.Vertical)
        splitter.setHandleWidth(1)
        splitter.addWidget(self.tabs_manager)
        splitter.addWidget(self.response_panel)
        splitter.setSizes([400, 400])
        layout.addWidget(splitter)
        return w

    # ── Menu ──────────────────────────────────────────────────────────────────

    def _build_menu(self):
        mb = self.menuBar()

        file_m = mb.addMenu("File")
        file_m.addAction("New Tab",             self.tabs_manager.new_tab).setShortcut("Ctrl+T")
        file_m.addAction("Close Tab",           self.tabs_manager.close_current_tab).setShortcut("Ctrl+W")
        file_m.addSeparator()
        file_m.addAction("Save to Collection",  self._save_to_collection).setShortcut("Ctrl+S")
        file_m.addSeparator()
        file_m.addAction("Import…",             self._import)
        file_m.addAction("Export Collection…",  self._export)
        file_m.addSeparator()
        file_m.addAction("Quit",                self.close).setShortcut("Ctrl+Q")

        view_m = mb.addMenu("View")
        view_m.addAction("Search Requests",     self._open_search).setShortcut("Ctrl+F")
        view_m.addAction("Toggle Theme",        self._toggle_theme)

        tools_m = mb.addMenu("Tools")
        tools_m.addAction("Environments…",      self._open_env_dialog).setShortcut("Ctrl+E")
        tools_m.addAction("Collection Runner…", self._open_runner)
        tools_m.addAction("Generate Code…",     self._open_code_gen)
        tools_m.addSeparator()
        tools_m.addAction("AI Assistant…",      self._open_ai_assistant)

        mb.addMenu("Help").addAction(
            f"About {APP_NAME}",
            lambda: QMessageBox.about(
                self, APP_NAME,
                f"<b>{APP_NAME}  v{APP_VERSION}</b><br>"
                "Enterprise-grade API testing tool with AI co-pilot.<br><br>"
                "Built with Python + PyQt6"
            )
        )

    def _build_shortcuts(self):
        QShortcut(QKeySequence("Ctrl+Return"), self, self._send_current)
        QShortcut(QKeySequence("Ctrl+T"),      self, self.tabs_manager.new_tab)
        QShortcut(QKeySequence("Ctrl+W"),      self, self.tabs_manager.close_current_tab)
        QShortcut(QKeySequence("Ctrl+S"),      self, self._save_to_collection)
        QShortcut(QKeySequence("Ctrl+F"),      self, self._open_search)
        QShortcut(QKeySequence("Ctrl+E"),          self, self._open_env_dialog)
        QShortcut(QKeySequence("Ctrl+Shift+A"),  self, self._toggle_ai_chat)
        QShortcut(QKeySequence("Escape"),        self, self._cancel_request)
        QShortcut(QKeySequence("Ctrl+Q"),        self, self.close)

    # ── Environment ───────────────────────────────────────────────────────────

    def _update_env_selector(self):
        combo = self.env_bar.env_combo
        combo.blockSignals(True)
        combo.clear()
        combo.addItem("No Environment", None)
        active = storage.get_active_environment()
        for e in storage.get_environments():
            combo.addItem(e.name, e.id)
            if active and e.id == active.id:
                combo.setCurrentIndex(combo.count() - 1)
        combo.blockSignals(False)

    def _on_env_changed(self, _):
        env_id = self.env_bar.env_combo.currentData()
        storage.set_active_environment(env_id)
        self._set_status(f"Environment: {self.env_bar.env_combo.currentText()}")

    def _get_active_variables(self) -> dict:
        env = storage.get_active_environment()
        return env.variables if env else {}

    # ── Sending ───────────────────────────────────────────────────────────────

    def _send_current(self):
        tab = self.tabs_manager.current_tab()
        if tab:
            self._send(tab.get_request())

    def _send(self, req: HttpRequest):
        if not req.url.strip():
            self._set_status("Enter a URL first", error=True)
            return
        if self._worker and self._worker.isRunning():
            self._worker.cancel()
            self._worker.wait(500)

        tab = self.tabs_manager.current_tab()
        if tab:
            tab.request_panel.send_btn.setEnabled(False)
            tab.request_panel.send_btn.setText("Sending…")

        self.response_panel.set_loading(True)
        storage.add_to_history(req)
        self.sidebar.refresh()

        self._worker = RequestWorker(req, self._get_active_variables())
        self._worker.finished.connect(self._on_response)
        self._worker.start()
        self._set_status(f"⟳  {req.method}  {req.url}")
        if self.chat_panel.isVisible():
            self.chat_panel.set_context(req=req, env_vars=self._get_active_variables())

    def _cancel_request(self):
        if self._worker and self._worker.isRunning():
            self._worker.cancel()
            self._worker.wait(500)
            self.response_panel.set_loading(False)
            self._restore_send_btn()
            self._set_status("Request cancelled")

    def _restore_send_btn(self):
        tab = self.tabs_manager.current_tab()
        if tab:
            tab.request_panel.send_btn.setEnabled(True)
            tab.request_panel.send_btn.setText("Send")

    def _on_response(self, resp, tests):
        self.response_panel.set_loading(False)
        self.response_panel.display(resp, tests)
        self._restore_send_btn()
        tab = self.tabs_manager.current_tab()
        req = tab.get_request() if tab else None
        if self.chat_panel.isVisible():
            self.chat_panel.set_context(req=req, resp=resp, env_vars=self._get_active_variables())
        if resp.error:
            self._set_status(f"✗  {resp.error}", error=True)
        else:
            size = resp.size_bytes or len(resp.body.encode())
            self._set_status(
                f"✓  {resp.status} {resp.reason}    {resp.elapsed_ms:.0f} ms    {_fmt_size(size)}"
            )

    # ── Request loading ───────────────────────────────────────────────────────

    def _load_request_in_tab(self, req: HttpRequest):
        self.tabs_manager.load_request_in_new_tab(req)
        if self.chat_panel.isVisible():
            self.chat_panel.set_context(req=req, env_vars=self._get_active_variables())

    # ── Save ──────────────────────────────────────────────────────────────────

    def _save_to_collection(self):
        cols = storage.get_collections()
        if not cols:
            QMessageBox.information(self, "No Collections",
                "Create a collection first using the + button in the sidebar.")
            return
        col_name, ok = QInputDialog.getItem(
            self, "Save Request", "Collection:", [c["name"] for c in cols], 0, False)
        if not ok:
            return
        col_id = next(c["id"] for c in cols if c["name"] == col_name)
        tab = self.tabs_manager.current_tab()
        if not tab:
            return
        req = tab.get_request()
        req_name, ok2 = QInputDialog.getText(
            self, "Request Name", "Name:", text=req.name or req.url or "New Request")
        if not ok2 or not req_name.strip():
            return
        req.name = req_name.strip()
        storage.save_request(col_id, req)
        self.tabs_manager.rename_current_tab(req.name)
        self.sidebar.refresh()
        self._set_status(f"✓  Saved '{req.name}'  →  '{col_name}'")

    # ── Dialogs ───────────────────────────────────────────────────────────────

    def _open_env_dialog(self):
        from app.ui.environment_dialog import EnvironmentDialog
        EnvironmentDialog(self).exec()
        self._update_env_selector()

    def _open_runner(self):
        from app.ui.collection_runner import CollectionRunnerDialog
        CollectionRunnerDialog(self).exec()

    def _open_code_gen(self):
        tab = self.tabs_manager.current_tab()
        if not tab:
            return
        from app.ui.code_gen_dialog import CodeGenDialog
        CodeGenDialog(tab.get_request(), self).exec()

    def _import(self):
        from app.ui.import_dialog import ImportDialog
        dlg = ImportDialog(self)
        dlg.exec()
        if dlg.imported_req:
            self.tabs_manager.load_request_in_new_tab(dlg.imported_req)
        self.sidebar.refresh()

    def _export(self):
        cols = storage.get_collections()
        if not cols:
            QMessageBox.information(self, "Nothing to Export", "No collections found.")
            return
        name, ok = QInputDialog.getItem(
            self, "Export Collection", "Collection:", [c["name"] for c in cols], 0, False)
        if not ok:
            return
        col_id = next(c["id"] for c in cols if c["name"] == name)
        path, _ = QFileDialog.getSaveFileName(
            self, "Save", f"{name}.json", "JSON (*.json)")
        if not path:
            return
        from app.core.exporter import export_collection
        with open(path, "w", encoding="utf-8") as f:
            f.write(export_collection(col_id))
        self._set_status(f"✓  Exported to {path}")

    def _open_search(self):
        from app.ui.search_dialog import SearchDialog
        dlg = SearchDialog(self)
        dlg.request_selected.connect(self._load_request_in_tab)
        dlg.exec()

    def _open_ai_assistant(self):
        from app.ui.ai_panel import AIAssistantDialog
        dlg = AIAssistantDialog(self)
        dlg.collection_imported.connect(self.sidebar.refresh)
        dlg.exec()
        self._update_env_selector()
        self.sidebar.refresh()

    def _toggle_ai_chat(self):
        if self.chat_panel.isVisible():
            self.chat_panel.hide()
        else:
            # Ensure the splitter gives the panel a visible width
            sizes = self._main_splitter.sizes()
            if sizes[2] < 50:
                total   = self._main_splitter.width() - 2  # 2 for handles
                chat_w  = 360
                side_w  = sizes[0]
                work_w  = max(total - side_w - chat_w, 400)
                self._main_splitter.setSizes([side_w, work_w, chat_w])
            self.chat_panel.show()
            tab = self.tabs_manager.current_tab()
            if tab:
                self.chat_panel.set_context(
                    req=tab.get_request(),
                    env_vars=self._get_active_variables()
                )

    def _ai_apply(self, atype: str, content: str):
        tab = self.tabs_manager.current_tab()
        if not tab:
            return
        rp = tab.request_panel
        if atype == "body":
            rp.apply_body(content)
        elif atype == "params":
            rp.apply_params(content)
        elif atype == "headers":
            rp.apply_headers(content)
        elif atype == "test":
            rp.apply_test_script(content)

    # ── Theme ─────────────────────────────────────────────────────────────────

    def _toggle_theme(self):
        toggle_theme(QApplication.instance())
        self.env_bar.theme_btn.setText("◑" if is_dark() else "◐")

    # ── Status Bar ────────────────────────────────────────────────────────────

    def _set_status(self, msg: str, error: bool = False):
        color = Colors.ERROR if error else Colors.TEXT_SECONDARY
        self._status_bar.setStyleSheet(
            f"QStatusBar {{ color: {color}; background: {Colors.BG_DARKEST}; "
            f"border-top: 1px solid {Colors.BORDER}; font-size: 11px; }}"
        )
        self._status_bar.showMessage(msg, 8000)


def _fmt_size(n: int) -> str:
    if n < 1024:         return f"{n} B"
    if n < 1024 * 1024:  return f"{n / 1024:.1f} KB"
    return f"{n / (1024 * 1024):.2f} MB"