DBClient/app/ui/explain_view.py

"""
EXPLAIN Plan Viewer
===================
Shows a database query's execution plan in two panes:

  Left  — Visual tree of nodes (parsed from each DB engine's EXPLAIN output).
  Right — Raw EXPLAIN results table (columns + rows from the driver).

Supports:
  • MySQL / MariaDB  : tabular EXPLAIN rows  (id, select_type, table, type, …)
  • PostgreSQL       : EXPLAIN ANALYZE text  (indented plan lines)
  • SQLite           : EXPLAIN QUERY PLAN    (id, parent, notused, detail)
  • MSSQL            : raw rows (minimal support)

Open from EditorTab via the "🔎 Explain" toolbar button.
The ExplainPanel is embeddable as a workspace tab.
"""
from __future__ import annotations

from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QSplitter,
    QTreeWidget, QTreeWidgetItem, QTableWidget, QTableWidgetItem,
    QLabel, QPushButton, QHeaderView, QPlainTextEdit, QStackedWidget,
    QAbstractItemView,
)
from PyQt6.QtCore import Qt, QThread, pyqtSignal
from PyQt6.QtGui import QColor, QFont, QBrush

from app.ui.syntax_highlighter import SQLHighlighter


# ── Cost-metric colour helpers ────────────────────────────────────────────────

def _cost_color(val_str: str) -> QColor:
    """Return a colour (green→yellow→red) based on a cost value string."""
    try:
        v = float(str(val_str).replace(",", ""))
    except (ValueError, TypeError):
        return QColor("#c6d0f5")   # neutral (Catppuccin text)
    if v <= 10:
        return QColor("#a6d189")   # green  — cheap
    if v <= 1_000:
        return QColor("#e5c890")   # yellow — moderate
    return QColor("#e78284")       # red    — expensive


# ── Background worker ─────────────────────────────────────────────────────────

class _ExplainWorker(QThread):
    finished = pyqtSignal(list, list)   # (cols, rows)
    error    = pyqtSignal(str)

    def __init__(self, driver, sql: str, parent=None):
        super().__init__(parent)
        self._driver = driver
        self._sql    = sql

    def run(self):
        try:
            cols, rows = self._driver.explain_query(self._sql)
            self.finished.emit(list(cols), [tuple(r) for r in rows])
        except Exception as e:
            self.error.emit(str(e))


# ── Tree builders per DB engine ───────────────────────────────────────────────

def _build_mysql_tree(tree: QTreeWidget, cols: list, rows: list):
    """
    MySQL EXPLAIN columns:
      id | select_type | table | partitions | type | possible_keys |
      key | key_len | ref | rows | filtered | Extra
    """
    tree.setColumnCount(5)
    tree.setHeaderLabels(["Table / Step", "Type", "Key", "Rows est.", "Extra"])
    tree.header().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents)
    tree.header().setSectionResizeMode(4, QHeaderView.ResizeMode.Stretch)

    col = {name.lower(): i for i, name in enumerate(cols)}
    prev_items: dict[int, QTreeWidgetItem] = {}

    for row in rows:
        def g(name: str) -> str:
            i = col.get(name)
            return "" if i is None else (str(row[i]) if row[i] is not None else "NULL")

        step_id    = g("id") or "?"
        sel_type   = g("select_type")
        table      = g("table")
        join_type  = g("type")
        key        = g("key")
        est_rows   = g("rows")
        extra      = g("extra")
        filtered   = g("filtered")

        label = f"[{step_id}] {table}" if table else f"[{step_id}] {sel_type}"
        item  = QTreeWidgetItem([label, join_type, key, est_rows, extra])

        # Colour the join-type column
        bad_types = {"all", "index", "range"}
        good_types = {"const", "eq_ref", "ref", "system"}
        jt_lower = join_type.lower()
        if jt_lower in bad_types:
            item.setForeground(1, QBrush(QColor("#e78284")))
        elif jt_lower in good_types:
            item.setForeground(1, QBrush(QColor("#a6d189")))

        # Colour rows estimate
        item.setForeground(3, QBrush(_cost_color(est_rows)))

        # Nest: use select_id as parent key (simplified: flat for now)
        try:
            sid = int(step_id)
        except ValueError:
            sid = 0
        parent_item = prev_items.get(sid - 1)
        if parent_item:
            parent_item.addChild(item)
        else:
            tree.addTopLevelItem(item)
        prev_items[sid] = item

    tree.expandAll()


def _build_postgres_tree(tree: QTreeWidget, cols: list, rows: list):
    """
    PostgreSQL EXPLAIN ANALYZE returns text lines in a single 'Plan' column.
    Parse indentation to build the tree.
    """
    tree.setColumnCount(2)
    tree.setHeaderLabels(["Plan Node", "Cost / Hint"])
    tree.header().setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch)
    tree.header().setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents)

    # rows is list of 1-tuples of text lines
    lines = [row[0] for row in rows if row]

    stack: list[tuple[int, QTreeWidgetItem]] = []   # (indent, item)

    for line in lines:
        if isinstance(line, (list, tuple)):
            line = str(line[0]) if line else ""
        text = str(line)
        stripped = text.lstrip("-> ").lstrip()
        indent   = len(text) - len(text.lstrip())

        # Split "Node  (cost=x..y rows=z ...)"
        hint = ""
        if "(cost=" in stripped or "(actual" in stripped:
            split_at = stripped.find("(")
            hint     = stripped[split_at:]
            stripped = stripped[:split_at].strip()

        item = QTreeWidgetItem([stripped.strip("->").strip(), hint])

        # Colour based on cost if available
        if "cost=" in hint:
            try:
                cost_str = hint.split("cost=")[1].split("..")[1].split()[0].rstrip(")")
                item.setForeground(1, QBrush(_cost_color(cost_str)))
            except (IndexError, ValueError):
                pass

        # Find parent by indent level
        while stack and stack[-1][0] >= indent:
            stack.pop()

        if stack:
            stack[-1][1].addChild(item)
        else:
            tree.addTopLevelItem(item)

        stack.append((indent, item))

    tree.expandAll()


def _build_sqlite_tree(tree: QTreeWidget, cols: list, rows: list):
    """
    SQLite EXPLAIN QUERY PLAN columns: id, parent, notused, detail
    (older SQLite: selectid, order, from, detail)
    """
    tree.setColumnCount(2)
    tree.setHeaderLabels(["Step", "Detail"])
    tree.header().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)

    col = {name.lower(): i for i, name in enumerate(cols)}

    # Try modern format first (id, parent, notused, detail)
    id_col     = col.get("id",     col.get("selectid", 0))
    parent_col = col.get("parent", col.get("order",    1))
    detail_col = col.get("detail", col.get("from",     3))

    items: dict = {}    # id → QTreeWidgetItem

    for row in rows:
        rid    = row[id_col]    if id_col < len(row) else 0
        parent = row[parent_col] if parent_col < len(row) else 0
        detail = row[detail_col] if detail_col < len(row) else str(row)

        item = QTreeWidgetItem([str(rid), str(detail)])
        # Colour SCAN (bad) vs SEARCH/INDEX (good)
        detail_str = str(detail).upper()
        if "SCAN" in detail_str and "INDEX" not in detail_str:
            item.setForeground(1, QBrush(QColor("#e78284")))
        elif "INDEX" in detail_str or "SEARCH" in detail_str:
            item.setForeground(1, QBrush(QColor("#a6d189")))

        items[rid] = item
        parent_item = items.get(parent)
        if parent_item and parent != rid:
            parent_item.addChild(item)
        else:
            tree.addTopLevelItem(item)

    tree.expandAll()


def _build_generic_tree(tree: QTreeWidget, cols: list, rows: list):
    """Fallback: show rows flat."""
    tree.setColumnCount(len(cols))
    tree.setHeaderLabels(cols)
    for row in rows:
        item = QTreeWidgetItem([str(v) if v is not None else "NULL" for v in row])
        tree.addTopLevelItem(item)
    tree.expandAll()


# ── Raw table helper ──────────────────────────────────────────────────────────

def _populate_raw_table(table: QTableWidget, cols: list, rows: list):
    table.setColumnCount(len(cols))
    table.setHorizontalHeaderLabels(cols)
    table.setRowCount(len(rows))
    for r, row in enumerate(rows):
        for c, val in enumerate(row):
            item = QTableWidgetItem(
                "NULL" if val is None else str(val)
            )
            item.setTextAlignment(Qt.AlignmentFlag.AlignCenter)
            table.setItem(r, c, item)
    table.resizeColumnsToContents()
    table.horizontalHeader().setStretchLastSection(True)


# ── Main panel ────────────────────────────────────────────────────────────────

class ExplainPanel(QWidget):
    """
    Embeddable widget that runs EXPLAIN on a SQL query and displays
    a visual tree (left) + raw tabular output (right).
    """

    status_message = pyqtSignal(str)

    def __init__(self, driver, sql: str, parent=None):
        super().__init__(parent)
        self._driver = driver
        self._sql    = sql
        self._worker: _ExplainWorker | None = None
        self._build_ui()
        self._run_explain()

    # ── UI ────────────────────────────────────────────────────────────────────

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

        # ── Header bar ────────────────────────────────────────────────────────
        hdr = QHBoxLayout()
        hdr.setContentsMargins(8, 6, 8, 4)
        hdr.setSpacing(8)

        title = QLabel("🔎  Query Execution Plan")
        f = title.font()
        f.setBold(True)
        f.setPointSize(11)
        title.setFont(f)
        title.setObjectName("structureTitle")

        self._refresh_btn = QPushButton("🔄 Re-run")
        self._refresh_btn.setFixedWidth(80)
        self._refresh_btn.clicked.connect(self._run_explain)

        self._status_lbl = QLabel("Running…")
        self._status_lbl.setObjectName("rowCountLbl")

        hdr.addWidget(title)
        hdr.addStretch()
        hdr.addWidget(self._status_lbl)
        hdr.addWidget(self._refresh_btn)
        root.addLayout(hdr)

        # ── SQL preview (collapsed) ───────────────────────────────────────────
        sql_lbl = QLabel("SQL:")
        sql_lbl.setContentsMargins(8, 0, 8, 0)
        self._sql_preview = QPlainTextEdit(self._sql)
        self._sql_preview.setReadOnly(True)
        self._sql_preview.setMaximumHeight(52)
        mono = QFont("Consolas", 10)
        self._sql_preview.setFont(mono)
        SQLHighlighter(self._sql_preview.document())
        root.addWidget(sql_lbl)
        root.addWidget(self._sql_preview)

        # ── Main splitter: tree / raw ─────────────────────────────────────────
        splitter = QSplitter(Qt.Orientation.Horizontal)

        # Left — visual plan tree
        left = QWidget()
        ll   = QVBoxLayout(left)
        ll.setContentsMargins(0, 0, 0, 0)
        ll.setSpacing(0)

        tree_hdr = QLabel("  Visual Plan")
        tree_hdr.setObjectName("sidebarTitle")
        tree_hdr.setContentsMargins(8, 4, 0, 4)
        f2 = tree_hdr.font()
        f2.setBold(True)
        tree_hdr.setFont(f2)
        ll.addWidget(tree_hdr)

        self._tree = QTreeWidget()
        self._tree.setAlternatingRowColors(True)
        self._tree.setAnimated(True)
        self._tree.setSelectionBehavior(
            QAbstractItemView.SelectionBehavior.SelectRows)
        ll.addWidget(self._tree)

        # Legend row
        legend_row = QHBoxLayout()
        legend_row.setContentsMargins(8, 2, 0, 4)
        for colour, label in [
            ("#a6d189", "Efficient"), ("#e5c890", "Moderate"),
            ("#e78284", "Expensive / SCAN"),
        ]:
            dot = QLabel("●")
            dot.setStyleSheet(f"color: {colour};")
            legend_row.addWidget(dot)
            legend_row.addWidget(QLabel(label))
            legend_row.addSpacing(12)
        legend_row.addStretch()
        ll.addLayout(legend_row)

        splitter.addWidget(left)

        # Right — raw results table
        right = QWidget()
        rl    = QVBoxLayout(right)
        rl.setContentsMargins(0, 0, 0, 0)
        rl.setSpacing(0)

        raw_hdr = QLabel("  Raw EXPLAIN Output")
        raw_hdr.setObjectName("sidebarTitle")
        raw_hdr.setContentsMargins(8, 4, 0, 4)
        f3 = raw_hdr.font()
        f3.setBold(True)
        raw_hdr.setFont(f3)
        rl.addWidget(raw_hdr)

        self._raw_table = QTableWidget()
        self._raw_table.setAlternatingRowColors(True)
        self._raw_table.setEditTriggers(QTableWidget.EditTrigger.NoEditTriggers)
        self._raw_table.setSelectionBehavior(
            QAbstractItemView.SelectionBehavior.SelectRows)
        self._raw_table.verticalHeader().setDefaultSectionSize(24)
        rl.addWidget(self._raw_table)

        splitter.addWidget(right)
        splitter.setSizes([450, 450])

        root.addWidget(splitter, 1)

    # ── Explain execution ─────────────────────────────────────────────────────

    def _run_explain(self):
        self._status_lbl.setText("Running EXPLAIN…")
        self._tree.clear()
        self._raw_table.setRowCount(0)
        self._refresh_btn.setEnabled(False)

        self._worker = _ExplainWorker(self._driver, self._sql, parent=self)
        self._worker.finished.connect(self._on_result)
        self._worker.error.connect(self._on_error)
        self._worker.start()

    def _on_result(self, cols: list, rows: list):
        self._refresh_btn.setEnabled(True)

        # ── Populate raw table ────────────────────────────────────────────────
        _populate_raw_table(self._raw_table, cols, rows)

        # ── Build visual tree based on DB type ────────────────────────────────
        db_type = getattr(self._driver, "db_type", "").lower()

        if db_type in ("mysql", "mariadb"):
            _build_mysql_tree(self._tree, cols, rows)
        elif db_type == "postgresql":
            _build_postgres_tree(self._tree, cols, rows)
        elif db_type == "sqlite":
            _build_sqlite_tree(self._tree, cols, rows)
        else:
            _build_generic_tree(self._tree, cols, rows)

        n = len(rows)
        self._status_lbl.setText(f"{n} plan node{'s' if n != 1 else ''}")
        self.status_message.emit(f"EXPLAIN complete — {n} nodes")

    def _on_error(self, msg: str):
        self._refresh_btn.setEnabled(True)
        self._status_lbl.setText(f"Error: {msg[:80]}")
        self._tree.clear()
        err_item = QTreeWidgetItem([f"⚠  {msg}"])
        err_item.setForeground(0, QBrush(QColor("#e78284")))
        self._tree.setColumnCount(1)
        self._tree.setHeaderLabels(["Error"])
        self._tree.addTopLevelItem(err_item)
        self.status_message.emit(f"EXPLAIN error: {msg}")