DBClient/app/ui/schema_browser.py

"""
Schema browser — left sidebar tree.

Node states:
  • saved_connection  — profile saved to disk but not yet connected (grey, dashed)
  • connection        — actively connected (coloured, expandable tree)

Tree structure when connected:
  Connection → Databases → (Tables, Views, Functions, Procedures, Triggers)
                            → Columns / Indexes / Foreign Keys
"""
from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QTreeWidget, QTreeWidgetItem,
    QLineEdit, QMenu, QMessageBox, QInputDialog,
)
from PyQt6.QtCore import Qt, pyqtSignal
from PyQt6.QtGui import QColor, QBrush

from app.utils.icons import get_icon
from app.utils.worker import SchemaWorker

# ── Node type codes stored in UserRole ────────────────────────────────────────
NT = {
    # saved but not connected
    "saved_connection": -1,
    # connected
    "connection": 0, "databases_folder": 1, "database": 2,
    "tables_folder": 3, "views_folder": 4, "functions_folder": 5,
    "procedures_folder": 6, "triggers_folder": 7,
    "table": 8, "view": 9, "function": 10, "procedure": 11, "trigger": 12,
    "columns_folder": 13, "indexes_folder": 14, "fks_folder": 15,
    "column": 16, "index": 17, "fk": 18,
}


class SchemaBrowser(QWidget):
    # ── Signals ───────────────────────────────────────────────────────────────
    open_table_viewer    = pyqtSignal(object, str, str)  # driver, db, table
    open_table_structure = pyqtSignal(object, str, str)
    open_sql_editor      = pyqtSignal(object, str)       # driver, db
    run_query_requested  = pyqtSignal(str)               # SQL snippet

    # Emitted so the main window can act on saved profile management
    connect_requested    = pyqtSignal(str)   # profile_id — user wants to connect
    edit_requested       = pyqtSignal(str)   # profile_id — user wants to edit
    delete_requested     = pyqtSignal(str)   # profile_id — user wants to delete

    def __init__(self, parent=None):
        super().__init__(parent)
        self._drivers:  dict = {}   # profile_id → driver (only connected)
        self._profiles: dict = {}   # profile_id → ConnectionProfile (all)
        self._items:    dict = {}   # profile_id → top-level QTreeWidgetItem
        self._workers:  list = []
        self._build_ui()

    # ── UI construction ───────────────────────────────────────────────────────

    def _build_ui(self):
        layout = QVBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)

        search_row = QHBoxLayout()
        search_row.setContentsMargins(6, 6, 6, 4)
        self._search = QLineEdit()
        self._search.setPlaceholderText("🔍  Filter objects…")
        self._search.textChanged.connect(self._filter)
        search_row.addWidget(self._search)
        layout.addLayout(search_row)

        self._tree = QTreeWidget()
        self._tree.setHeaderHidden(True)
        self._tree.setAnimated(True)
        self._tree.setIndentation(16)
        self._tree.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self._tree.customContextMenuRequested.connect(self._context_menu)
        self._tree.itemDoubleClicked.connect(self._on_double_click)
        self._tree.itemExpanded.connect(self._on_expanded)
        layout.addWidget(self._tree)

    # ── Public API ────────────────────────────────────────────────────────────

    def add_saved_profile(self, profile) -> None:
        """Add a profile that is saved but not yet connected (greyed out)."""
        pid = profile.id
        self._profiles[pid] = profile

        # If it already exists as a tree item (e.g., being refreshed), remove it first
        self._remove_item(pid)

        item = QTreeWidgetItem()
        item.setText(0, f"  {profile.name}  ·  {profile.db_type_display}")
        item.setData(0, Qt.ItemDataRole.UserRole,
                     (NT["saved_connection"], pid, "", ""))

        # Dim colour — show the profile colour but muted
        color = QColor(profile.color)
        color.setAlpha(140)
        item.setForeground(0, QBrush(QColor("#737994")))   # greyed out
        item.setIcon(0, get_icon(profile.db_type, 16))
        item.setToolTip(0, self._connection_tooltip(profile))

        f = item.font(0)
        f.setItalic(True)
        item.setFont(0, f)

        self._tree.addTopLevelItem(item)
        self._items[pid] = item

    def add_connection(self, profile, driver) -> None:
        """
        Upgrade a saved (disconnected) profile to a live connected node,
        or add it fresh if it was never in the tree before.
        """
        pid = profile.id
        self._drivers[pid]  = driver
        self._profiles[pid] = profile

        # Reuse existing item if already in tree (upgrade from saved_connection)
        existing = self._items.get(pid)
        if existing:
            self._tree.takeTopLevelItem(
                self._tree.indexOfTopLevelItem(existing))

        item = QTreeWidgetItem()
        item.setText(0, f"  {profile.name}")
        item.setData(0, Qt.ItemDataRole.UserRole, (NT["connection"], pid, "", ""))
        item.setForeground(0, QBrush(QColor(profile.color)))

        f = item.font(0)
        f.setBold(True)
        f.setItalic(False)
        item.setFont(0, f)
        item.setIcon(0, get_icon(profile.db_type, 16))
        item.setToolTip(0, self._connection_tooltip(profile))
        item.addChild(QTreeWidgetItem(["Loading…"]))   # triggers expand arrow

        self._tree.addTopLevelItem(item)
        self._items[pid] = item
        item.setExpanded(True)

    def remove_connection(self, profile_id: str, keep_saved: bool = True) -> None:
        """
        Disconnect: revert the node to 'saved' state (keep_saved=True)
        or remove it entirely (keep_saved=False).
        """
        driver = self._drivers.pop(profile_id, None)
        if driver:
            try:
                driver.disconnect()
            except Exception:
                pass

        if keep_saved and profile_id in self._profiles:
            profile = self._profiles[profile_id]
            self._remove_item(profile_id)
            self.add_saved_profile(profile)
        else:
            self._profiles.pop(profile_id, None)
            self._remove_item(profile_id)

    def update_saved_profile(self, profile) -> None:
        """Called after an edit — refresh the sidebar label."""
        pid = profile.id
        self._profiles[pid] = profile
        if pid in self._drivers:
            # Currently connected — just update the label
            item = self._items.get(pid)
            if item:
                item.setText(0, f"  {profile.name}")
                item.setForeground(0, QBrush(QColor(profile.color)))
        else:
            # Saved only — rebuild the item
            self._remove_item(pid)
            self.add_saved_profile(profile)

    def is_connected(self, profile_id: str) -> bool:
        return profile_id in self._drivers

    # ── Helpers ───────────────────────────────────────────────────────────────

    def _remove_item(self, profile_id: str) -> None:
        item = self._items.pop(profile_id, None)
        if item:
            idx = self._tree.indexOfTopLevelItem(item)
            if idx >= 0:
                self._tree.takeTopLevelItem(idx)

    @staticmethod
    def _connection_tooltip(profile) -> str:
        if profile.db_type == "sqlite":
            return profile.database
        return (f"{profile.db_type_display}  ·  "
                f"{profile.username}@{profile.host}:{profile.port}"
                + (f"/{profile.database}" if profile.database else ""))

    # ── Tree population (lazy-load on expand) ─────────────────────────────────

    def _on_expanded(self, item: QTreeWidgetItem):
        data = item.data(0, Qt.ItemDataRole.UserRole)
        if not data:
            return
        node_type, pid, db, name = data

        if node_type == NT["connection"]:
            self._load_databases(item, pid)
        elif node_type == NT["database"]:
            self._build_db_folders(item, pid, db)
        elif node_type == NT["tables_folder"]:
            self._load_tables(item, pid, db)
        elif node_type == NT["views_folder"]:
            self._load_views(item, pid, db)
        elif node_type == NT["functions_folder"]:
            self._load_generic(item, pid, db, "function",
                               lambda d: self._drivers[pid].get_functions(d))
        elif node_type == NT["procedures_folder"]:
            self._load_generic(item, pid, db, "procedure",
                               lambda d: self._drivers[pid].get_stored_procedures(d))
        elif node_type == NT["triggers_folder"]:
            self._load_generic(item, pid, db, "trigger",
                               lambda d: self._drivers[pid].get_triggers(d))
        elif node_type == NT["table"]:
            self._build_table_folders(item, pid, db, name)
        elif node_type == NT["columns_folder"]:
            self._load_columns(item, pid, db, name)
        elif node_type == NT["indexes_folder"]:
            self._load_indexes(item, pid, db, name)
        elif node_type == NT["fks_folder"]:
            self._load_fks(item, pid, db, name)

    def _placeholder(self, item: QTreeWidgetItem) -> bool:
        if item.childCount() == 1 and item.child(0).text(0) in ("Loading…", ""):
            item.takeChild(0)
            return True
        return False

    # databases
    def _load_databases(self, item, pid):
        if not self._placeholder(item):
            return
        w = SchemaWorker(self._drivers[pid].get_databases, parent=self)
        w.result.connect(lambda dbs: self._populate_databases(item, pid, dbs))
        w.error.connect(lambda e: self._show_error(item, e))
        self._workers.append(w)
        w.start()

    def _populate_databases(self, item, pid, dbs):
        for db in dbs:
            child = QTreeWidgetItem([f" 🗄  {db}"])
            child.setData(0, Qt.ItemDataRole.UserRole, (NT["database"], pid, db, ""))
            child.addChild(QTreeWidgetItem(["Loading…"]))
            item.addChild(child)

    # database folders
    def _build_db_folders(self, item, pid, db):
        if not self._placeholder(item):
            return
        for label, nt in [
            ("📂 Tables",     NT["tables_folder"]),
            ("📂 Views",      NT["views_folder"]),
            ("📂 Functions",  NT["functions_folder"]),
            ("📂 Procedures", NT["procedures_folder"]),
            ("📂 Triggers",   NT["triggers_folder"]),
        ]:
            child = QTreeWidgetItem([label])
            child.setData(0, Qt.ItemDataRole.UserRole, (nt, pid, db, ""))
            child.addChild(QTreeWidgetItem(["Loading…"]))
            item.addChild(child)

    # tables
    def _load_tables(self, item, pid, db):
        if not self._placeholder(item):
            return
        w = SchemaWorker(self._drivers[pid].get_tables, db, parent=self)
        w.result.connect(lambda t: self._populate_tables(item, pid, db, t))
        w.error.connect(lambda e: self._show_error(item, e))
        self._workers.append(w)
        w.start()

    def _populate_tables(self, item, pid, db, tables):
        for t in tables:
            label = f" 📋 {t.name}"
            if t.row_count:
                label += f"  ({t.row_count:,})"
            child = QTreeWidgetItem([label])
            child.setData(0, Qt.ItemDataRole.UserRole, (NT["table"], pid, db, t.name))
            child.setToolTip(0, t.comment or t.name)
            child.addChild(QTreeWidgetItem(["Loading…"]))
            item.addChild(child)

    # views
    def _load_views(self, item, pid, db):
        if not self._placeholder(item):
            return
        w = SchemaWorker(self._drivers[pid].get_views, db, parent=self)
        w.result.connect(
            lambda v: self._populate_generic(item, pid, db, v, "view", "👁"))
        w.error.connect(lambda e: self._show_error(item, e))
        self._workers.append(w)
        w.start()

    # generic (functions / procedures / triggers)
    def _load_generic(self, item, pid, db, kind, fn):
        if not self._placeholder(item):
            return
        icons = {"function": "⚡", "procedure": "📦", "trigger": "⚙"}
        icon = icons.get(kind, "📄")
        w = SchemaWorker(fn, db, parent=self)
        w.result.connect(
            lambda items: self._populate_generic(item, pid, db, items, kind, icon))
        w.error.connect(lambda e: self._show_error(item, e))
        self._workers.append(w)
        w.start()

    def _populate_generic(self, item, pid, db, names, kind, icon):
        nt = NT.get(kind, NT["view"])
        for name in names:
            child = QTreeWidgetItem([f" {icon} {name}"])
            child.setData(0, Qt.ItemDataRole.UserRole, (nt, pid, db, name))
            item.addChild(child)

    # table sub-folders
    def _build_table_folders(self, item, pid, db, table):
        if not self._placeholder(item):
            return
        for label, nt in [
            ("📊 Columns",      NT["columns_folder"]),
            ("🔍 Indexes",      NT["indexes_folder"]),
            ("🔗 Foreign Keys", NT["fks_folder"]),
        ]:
            child = QTreeWidgetItem([label])
            child.setData(0, Qt.ItemDataRole.UserRole, (nt, pid, db, table))
            child.addChild(QTreeWidgetItem(["Loading…"]))
            item.addChild(child)

    def _load_columns(self, item, pid, db, table):
        if not self._placeholder(item):
            return
        w = SchemaWorker(self._drivers[pid].get_columns, db, table, parent=self)
        w.result.connect(lambda c: self._populate_columns(item, c))
        w.error.connect(lambda e: self._show_error(item, e))
        self._workers.append(w)
        w.start()

    def _populate_columns(self, item, cols):
        for c in cols:
            icon = "🔑" if c.is_primary_key else ("🔗" if c.is_foreign_key else "·")
            null = "NULL" if c.nullable else "NOT NULL"
            label = f"  {icon} {c.name}   {c.data_type}  {null}"
            child = QTreeWidgetItem([label])
            child.setData(0, Qt.ItemDataRole.UserRole, (NT["column"], "", "", c.name))
            if c.is_primary_key:
                child.setForeground(0, QBrush(QColor("#e5c890")))
            elif c.is_foreign_key:
                child.setForeground(0, QBrush(QColor("#99d1db")))
            item.addChild(child)

    def _load_indexes(self, item, pid, db, table):
        if not self._placeholder(item):
            return
        w = SchemaWorker(self._drivers[pid].get_indexes, db, table, parent=self)
        w.result.connect(lambda i: self._populate_indexes(item, i))
        w.error.connect(lambda e: self._show_error(item, e))
        self._workers.append(w)
        w.start()

    def _populate_indexes(self, item, indexes):
        for idx in indexes:
            u = " [UNIQUE]" if idx.is_unique else ""
            label = f" 🔍 {idx.name}{u}  ({', '.join(idx.columns)})"
            item.addChild(QTreeWidgetItem([label]))

    def _load_fks(self, item, pid, db, table):
        if not self._placeholder(item):
            return
        w = SchemaWorker(self._drivers[pid].get_foreign_keys, db, table, parent=self)
        w.result.connect(lambda f: self._populate_fks(item, f))
        w.error.connect(lambda e: self._show_error(item, e))
        self._workers.append(w)
        w.start()

    def _populate_fks(self, item, fks):
        for fk in fks:
            label = f" 🔗 {fk.column} → {fk.ref_table}.{fk.ref_column}"
            item.addChild(QTreeWidgetItem([label]))

    def _show_error(self, item, msg):
        err = QTreeWidgetItem([f"❌ {msg}"])
        err.setForeground(0, QBrush(QColor("#e78284")))
        item.addChild(err)

    # ── Context menu ──────────────────────────────────────────────────────────

    def _context_menu(self, pos):
        item = self._tree.itemAt(pos)
        if not item:
            return
        data = item.data(0, Qt.ItemDataRole.UserRole)
        if not data:
            return
        node_type, pid, db, name = data

        menu = QMenu(self)

        # ── Saved (disconnected) connection ───────────────────────────────────
        if node_type == NT["saved_connection"]:
            menu.addAction("🔌 Connect", lambda: self.connect_requested.emit(pid))
            menu.addSeparator()
            menu.addAction("✏️ Edit Connection…", lambda: self.edit_requested.emit(pid))
            menu.addAction("🗑️ Delete Connection", lambda: self.delete_requested.emit(pid))

        # ── Active connection ─────────────────────────────────────────────────
        elif node_type == NT["connection"]:
            menu.addAction("🔄 Refresh",            lambda: self._refresh_node(item))
            menu.addAction("🗄️ New Database…",      lambda: self._create_db(pid))
            menu.addSeparator()
            menu.addAction("✏️ Edit Connection…",   lambda: self.edit_requested.emit(pid))
            menu.addAction("🗑️ Delete Connection…", lambda: self.delete_requested.emit(pid))
            menu.addSeparator()
            menu.addAction("🔌 Disconnect",         lambda: self._disconnect(pid))

        # ── Database ──────────────────────────────────────────────────────────
        elif node_type == NT["database"]:
            menu.addAction("✏️ New SQL Tab",
                           lambda: self.open_sql_editor.emit(self._drivers[pid], db))
            menu.addAction("🔄 Refresh", lambda: self._refresh_node(item))

        # ── Table ──────────────────────────────────────────────────────────────
        elif node_type == NT["table"]:
            d = self._drivers[pid]
            menu.addAction("📋 Open Data",
                           lambda: self.open_table_viewer.emit(d, db, name))
            menu.addAction("🏗️ Table Structure",
                           lambda: self.open_table_structure.emit(d, db, name))
            menu.addSeparator()
            menu.addAction("📑 Copy SELECT *",
                           lambda: self.run_query_requested.emit(
                               f"SELECT * FROM {name} LIMIT 100;"))
            menu.addAction("🔄 Refresh", lambda: self._refresh_node(item))
            menu.addSeparator()
            menu.addAction("🗑️ Drop Table…", lambda: self._drop_table(pid, db, name))

        elif node_type in (NT["tables_folder"], NT["views_folder"]):
            menu.addAction("🔄 Refresh", lambda: self._refresh_node(item))

        if menu.actions():
            menu.exec(self._tree.viewport().mapToGlobal(pos))

    # ── Context menu actions ──────────────────────────────────────────────────

    def _disconnect(self, pid: str):
        self.remove_connection(pid, keep_saved=True)

    def _refresh_node(self, item: QTreeWidgetItem):
        item.takeChildren()
        item.addChild(QTreeWidgetItem(["Loading…"]))
        item.setExpanded(False)
        item.setExpanded(True)

    def _create_db(self, pid: str):
        name, ok = QInputDialog.getText(self, "New Database", "Database name:")
        if ok and name.strip():
            try:
                self._drivers[pid].create_database(name.strip())
                QMessageBox.information(self, "Success",
                                        f"Database '{name.strip()}' created.")
            except Exception as e:
                QMessageBox.critical(self, "Error", str(e))

    def _drop_table(self, pid: str, _db: str, table: str):
        btn = QMessageBox.warning(
            self, "Drop Table",
            f"Are you sure you want to DROP TABLE '{table}'?\nThis cannot be undone.",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No,
        )
        if btn == QMessageBox.StandardButton.Yes:
            try:
                self._drivers[pid].execute_query(f"DROP TABLE `{table}`")
                QMessageBox.information(self, "Dropped", f"Table '{table}' dropped.")
            except Exception as e:
                QMessageBox.critical(self, "Error", str(e))

    # ── Filter ────────────────────────────────────────────────────────────────

    def _filter(self, text: str):
        text = text.lower()
        for i in range(self._tree.topLevelItemCount()):
            self._filter_item(self._tree.topLevelItem(i), text)

    def _filter_item(self, item: QTreeWidgetItem, text: str) -> bool:
        matches = text in item.text(0).lower()
        child_match = any(
            self._filter_item(item.child(i), text)
            for i in range(item.childCount())
        )
        visible = matches or child_match
        item.setHidden(not visible)
        return visible

    # ── Double-click ──────────────────────────────────────────────────────────

    def _on_double_click(self, item: QTreeWidgetItem, _col: int):
        data = item.data(0, Qt.ItemDataRole.UserRole)
        if not data:
            return
        node_type, pid, db, name = data

        if node_type == NT["saved_connection"]:
            # Double-click on disconnected profile → try to connect
            self.connect_requested.emit(pid)

        elif node_type == NT["table"]:
            driver = self._drivers.get(pid)
            if driver:
                self.open_table_viewer.emit(driver, db, name)

        elif node_type == NT["database"]:
            driver = self._drivers.get(pid)
            if driver:
                self.open_sql_editor.emit(driver, db)