DBClient/app/ui/sql_completer.py

"""
Schema-aware SQL auto-completer for CodeEditor.

Provides keyword completions from the syntax highlighter lists, plus
table/view/column names loaded asynchronously from the active driver.
Trigger: type 2+ characters, or press Ctrl+Space to force-show.
"""
import re
from PyQt6.QtCore import Qt, QObject
from PyQt6.QtGui import QFont, QFontMetrics, QTextCursor
from PyQt6.QtWidgets import QApplication, QListWidget, QListWidgetItem

from app.ui.syntax_highlighter import _KEYWORDS, _TYPES, _FUNCTIONS
from app.utils.worker import SchemaWorker

# Combined, deduplicated keyword list used as the static completion pool.
SQL_KEYWORDS: list[str] = list(dict.fromkeys(_KEYWORDS + _TYPES + _FUNCTIONS))

# Matches either "table.col_prefix" or bare "word" at end of a line fragment.
_PREFIX_RE = re.compile(r'([\w]+)\.([\w]*)$|([\w]+)$', re.IGNORECASE)


def _cursor_in_string(left: str) -> bool:
    """Return True if an odd number of unescaped quotes precede the cursor."""
    count, i = 0, 0
    while i < len(left):
        if left[i] == '\\':
            i += 2
            continue
        if left[i] in ("'", '"'):
            count += 1
        i += 1
    return count % 2 == 1


# ── Floating popup ─────────────────────────────────────────────────────────────

class _CompletionPopup(QListWidget):
    """Frameless floating list that shows completion candidates."""

    def __init__(self, editor):
        super().__init__(None)
        self._editor = editor

        self.setWindowFlags(
            Qt.WindowType.ToolTip | Qt.WindowType.FramelessWindowHint
        )
        self.setFocusProxy(editor)

        font = QFont("Consolas", 11)
        self.setFont(font)
        self._row_h = QFontMetrics(font).height() + 6

        self.setStyleSheet("""
            QListWidget {
                background: #1e1e2e;
                color: #cdd6f4;
                border: 1px solid #45475a;
                border-radius: 4px;
                padding: 2px;
                outline: none;
            }
            QListWidget::item { padding: 2px 8px; }
            QListWidget::item:selected {
                background: #313244;
                color: #cdd6f4;
            }
            QListWidget::item:hover { background: #2a2a3c; }
        """)
        self.setVerticalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAsNeeded)
        self.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)

        self.itemClicked.connect(
            lambda item: self._editor._completer._accept_completion(item.text())
        )

    def show_at(self, items: list[str], cursor_rect) -> None:
        self.clear()
        for word in items:
            self.addItem(QListWidgetItem(word))

        visible = min(len(items), 10)
        fm = QFontMetrics(self.font())
        w = max(180, max(fm.horizontalAdvance(t) for t in items) + 28)
        h = visible * self._row_h + 8

        gp = self._editor.viewport().mapToGlobal(cursor_rect.bottomLeft())
        gx, gy = gp.x(), gp.y() + 2

        screen = QApplication.primaryScreen().availableGeometry()
        if gy + h > screen.bottom():
            top = self._editor.viewport().mapToGlobal(cursor_rect.topLeft())
            gy = top.y() - h - 2
        if gx + w > screen.right():
            gx = screen.right() - w

        self.setGeometry(gx, gy, w, h)
        self.setCurrentRow(0)
        self.show()

    def move_selection(self, delta: int) -> None:
        n = self.count()
        if n:
            self.setCurrentRow((self.currentRow() + delta) % n)

    def current_text(self) -> str | None:
        item = self.currentItem()
        return item.text() if item else None

    def hide_popup(self) -> None:
        self.hide()
        self.clear()


# ── Main completer ─────────────────────────────────────────────────────────────

class SqlCompleter(QObject):
    """
    Attaches to a CodeEditor and provides schema-aware completions.

    Usage (from EditorTab):
        self._completer = SqlCompleter(self._editor, parent=self)
        self._editor.set_completer(self._completer)
        self._completer.set_context(driver, database)
    """

    def __init__(self, editor, parent=None):
        super().__init__(parent)
        self._editor   = editor
        self._driver   = None
        self._database = ""
        self._load_gen = 0
        self._workers: list[SchemaWorker] = []
        self._pending_cols: set[str] = set()

        # Schema cache
        self._tables: list[str] = []
        self._views:  list[str] = []
        self._cols:   dict[str, list[str]] = {}

        self._popup = _CompletionPopup(editor)

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

    def set_context(self, driver, database: str) -> None:
        """Update the active driver/database and refresh schema cache."""
        self._driver   = driver
        self._database = database
        self._reload_schema()

    def invalidate(self) -> None:
        """Discard cached schema and reload — call after DDL statements."""
        self._tables.clear()
        self._views.clear()
        self._cols.clear()
        self._pending_cols.clear()
        if self._driver:
            self._reload_schema()

    def trigger_completion(self, force: bool = False) -> None:
        """Show/update the popup. force=True ignores the 2-char minimum."""
        prefix, qualifier = self._extract_prefix()

        if not force and len(prefix) < 2 and qualifier is None:
            self._popup.hide_popup()
            return

        candidates = self._build_candidates(prefix, qualifier)
        if not candidates:
            self._popup.hide_popup()
            return

        self._popup.show_at(candidates, self._editor.cursorRect())

    # ── Schema loading ────────────────────────────────────────────────────────

    def _reload_schema(self) -> None:
        self._load_gen += 1
        self._tables.clear()
        self._views.clear()
        self._cols.clear()
        self._pending_cols.clear()

        if not self._driver or not self._database:
            return

        gen = self._load_gen
        for fn, cb in [
            (self._driver.get_tables, lambda d: self._on_tables(d, gen)),
            (self._driver.get_views,  lambda d: self._on_views(d, gen)),
        ]:
            w = SchemaWorker(fn, self._database)
            w.result.connect(cb)
            w.start()
            self._workers.append(w)

    def _on_tables(self, data, gen: int) -> None:
        if gen == self._load_gen:
            self._tables = [t.name for t in data]

    def _on_views(self, data, gen: int) -> None:
        if gen == self._load_gen:
            self._views = list(data)

    def _load_cols_for(self, table: str) -> None:
        if table in self._cols or table in self._pending_cols or not self._driver:
            return
        self._pending_cols.add(table)
        gen = self._load_gen
        w = SchemaWorker(self._driver.get_columns, self._database, table)
        w.result.connect(lambda cols: self._on_cols(table, cols, gen))
        w.start()
        self._workers.append(w)

    def _on_cols(self, table: str, cols, gen: int) -> None:
        self._pending_cols.discard(table)
        if gen == self._load_gen:
            self._cols[table] = [c.name for c in cols]

    # ── Prefix extraction ─────────────────────────────────────────────────────

    def _extract_prefix(self) -> tuple[str, str | None]:
        """
        Return (prefix, table_qualifier).
        - "SELECT" typed → ("SELECT", None)
        - "users.na" typed → ("na", "users")
        - "users." typed → ("", "users")
        - cursor inside string → ("", None)
        """
        tc = self._editor.textCursor()
        left = tc.block().text()[:tc.positionInBlock()]

        if _cursor_in_string(left):
            return ("", None)

        m = _PREFIX_RE.search(left)
        if not m:
            return ("", None)

        if m.group(1) is not None:
            return (m.group(2), m.group(1))   # col_prefix, table
        return (m.group(3), None)              # bare_prefix

    # ── Candidate building ────────────────────────────────────────────────────

    def _build_candidates(self, prefix: str, qualifier: str | None) -> list[str]:
        p = prefix.lower()

        if qualifier is not None:
            # Dot-context: column completions for the qualifying table/view
            canonical = next(
                (t for t in self._tables + self._views
                 if t.lower() == qualifier.lower()),
                qualifier,
            )
            self._load_cols_for(canonical)
            return [c for c in self._cols.get(canonical, [])
                    if c.lower().startswith(p)]

        # General: keywords + tables + views + all cached columns
        pool = SQL_KEYWORDS + self._tables + self._views
        for col_list in self._cols.values():
            pool.extend(col_list)

        seen: dict[str, None] = {}
        for word in pool:
            if word.lower().startswith(p):
                seen[word] = None
        return list(seen.keys())

    # ── Insertion ─────────────────────────────────────────────────────────────

    def _accept_completion(self, text: str | None) -> None:
        if not text:
            return
        self._popup.hide_popup()

        prefix, _ = self._extract_prefix()
        tc = self._editor.textCursor()
        tc.movePosition(
            QTextCursor.MoveOperation.Left,
            QTextCursor.MoveMode.KeepAnchor,
            len(prefix),
        )
        tc.insertText(text)
        self._editor.setTextCursor(tc)
        self._editor.setFocus()