server-manager/gui/tabs/query_tab.py

"""
Query tab — SQL database interaction with tree explorer, editor, results grid, and export.
"""

import csv
import time
import threading
import tkinter as tk
from tkinter import ttk, filedialog

import customtkinter as ctk

from core.i18n import t
from core.icons import icon_text
from core.sql_client import SQLClient

_TREE_THEME_APPLIED = False


def _apply_db_tree_theme():
    global _TREE_THEME_APPLIED
    if _TREE_THEME_APPLIED:
        return
    _TREE_THEME_APPLIED = True

    style = ttk.Style()
    try:
        style.theme_use("clam")
    except Exception:
        pass
    style.configure(
        "DBTree.Treeview",
        background="#1e1e1e",
        foreground="#dcdcdc",
        fieldbackground="#1e1e1e",
        borderwidth=0,
        font=("Consolas", 11),
        rowheight=24,
    )
    style.map(
        "DBTree.Treeview",
        background=[("selected", "#3b82f6")],
        foreground=[("selected", "#ffffff")],
    )
    style.layout("DBTree.Treeview", [
        ("DBTree.Treeview.treearea", {"sticky": "nswe"}),
    ])
    style.configure(
        "Results.Treeview",
        background="#1e1e1e",
        foreground="#dcdcdc",
        fieldbackground="#1e1e1e",
        borderwidth=0,
        font=("Consolas", 11),
        rowheight=24,
    )
    style.configure(
        "Results.Treeview.Heading",
        background="#2b2b2b",
        foreground="#9ca3af",
        borderwidth=0,
        font=("Segoe UI", 10, "bold"),
        relief="flat",
    )
    style.map(
        "Results.Treeview",
        background=[("selected", "#3b82f6")],
        foreground=[("selected", "#ffffff")],
    )
    style.map(
        "Results.Treeview.Heading",
        background=[("active", "#333333")],
    )
    style.layout("Results.Treeview", [
        ("Results.Treeview.treearea", {"sticky": "nswe"}),
    ])


class QueryTab(ctk.CTkFrame):
    def __init__(self, master, store):
        super().__init__(master, fg_color="transparent")
        self._current_alias: str | None = None
        self.store = store
        self._client: SQLClient | None = None
        self._db_lock = threading.Lock()
        self._results: list[list] = []
        self._columns: list[str] = []
        self._executing = False
        self._current_db: str = ""

        _apply_db_tree_theme()
        self._build_ui()

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

    def _build_ui(self):
        # === Main PanedWindow: tree | editor+results ===
        self._paned = ttk.PanedWindow(self, orient="horizontal")
        self._paned.pack(fill="both", expand=True, padx=5, pady=5)

        # ── Left panel: Database tree ──
        left_frame = tk.Frame(self._paned, bg="#1a1a1a", width=260)
        self._paned.add(left_frame, weight=0)

        # Tree toolbar
        toolbar = tk.Frame(left_frame, bg="#1a1a1a")
        toolbar.pack(fill="x", padx=2, pady=(4, 2))

        self._tree_title = tk.Label(
            toolbar, text=t("tree_databases"), bg="#1a1a1a", fg="#9ca3af",
            font=("Segoe UI", 10, "bold"), anchor="w",
        )
        self._tree_title.pack(side="left", padx=4)

        self._refresh_btn = tk.Button(
            toolbar, text="\u21bb", bg="#2b2b2b", fg="#9ca3af",
            font=("Segoe UI", 11), bd=0, padx=6, pady=0,
            activebackground="#3b3b3b", activeforeground="#dcdcdc",
            command=self._refresh_tree,
        )
        self._refresh_btn.pack(side="right", padx=4)

        # Tree + scrollbar
        tree_container = tk.Frame(left_frame, bg="#1e1e1e")
        tree_container.pack(fill="both", expand=True, padx=2, pady=(0, 2))

        self._tree_scroll = ttk.Scrollbar(tree_container, orient="vertical")
        self._tree_scroll.pack(side="right", fill="y")

        self._db_tree = ttk.Treeview(
            tree_container, show="tree", selectmode="browse",
            style="DBTree.Treeview",
            yscrollcommand=self._tree_scroll.set,
        )
        self._db_tree.pack(fill="both", expand=True)
        self._tree_scroll.config(command=self._db_tree.yview)

        # Tree bindings
        self._db_tree.bind("<<TreeviewOpen>>", self._on_tree_expand)
        self._db_tree.bind("<Double-1>", self._on_tree_dblclick)
        self._db_tree.bind("<Button-3>", self._on_tree_rightclick)
        self._db_tree.bind("<<TreeviewSelect>>", self._on_tree_select)

        # ── Right panel: editor + buttons + results ──
        right_frame = tk.Frame(self._paned, bg="#1a1a1a")
        self._paned.add(right_frame, weight=1)

        right_ctk = ctk.CTkFrame(right_frame, fg_color="transparent")
        right_ctk.pack(fill="both", expand=True)

        # === SQL Editor ===
        editor_frame = ctk.CTkFrame(right_ctk, fg_color="transparent")
        editor_frame.pack(fill="x", padx=8, pady=(8, 4))

        self._editor = ctk.CTkTextbox(
            editor_frame,
            font=ctk.CTkFont(family="Consolas", size=13),
            height=140,
            wrap="none",
        )
        self._editor.pack(fill="both", expand=True)
        self._editor.insert("0.0", t("query_editor_placeholder"))
        self._editor.bind("<FocusIn>", self._on_editor_focus)
        self._editor.bind("<F5>", lambda e: self._execute_query())
        self._editor.bind("<Control-Return>", lambda e: self._execute_query())

        # === Button row ===
        btn_row = ctk.CTkFrame(right_ctk, fg_color="transparent")
        btn_row.pack(fill="x", padx=8, pady=4)

        self._exec_btn = ctk.CTkButton(
            btn_row,
            text=icon_text("execute", t("query_execute")),
            command=self._execute_query,
            width=130,
            fg_color="#2563eb",
            hover_color="#1d4ed8",
        )
        self._exec_btn.pack(side="left", padx=(0, 6))

        self._clear_btn = ctk.CTkButton(
            btn_row,
            text=icon_text("clear", t("query_clear")),
            command=self._clear_all,
            width=80,
            fg_color="#6b7280",
            hover_color="#4b5563",
        )
        self._clear_btn.pack(side="left", padx=(0, 6))

        self._export_btn = ctk.CTkButton(
            btn_row,
            text=icon_text("save", t("query_export_csv")),
            command=self._export_csv,
            width=110,
            fg_color="#059669",
            hover_color="#047857",
        )
        self._export_btn.pack(side="left")

        # === Results area (Treeview) ===
        results_frame = ctk.CTkFrame(right_ctk, fg_color="transparent")
        results_frame.pack(fill="both", expand=True, padx=8, pady=(4, 4))

        self._res_xscroll = ttk.Scrollbar(results_frame, orient="horizontal")
        self._res_xscroll.pack(side="bottom", fill="x")

        self._res_yscroll = ttk.Scrollbar(results_frame, orient="vertical")
        self._res_yscroll.pack(side="right", fill="y")

        self._results_tree = ttk.Treeview(
            results_frame,
            show="headings",
            style="Results.Treeview",
            xscrollcommand=self._res_xscroll.set,
            yscrollcommand=self._res_yscroll.set,
        )
        self._results_tree.pack(fill="both", expand=True)

        self._res_xscroll.config(command=self._results_tree.xview)
        self._res_yscroll.config(command=self._results_tree.yview)

        # === Status bar ===
        self._status_label = ctk.CTkLabel(
            right_ctk, text="", anchor="w",
            font=ctk.CTkFont(size=12), text_color="#9ca3af",
        )
        self._status_label.pack(fill="x", padx=10, pady=(0, 6))

    # ── Editor placeholder logic ───────────────────────────────────

    def _on_editor_focus(self, event=None):
        content = self._editor.get("0.0", "end").strip()
        if content == t("query_editor_placeholder"):
            self._editor.delete("0.0", "end")

    # ── Connection management (thread-safe) ────────────────────────

    def _ensure_connected(self):
        """Check connection, reconnect if broken. MUST be called under _db_lock."""
        if not self._client:
            raise ConnectionError("No client")
        try:
            if not self._client.check_connection():
                raise Exception("lost")
        except Exception:
            server = self.store.get_server(self._current_alias)
            if not server:
                raise ConnectionError("Server not found")
            self._client.disconnect()
            self._client = SQLClient(server)
            if not self._client.connect():
                self._client = None
                raise ConnectionError("Reconnect failed")
            if self._current_db:
                try:
                    self._client.switch_database(self._current_db)
                except Exception:
                    pass

    # ── Server / database connection ───────────────────────────────

    def set_server(self, alias: str | None):
        """Called when user selects a server in the sidebar."""
        self._current_alias = alias
        self._disconnect()
        self._clear_results()
        self._clear_tree()
        self._set_status("")
        self._current_db = ""

        if not alias:
            return

        self._set_status(f"Connecting to {alias}...")
        threading.Thread(
            target=self._connect_and_load_tree,
            args=(alias,),
            daemon=True,
        ).start()

    def _connect_and_load_tree(self, alias: str):
        """Background: create SQLClient, fetch databases, populate tree."""
        try:
            server = self.store.get_server(alias)
            if not server:
                self.after(0, lambda: self._set_status(t("query_error"), error=True))
                return

            client = SQLClient(server)
            if not client.connect():
                self.after(0, lambda: self._set_status(
                    t("query_error") + ": connection failed", error=True))
                return

            self._db_lock.acquire(timeout=10)
            try:
                self._client = client
                databases = client.list_databases()
            finally:
                self._db_lock.release()

            def _update():
                if self._current_alias != alias:
                    return
                self._populate_tree_databases(databases)
                if databases:
                    self._current_db = databases[0]
                self._set_status(t("tree_connected"))

            self.after(0, _update)
        except Exception as exc:
            self.after(0, lambda: self._set_status(str(exc), error=True))

    def _disconnect(self):
        with self._db_lock:
            if self._client:
                try:
                    self._client.disconnect()
                except Exception:
                    pass
                self._client = None

    # ── Database tree ──────────────────────────────────────────────

    def _clear_tree(self):
        self._db_tree.delete(*self._db_tree.get_children())

    def _populate_tree_databases(self, databases: list[str]):
        self._db_tree.delete(*self._db_tree.get_children())
        for db in databases:
            iid = f"db:{db}"
            self._db_tree.insert(
                "", "end", iid=iid,
                text=f"  \U0001f5c4  {db}",
                open=False,
            )
            self._db_tree.insert(iid, "end", iid=f"_dummy:{db}",
                                 text=f"  {t('tree_loading')}")

    def _on_tree_expand(self, event):
        """Lazy-load children when a node is expanded for the first time."""
        node = self._db_tree.focus()
        if not node:
            return
        children = self._db_tree.get_children(node)
        if len(children) == 1 and str(children[0]).startswith("_dummy:"):
            if node.startswith("db:"):
                self._load_tables(node)
            elif node.startswith("tbl:"):
                self._load_columns(node)

    def _load_tables(self, db_node: str):
        """Background: load tables for a database node."""
        db_name = db_node[3:]

        def _do():
            try:
                if not self._db_lock.acquire(timeout=10):
                    return
                try:
                    self._ensure_connected()
                    tables = self._client.list_tables(db_name)
                finally:
                    self._db_lock.release()

                def _update():
                    for ch in self._db_tree.get_children(db_node):
                        self._db_tree.delete(ch)
                    if not tables:
                        self._db_tree.insert(
                            db_node, "end",
                            text=f"  {t('tree_no_tables')}",
                        )
                        return
                    for tbl in tables:
                        tbl_iid = f"tbl:{db_name}.{tbl}"
                        self._db_tree.insert(
                            db_node, "end", iid=tbl_iid,
                            text=f"  \U0001f4cb  {tbl}",
                            open=False,
                        )
                        self._db_tree.insert(
                            tbl_iid, "end",
                            iid=f"_dummy:{db_name}.{tbl}",
                            text=f"  {t('tree_loading')}",
                        )

                self.after(0, _update)
            except Exception as exc:
                self.after(0, lambda: self._set_status(f"Error: {exc}", error=True))

        threading.Thread(target=_do, daemon=True).start()

    def _load_columns(self, tbl_node: str):
        """Background: load columns for a table node."""
        parts = tbl_node[4:].split(".", 1)
        if len(parts) != 2:
            return
        db_name, tbl_name = parts

        def _do():
            try:
                if not self._db_lock.acquire(timeout=10):
                    return
                try:
                    self._ensure_connected()
                    self._client.switch_database(db_name)
                    columns = self._client.describe_table(tbl_name)
                finally:
                    self._db_lock.release()

                def _update():
                    for ch in self._db_tree.get_children(tbl_node):
                        self._db_tree.delete(ch)
                    if not columns:
                        self._db_tree.insert(
                            tbl_node, "end",
                            text=f"  {t('tree_no_columns')}",
                        )
                        return
                    for col in columns:
                        name = col.get("name", "?")
                        dtype = col.get("type", "")
                        key = col.get("key", "")
                        icon = "\U0001f511" if "PRI" in str(key) else "\U0001f4dd"
                        label = f"  {icon}  {name}  ({dtype})"
                        col_iid = f"col:{db_name}.{tbl_name}.{name}"
                        self._db_tree.insert(
                            tbl_node, "end", iid=col_iid,
                            text=label,
                        )

                self.after(0, _update)
            except Exception as exc:
                self.after(0, lambda: self._set_status(f"Error: {exc}", error=True))

        threading.Thread(target=_do, daemon=True).start()

    def _on_tree_select(self, event):
        """When user clicks a node, just remember the active DB (no network calls)."""
        node = self._db_tree.focus()
        if not node:
            return
        if node.startswith("db:"):
            self._current_db = node[3:]
        elif node.startswith("tbl:"):
            self._current_db = node[4:].split(".", 1)[0]

    def _on_tree_dblclick(self, event):
        """Double-click on table → SELECT * LIMIT 1000."""
        node = self._db_tree.identify_row(event.y)
        if not node:
            return

        if node.startswith("tbl:"):
            parts = node[4:].split(".", 1)
            if len(parts) == 2:
                db_name, tbl_name = parts
                self._select_top(tbl_name, db_name)
        elif node.startswith("col:"):
            parts = node[4:].split(".")
            if len(parts) >= 3:
                col_name = parts[-1]
                self._editor.insert("insert", f"`{col_name}`")

    def _on_tree_rightclick(self, event):
        """Context menu on tree nodes."""
        node = self._db_tree.identify_row(event.y)
        if not node:
            return
        self._db_tree.selection_set(node)

        menu = tk.Menu(self, tearoff=0, bg="#2b2b2b", fg="#dcdcdc",
                       activebackground="#3b82f6", activeforeground="#ffffff",
                       font=("Segoe UI", 10))

        if node.startswith("db:"):
            db_name = node[3:]
            menu.add_command(label=t("tree_use_db"),
                             command=lambda: self._use_database(db_name))
            menu.add_command(label=t("tree_refresh"),
                             command=lambda: self._refresh_node(node))
            menu.add_separator()
            menu.add_command(label=t("tree_copy_name"),
                             command=lambda: self._copy_to_clipboard(db_name))
        elif node.startswith("tbl:"):
            parts = node[4:].split(".", 1)
            if len(parts) == 2:
                db_name, tbl_name = parts
                menu.add_command(label=t("tree_select_top"),
                                 command=lambda: self._select_top(tbl_name, db_name))
                menu.add_command(label=t("tree_describe"),
                                 command=lambda: self._describe_table_query(tbl_name, db_name))
                menu.add_separator()
                menu.add_command(label=t("tree_copy_name"),
                                 command=lambda: self._copy_to_clipboard(tbl_name))
                menu.add_command(label=t("tree_refresh"),
                                 command=lambda: self._refresh_node(node))
        elif node.startswith("col:"):
            parts = node[4:].split(".")
            if len(parts) >= 3:
                col_name = parts[-1]
                menu.add_command(label=t("tree_copy_name"),
                                 command=lambda: self._copy_to_clipboard(col_name))
                menu.add_command(
                    label=t("tree_select_column"),
                    command=lambda: self._insert_text(f"`{col_name}`"),
                )

        menu.tk_popup(event.x_root, event.y_root)

    def _refresh_tree(self):
        """Reload entire tree from server."""
        if not self._client:
            return
        self._set_status(t("tree_loading"))

        def _do():
            try:
                if not self._db_lock.acquire(timeout=10):
                    self.after(0, lambda: self._set_status("Database busy", error=True))
                    return
                try:
                    self._ensure_connected()
                    dbs = self._client.list_databases()
                finally:
                    self._db_lock.release()
                self.after(0, lambda: self._populate_tree_databases(dbs))
                self.after(0, lambda: self._set_status(t("tree_connected")))
            except Exception as exc:
                self.after(0, lambda: self._set_status(str(exc), error=True))

        threading.Thread(target=_do, daemon=True).start()

    def _refresh_node(self, node: str):
        """Refresh a specific node (database or table)."""
        if node.startswith("db:"):
            for ch in self._db_tree.get_children(node):
                self._db_tree.delete(ch)
            db_name = node[3:]
            self._db_tree.insert(node, "end", iid=f"_dummy:{db_name}",
                                 text=f"  {t('tree_loading')}")
            self._db_tree.item(node, open=True)
            self._load_tables(node)
        elif node.startswith("tbl:"):
            for ch in self._db_tree.get_children(node):
                self._db_tree.delete(ch)
            parts = node[4:]
            self._db_tree.insert(node, "end", iid=f"_dummy:{parts}",
                                 text=f"  {t('tree_loading')}")
            self._db_tree.item(node, open=True)
            self._load_columns(node)

    # ── Tree actions ───────────────────────────────────────────────

    def _use_database(self, db_name: str):
        self._current_db = db_name
        self._set_status(f"Database: {db_name}")

    def _select_top(self, table: str, db_name: str = ""):
        """Insert SELECT * FROM table LIMIT 1000 and execute."""
        if db_name:
            self._current_db = db_name
        sql = f"SELECT * FROM `{table}` LIMIT 1000"
        self._editor.delete("0.0", "end")
        self._editor.insert("0.0", sql)
        self._execute_query()

    def _describe_table_query(self, table: str, db_name: str = ""):
        """Insert DESCRIBE/SHOW COLUMNS and execute."""
        if db_name:
            self._current_db = db_name
        if self._client and self._client._type in ("mariadb", "mysql"):
            sql = f"SHOW COLUMNS FROM `{table}`"
        else:
            sql = f"SELECT * FROM `{table}` WHERE 1=0"
        self._editor.delete("0.0", "end")
        self._editor.insert("0.0", sql)
        self._execute_query()

    def _copy_to_clipboard(self, text: str):
        self.clipboard_clear()
        self.clipboard_append(text)

    def _insert_text(self, text: str):
        self._editor.insert("insert", text)

    # ── Query execution ────────────────────────────────────────────

    def _execute_query(self):
        """Run the SQL query in a background thread."""
        if self._executing or not self._client:
            return

        sql = self._editor.get("0.0", "end").strip()
        if not sql or sql == t("query_editor_placeholder"):
            return

        self._executing = True
        self._exec_btn.configure(state="disabled")
        self._set_status("Executing...")

        threading.Thread(
            target=self._run_query,
            args=(sql,),
            daemon=True,
        ).start()

    def _run_query(self, sql: str):
        """Background thread: execute SQL with lock, auto-reconnect, auto-switch DB."""
        start = time.perf_counter()
        try:
            acquired = self._db_lock.acquire(timeout=10)
            if not acquired:
                self.after(0, lambda: self._set_status("Database busy, try again", error=True))
                return
            try:
                self._ensure_connected()
                if self._current_db:
                    self._client.switch_database(self._current_db)
                result = self._client.execute_query(sql)
            finally:
                self._db_lock.release()

            columns = result["columns"]
            rows = result["rows"]
            elapsed = time.perf_counter() - start

            def _update():
                self._columns = columns
                self._results = rows
                self._populate_results(columns, rows)
                self._set_status(
                    t("query_status_rows").format(
                        rows=len(rows), time=f"{elapsed:.3f}"
                    )
                )

            self.after(0, _update)
        except Exception as exc:
            self.after(0, lambda: self._set_status(
                f"{t('query_error')}: {exc}", error=True))
        finally:
            def _reset():
                self._executing = False
                self._exec_btn.configure(state="normal")
            self.after(0, _reset)

    # ── Results Treeview population ────────────────────────────────

    def _populate_results(self, columns: list[str], rows: list[list]):
        """Clear and populate the results Treeview."""
        self._results_tree.delete(*self._results_tree.get_children())

        if not columns:
            self._results_tree["columns"] = ()
            return

        self._results_tree["columns"] = columns
        for col in columns:
            self._results_tree.heading(col, text=col, anchor="w")
            self._results_tree.column(col, width=120, minwidth=60, anchor="w")

        for row in rows:
            display = [str(v) if v is not None else "NULL" for v in row]
            self._results_tree.insert("", "end", values=display)

    def _clear_results(self):
        self._results_tree.delete(*self._results_tree.get_children())
        self._results_tree["columns"] = ()
        self._columns = []
        self._results = []

    # ── Button actions ─────────────────────────────────────────────

    def _clear_all(self):
        self._editor.delete("0.0", "end")
        self._clear_results()
        self._set_status("")

    def _export_csv(self):
        if not self._columns or not self._results:
            return

        path = filedialog.asksaveasfilename(
            defaultextension=".csv",
            filetypes=[("CSV files", "*.csv"), ("All files", "*.*")],
            title=t("query_export_csv"),
        )
        if not path:
            return

        try:
            with open(path, "w", newline="", encoding="utf-8") as f:
                writer = csv.writer(f)
                writer.writerow(self._columns)
                for row in self._results:
                    writer.writerow(
                        [str(v) if v is not None else "" for v in row]
                    )
            self._set_status(f"Exported {len(self._results)} rows to {path}")
        except Exception as exc:
            self._set_status(f"{t('query_error')}: {exc}", error=True)

    # ── Status bar ─────────────────────────────────────────────────

    def _set_status(self, text: str, error: bool = False):
        color = "#ef4444" if error else "#9ca3af"
        self._status_label.configure(text=text, text_color=color)