"""Query execution helpers and actions."""

from __future__ import annotations

import re

from typing import TYPE_CHECKING, Any, Callable

from sqlit.domains.explorer.ui.tree import db_switching as tree_db_switching
from sqlit.domains.process_worker.ui.mixins.process_worker_lifecycle import (
    ProcessWorkerLifecycleMixin,
)
from sqlit.shared.ui.protocols import QueryMixinHost
from sqlit.shared.ui.spinner import Spinner

from .query_constants import MAX_FETCH_ROWS

if TYPE_CHECKING:
    from textual.worker import Worker

    from sqlit.domains.query.app.cancellable import CancellableQuery
    from sqlit.domains.query.app.query_service import QueryService
    from sqlit.domains.query.app.transaction import TransactionExecutor


_SCHEMA_CHANGE_RE = re.compile(
    r"\b(create|alter|drop|truncate|rename|comment|grant|revoke)\b",
    re.IGNORECASE,
)
_SQL_COMMENT_RE = re.compile(r"(--[^\n]*|/\*.*?\*/)", re.DOTALL)
_SQL_LITERAL_RE = re.compile(r"('([^']|'')*'|\"([^\"]|\"\")*\"|`[^`]*`|\[[^\]]*\])", re.DOTALL)


def _strip_sql_comments_and_literals(sql: str) -> str:
    sql = _SQL_COMMENT_RE.sub(" ", sql)
    return _SQL_LITERAL_RE.sub(" ", sql)


class QueryExecutionMixin(ProcessWorkerLifecycleMixin):
    """Mixin providing query execution actions."""

    _query_service: QueryService | None = None
    _query_service_db_type: str | None = None
    _query_worker: Worker[Any] | None = None
    _cancellable_query: CancellableQuery | None = None
    _transaction_executor: TransactionExecutor | None = None
    _transaction_executor_config: Any | None = None
    _query_spinner: Spinner | None = None
    _query_cursor_cache: dict[str, tuple[int, int]] | None = None
    _query_target_database: str | None = None
    _schema_worker: Any | None = None
    _schema_indexing: bool = False
    _pending_telescope_query: tuple[str, str] | None = None
    _telescope_auto_filter: bool = False

    def action_execute_query(self: QueryMixinHost) -> None:
        """Execute the current query."""
        self._execute_query_common(keep_insert_mode=False)

    def action_execute_query_insert(self: QueryMixinHost) -> None:
        """Execute query in INSERT mode without leaving it."""
        self._execute_query_common(keep_insert_mode=True)

    def action_execute_query_atomic(self: QueryMixinHost) -> None:
        """Execute query atomically (wrapped in BEGIN/COMMIT with rollback on error)."""
        if self.current_connection is None or self.current_provider is None:
            self.notify("Connect to a server to execute queries", severity="warning")
            return

        query = self._get_query_to_execute()

        if not query:
            self.notify("No query to execute", severity="warning")
            return

        def _proceed() -> None:
            if hasattr(self, "_query_worker") and self._query_worker is not None:
                self._query_worker.cancel()

            self._start_query_spinner()

            self._query_worker = self.run_worker(
                self._run_query_atomic_async(query),
                name="query_execution_atomic",
                exclusive=True,
            )

        self._maybe_confirm_query(query, _proceed)

    def action_execute_single_statement(self: QueryMixinHost) -> None:
        """Execute only the SQL statement at the current cursor position."""
        from sqlit.domains.query.app.multi_statement import find_statement_at_cursor

        if self.current_connection is None or self.current_provider is None:
            self.notify("Connect to a server to execute queries", severity="warning")
            return

        full_query = self.query_input.text
        if not full_query or not full_query.strip():
            self.notify("No query to execute", severity="warning")
            return

        # Get cursor position and find statement
        row, col = self.query_input.cursor_location
        result = find_statement_at_cursor(full_query, row, col)

        if result is None:
            self.notify("No statement found at cursor", severity="warning")
            return

        statement, _, _ = result

        if not statement.strip():
            self.notify("No statement found at cursor", severity="warning")
            return

        def _proceed() -> None:
            if hasattr(self, "_query_worker") and self._query_worker is not None:
                self._query_worker.cancel()

            self._start_query_spinner()

            self._query_worker = self.run_worker(
                self._run_query_async(statement, keep_insert_mode=False),
                name="query_execution_single",
                exclusive=True,
            )

        self._maybe_confirm_query(statement, _proceed)

    def _get_query_to_execute(self: QueryMixinHost) -> str:
        """Return selected text if present, otherwise the full query."""
        selection = self.query_input.selection
        if selection.start != selection.end:
            from sqlit.domains.query.editing import get_selection_text

            start = selection.start
            end = selection.end
            if start > end:
                start, end = end, start

            selected_sql = get_selection_text(
                self.query_input.text,
                start[0],
                start[1],
                end[0],
                end[1],
            )
            if selected_sql and selected_sql.strip():
                return selected_sql.strip()

        return self.query_input.text.strip()

    def _execute_query_common(self: QueryMixinHost, keep_insert_mode: bool) -> None:
        """Common query execution logic."""
        if self.current_connection is None or self.current_provider is None:
            self.notify("Connect to a server to execute queries", severity="warning")
            return

        query = self._get_query_to_execute()

        if not query:
            self.notify("No query to execute", severity="warning")
            return

        def _proceed() -> None:
            if hasattr(self, "_query_worker") and self._query_worker is not None:
                self._query_worker.cancel()

            self._start_query_spinner()

            self._query_worker = self.run_worker(
                self._run_query_async(query, keep_insert_mode),
                name="query_execution",
                exclusive=True,
            )

        self._maybe_confirm_query(query, _proceed)

    def _maybe_confirm_query(self: QueryMixinHost, query: str, proceed: Callable[[], None]) -> None:
        """Confirm query execution based on alert mode, then call proceed."""
        from sqlit.domains.query.app.alerts import (
            AlertMode,
            AlertSeverity,
            CONNECTION_ALERT_OPTION,
            DATABASE_ALERT_SETTING,
            classify_query_alert,
            format_alert_mode,
            lookup_database_override,
            resolve_alert_mode,
            should_confirm,
        )
        from sqlit.domains.query.app.multi_statement import get_executable_sql
        from sqlit.shared.ui.screens.confirm import ConfirmScreen

        global_mode = getattr(self.services.runtime, "query_alert_mode", 0) or 0
        connection_option = None
        config = getattr(self, "current_config", None)
        if config is not None:
            connection_option = config.get_option(CONNECTION_ALERT_OPTION)
        database: str | None = None
        get_db = getattr(self, "_get_effective_database", None)
        if callable(get_db):
            try:
                value = get_db()
                database = str(value) if value else None
            except Exception:
                database = None
        overrides: dict[str, Any] = {}
        try:
            raw_overrides = self.services.settings_store.get(DATABASE_ALERT_SETTING)
            if isinstance(raw_overrides, dict):
                overrides = raw_overrides
        except Exception:
            overrides = {}
        connection_name = getattr(config, "name", None) if config is not None else None
        database_override = lookup_database_override(overrides, connection_name, database)
        mode, _ = resolve_alert_mode(
            global_mode=global_mode,
            connection_option=connection_option,
            database_override=database_override,
        )

        if mode == AlertMode.OFF:
            proceed()
            return

        # Get the SQL that will actually execute (without comment-only statements)
        executable_sql = get_executable_sql(query)
        if not executable_sql:
            # Nothing to execute after filtering comments
            proceed()
            return

        severity = classify_query_alert(executable_sql)
        if severity == AlertSeverity.NONE or not should_confirm(mode, severity):
            proceed()
            return

        title = "Confirm query"
        if severity == AlertSeverity.DELETE:
            title = "Confirm DELETE query"
        elif severity == AlertSeverity.WRITE:
            title = "Confirm write query"

        def _on_result(confirmed: bool | None) -> None:
            if confirmed:
                proceed()
                return
            self.notify(
                f"Query cancelled (alert mode: {format_alert_mode(mode)})",
                severity="warning",
            )

        self.push_screen(
            ConfirmScreen(title, executable_sql, yes_label="Yes", no_label="No"),
            _on_result,
        )

    def _query_changes_schema(self: QueryMixinHost, query: str) -> bool:
        cleaned = _strip_sql_comments_and_literals(query)
        return bool(_SCHEMA_CHANGE_RE.search(cleaned))

    def _maybe_refresh_explorer_after_query(self: QueryMixinHost, query: str) -> None:
        if not self._query_changes_schema(query):
            return
        refresh = getattr(self, "_refresh_tree_after_schema_change", None)
        if callable(refresh):
            refresh()
            return
        action = getattr(self, "action_refresh_tree", None)
        if callable(action):
            action()

    def _start_query_spinner(self: QueryMixinHost) -> None:
        """Start the query execution spinner animation."""
        import time

        self.query_executing = True
        self._query_start_time = time.perf_counter()
        if self._query_spinner is not None:
            self._query_spinner.stop()
        self._query_spinner = Spinner(self, on_tick=lambda _: self._update_status_bar(), fps=30)
        self._query_spinner.start()
        self._update_footer_bindings()

    def _stop_query_spinner(self: QueryMixinHost) -> None:
        """Stop the query execution spinner animation."""
        self.query_executing = False
        if self._query_spinner is not None:
            self._query_spinner.stop()
            self._query_spinner = None
        if getattr(self, "_defer_schema_load", False):
            setattr(self, "_defer_schema_load", False)
            loader = getattr(self, "_load_schema_cache", None)
            if callable(loader):
                try:
                    loader()
                except Exception:
                    pass

    def _get_history_store(self: QueryMixinHost) -> Any:
        store = getattr(self, "_history_store", None)
        if store is not None:
            return store
        return self.services.history_store

    def _get_unsaved_history_store(self: QueryMixinHost) -> Any:
        store = getattr(self, "_unsaved_history_store", None)
        if store is None:
            from sqlit.domains.query.store.memory import InMemoryHistoryStore

            store = InMemoryHistoryStore()
            self._unsaved_history_store = store
        return store

    def _should_save_query_history(self: QueryMixinHost, config: Any) -> bool:
        """Return True if the connection is saved and history should be persisted."""
        name = getattr(config, "name", "")
        if not name:
            return False
        connections = getattr(self, "connections", None) or []
        return any(getattr(conn, "name", None) == name for conn in connections)

    def _save_query_history(self: QueryMixinHost, config: Any, query: str) -> None:
        """Save query history only for saved connections."""
        database = getattr(self, "_active_database", None) or ""
        if not database:
            endpoint = getattr(config, "tcp_endpoint", None)
            if endpoint:
                database = getattr(endpoint, "database", "") or ""
        if self._should_save_query_history(config):
            self._get_history_store().save_query(config.name, query, database=database)
            return
        self._get_unsaved_history_store().save_query(config.name, query)

    def _get_query_service(self: QueryMixinHost, provider: Any) -> QueryService:
        if self._query_service is None or (
            self._query_service_db_type is not None
            and self._query_service_db_type != provider.metadata.db_type
        ):
            from sqlit.domains.query.app.query_service import DialectQueryAnalyzer, QueryService

            self._query_service = QueryService(
                self._get_history_store(),
                analyzer=DialectQueryAnalyzer(provider.dialect),
            )
            self._query_service_db_type = provider.metadata.db_type
        return self._query_service

    def _get_transaction_executor(self: QueryMixinHost, config: Any, provider: Any) -> Any:
        """Get or create a TransactionExecutor for the current connection."""
        from sqlit.domains.query.app.transaction import TransactionExecutor

        # Create new executor if none exists or if config changed
        if self._transaction_executor is None or self._transaction_executor_config != config:
            if self._transaction_executor is not None:
                self._transaction_executor.close()
            self._transaction_executor = TransactionExecutor(config=config, provider=provider)
            self._transaction_executor_config = config
        return self._transaction_executor

    def _reset_transaction_executor(self: QueryMixinHost) -> None:
        """Reset the transaction executor (e.g., on disconnect)."""
        if self._transaction_executor is not None:
            self._transaction_executor.close()
            self._transaction_executor = None
        self._transaction_executor_config = None

    def _on_disconnect(self: QueryMixinHost) -> None:
        """Handle disconnect lifecycle event."""
        parent_disconnect = getattr(super(), "_on_disconnect", None)
        if callable(parent_disconnect):
            parent_disconnect()
        self._reset_transaction_executor()

    def _on_connect(self: QueryMixinHost) -> None:
        """Handle connect lifecycle event."""
        parent_connect = getattr(super(), "_on_connect", None)
        if callable(parent_connect):
            parent_connect()

        self._maybe_run_pending_telescope_query()

    def watch_current_connection(self: QueryMixinHost, old_value: Any, new_value: Any) -> None:
        self._maybe_run_pending_telescope_query()

    def watch_current_provider(self: QueryMixinHost, old_value: Any, new_value: Any) -> None:
        self._maybe_run_pending_telescope_query()

    def _on_connect_failed(self: QueryMixinHost, config: Any) -> None:
        pending = getattr(self, "_pending_telescope_query", None)
        if not pending:
            return
        if getattr(config, "name", None) == pending[0]:
            self._pending_telescope_query = None

    def _maybe_run_pending_telescope_query(self: QueryMixinHost) -> None:
        if not getattr(self, "_screen_stack", None):
            return
        pending = getattr(self, "_pending_telescope_query", None)
        if not pending:
            return
        if (
            self.current_connection is None
            or self.current_provider is None
            or self.current_config is None
        ):
            return
        connection_name = pending[0]
        query = pending[1]
        database = pending[2] if len(pending) > 2 else ""
        if self.current_config.name != connection_name:
            return
        self._pending_telescope_query = None
        if database:
            self._active_database = database
        self._apply_history_query(query)

    @property
    def in_transaction(self: QueryMixinHost) -> bool:
        """Whether we're currently in a transaction."""
        if self._transaction_executor is not None:
            return bool(self._transaction_executor.in_transaction)
        return False

    async def _run_query_async(self: QueryMixinHost, query: str, keep_insert_mode: bool) -> None:
        """Run query asynchronously using TransactionExecutor for transaction support."""
        import asyncio
        import time

        from sqlit.domains.query.app.multi_statement import (
            MultiStatementExecutor,
            split_statements,
        )
        from sqlit.domains.query.app.query_service import QueryResult, parse_use_statement
        from sqlit.domains.query.app.transaction import is_transaction_end, is_transaction_start

        provider = self.current_provider
        config = self.current_config

        if not provider or not config:
            self._display_query_error("Not connected")
            self._stop_query_spinner()
            return

        # If we have a target database from clicking a table in the tree,
        # use that database for the query execution (needed for Azure SQL)
        target_db = getattr(self, "_query_target_database", None)
        endpoint = config.tcp_endpoint
        current_db = endpoint.database if endpoint else ""
        if target_db and target_db != current_db:
            config = provider.apply_database_override(config, target_db)
        # Clear target database after use - it's only for the auto-generated query
        self._query_target_database = None

        # Apply active database to query execution (from USE statement or 'u' key)
        active_db = None
        if hasattr(self, "_get_effective_database"):
            active_db = self._get_effective_database()
        endpoint = config.tcp_endpoint
        current_db = endpoint.database if endpoint else ""
        if active_db and active_db != current_db and not target_db:
            config = provider.apply_database_override(config, active_db)

        # Handle USE database statements
        db_name = parse_use_statement(query)
        if db_name is not None:
            self._stop_query_spinner()
            self._display_non_query_result(0, 0)
            tree_db_switching.set_default_database(self, db_name)
            if keep_insert_mode:
                self._restore_insert_mode()
            return

        # Use TransactionExecutor for transaction-aware query execution
        executor = self._get_transaction_executor(config, provider)
        # Check if this is a multi-statement query (after filtering comment-only statements)
        from sqlit.domains.query.editing.comments import is_comment_only_statement

        statements = split_statements(query)
        executable_statements = [s for s in statements if not is_comment_only_statement(s)]
        is_multi_statement = len(executable_statements) > 1

        if is_multi_statement:
            self._pending_result_table_info = None
        elif executable_statements:
            if getattr(self, "_pending_result_table_info", None) is None:
                table_info = self._infer_result_table_info(executable_statements[0])
                if table_info is not None:
                    self._pending_result_table_info = table_info
                    prime = getattr(self, "_prime_result_table_columns", None)
                    if callable(prime):
                        prime(table_info)

        try:
            start_time = time.perf_counter()
            max_rows = self.services.runtime.max_rows or MAX_FETCH_ROWS

            use_process_worker = self._use_process_worker(provider)
            if use_process_worker and executable_statements:
                statement = executable_statements[0].strip()
                if self.in_transaction or is_transaction_start(statement) or is_transaction_end(statement):
                    use_process_worker = False

            if use_process_worker and not is_multi_statement:
                client = await self._get_process_worker_client_async()
                if client is None:
                    error = getattr(self, "_process_worker_client_error", None)
                    if error:
                        self.notify(
                            f"Process worker unavailable; falling back. ({error})",
                            severity="error",
                            title="Process Worker",
                        )
                    else:
                        self.notify(
                            "Process worker unavailable; falling back.",
                            severity="error",
                            title="Process Worker",
                        )
                    use_process_worker = False

                if use_process_worker:
                    outcome = await asyncio.to_thread(client.execute, query, config, max_rows)
                    if outcome.cancelled:
                        return
                    if outcome.error:
                        self._display_query_error(outcome.error)
                        return

                    try:
                        await asyncio.to_thread(self._save_query_history, config, query)
                    except Exception:
                        pass
                    result = outcome.result
                    elapsed_ms = outcome.elapsed_ms

                    if isinstance(result, QueryResult):
                        await self._display_query_results(
                            result.columns,
                            result.rows,
                            result.row_count,
                            result.truncated,
                            elapsed_ms,
                        )
                    else:
                        self._display_non_query_result(result.rows_affected, elapsed_ms)
                    self._maybe_refresh_explorer_after_query(query)
                    if keep_insert_mode:
                        self._restore_insert_mode()
                    return

            if is_multi_statement:
                # Multi-statement execution with stacked results
                multi_executor = MultiStatementExecutor(executor)
                multi_result = await asyncio.to_thread(
                    multi_executor.execute,
                    query,
                    max_rows,
                )
                elapsed_ms = (time.perf_counter() - start_time) * 1000

                try:
                    await asyncio.to_thread(self._save_query_history, config, query)
                except Exception:
                    pass
                self._display_multi_statement_results(multi_result, elapsed_ms)
                self._maybe_refresh_explorer_after_query(query)
            else:
                # Single statement - existing behavior
                result = await asyncio.to_thread(
                    executor.execute,
                    query,
                    max_rows,
                )
                elapsed_ms = (time.perf_counter() - start_time) * 1000

                try:
                    await asyncio.to_thread(self._save_query_history, config, query)
                except Exception:
                    pass

                if isinstance(result, QueryResult):
                    await self._display_query_results(
                        result.columns, result.rows, result.row_count, result.truncated, elapsed_ms
                    )
                else:
                    self._display_non_query_result(result.rows_affected, elapsed_ms)
                self._maybe_refresh_explorer_after_query(query)

            if keep_insert_mode:
                self._restore_insert_mode()

        except RuntimeError as e:
            if "cancelled" in str(e).lower():
                pass  # Already handled by action_cancel_query
            else:
                self._display_query_error(str(e))
        except Exception as e:
            self._display_query_error(str(e))
        finally:
            self._stop_query_spinner()

    async def _run_query_atomic_async(self: QueryMixinHost, query: str) -> None:
        """Run query atomically (BEGIN/COMMIT with rollback on error)."""
        import asyncio
        import time

        from sqlit.domains.query.app.multi_statement import MultiStatementResult
        from sqlit.domains.query.app.query_service import QueryResult
        from sqlit.domains.query.app.transaction import TransactionExecutor

        provider = self.current_provider
        config = self.current_config

        if not provider or not config:
            self._display_query_error("Not connected")
            self._stop_query_spinner()
            return

        # Apply active database if set
        active_db = None
        if hasattr(self, "_get_effective_database"):
            active_db = self._get_effective_database()
        endpoint = config.tcp_endpoint
        current_db = endpoint.database if endpoint else ""
        if active_db and active_db != current_db:
            config = provider.apply_database_override(config, active_db)

        # Create a dedicated executor for atomic execution
        executor = TransactionExecutor(config=config, provider=provider)
        try:
            start_time = time.perf_counter()
            max_rows = self.services.runtime.max_rows or MAX_FETCH_ROWS
            result = await asyncio.to_thread(
                executor.atomic_execute,
                query,
                max_rows,
            )
            elapsed_ms = (time.perf_counter() - start_time) * 1000

            try:
                await asyncio.to_thread(self._save_query_history, config, query)
            except Exception:
                pass

            if isinstance(result, MultiStatementResult):
                # Multi-statement atomic execution
                self._display_multi_statement_results(result, elapsed_ms)
                if result.has_error:
                    self.notify("Transaction rolled back (error in statement)", severity="error")
                else:
                    self.notify("Query executed atomically (committed)", severity="information")
                    self._maybe_refresh_explorer_after_query(query)
            elif isinstance(result, QueryResult):
                await self._display_query_results(
                    result.columns, result.rows, result.row_count, result.truncated, elapsed_ms
                )
                self.notify("Query executed atomically (committed)", severity="information")
                self._maybe_refresh_explorer_after_query(query)
            else:
                self._display_non_query_result(result.rows_affected, elapsed_ms)
                self.notify("Query executed atomically (committed)", severity="information")
                self._maybe_refresh_explorer_after_query(query)

        except Exception as e:
            self._display_query_error(f"Transaction rolled back: {e}")
        finally:
            executor.close()
            self._stop_query_spinner()

    def _restore_insert_mode(self: QueryMixinHost) -> None:
        """Restore INSERT mode after query execution (called on main thread)."""
        from sqlit.core.vim import VimMode

        self.vim_mode = VimMode.INSERT
        self.query_input.read_only = False
        self.query_input.focus()
        self._update_footer_bindings()
        self._update_vim_mode_visuals()

    def action_cancel_query(self: QueryMixinHost) -> None:
        """Cancel the currently running query."""
        if not getattr(self, "query_executing", False):
            self.notify("No query running")
            return

        client = getattr(self, "_process_worker_client", None)
        if client is not None:
            try:
                client.cancel_current()
            except Exception:
                pass

        if hasattr(self, "_cancellable_query") and self._cancellable_query is not None:
            self._cancellable_query.cancel()

        if hasattr(self, "_query_worker") and self._query_worker is not None:
            self._query_worker.cancel()
            self._query_worker = None

        self._stop_query_spinner()

        self._replace_results_table(["Status"], [("Query cancelled",)])

        self.notify("Query cancelled", severity="warning")

    def action_cancel_operation(self: QueryMixinHost) -> None:
        """Cancel any running operation (query or schema indexing)."""
        cancelled = False

        # Cancel query if running
        if getattr(self, "query_executing", False):
            # Cancel the cancellable query (closes dedicated connection)
            client = getattr(self, "_process_worker_client", None)
            if client is not None:
                try:
                    client.cancel_current()
                except Exception:
                    pass
            if hasattr(self, "_cancellable_query") and self._cancellable_query is not None:
                self._cancellable_query.cancel()

            if hasattr(self, "_query_worker") and self._query_worker is not None:
                self._query_worker.cancel()
                self._query_worker = None
            self._stop_query_spinner()

            # Update results table to show cancelled state
            self._replace_results_table(["Status"], [("Query cancelled",)])
            cancelled = True

        # Cancel schema indexing if running
        if getattr(self, "_schema_indexing", False):
            if hasattr(self, "_schema_worker") and self._schema_worker is not None:
                self._schema_worker.cancel()
                self._schema_worker = None
            self._stop_schema_spinner()
            cancelled = True

        if cancelled:
            self.notify("Operation cancelled", severity="warning")
        else:
            self.notify("No operation running")

    def action_clear_query(self: QueryMixinHost) -> None:
        """Clear the query input."""
        self.query_input.text = ""

    def action_new_query(self: QueryMixinHost) -> None:
        """Start a new query (clear input and results)."""
        self.query_input.text = ""
        self._replace_results_table([], [])

    def _apply_history_query(self: QueryMixinHost, query: str) -> None:
        """Load a query into the editor and restore cursor position if possible."""
        # Initialize cursor cache if needed
        if self._query_cursor_cache is None:
            self._query_cursor_cache = {}
        cursor_cache = self._query_cursor_cache

        # Save current query's cursor position before switching
        current_query = self.query_input.text
        if current_query:
            cursor_cache[current_query] = self.query_input.cursor_location

        # Set new query text
        self.query_input.text = query

        # Restore cursor position if we have it cached, otherwise go to end
        if query in cursor_cache:
            self.query_input.cursor_location = cursor_cache[query]
        else:
            lines = query.split("\n")
            last_line = len(lines) - 1
            last_col = len(lines[-1]) if lines else 0
            self.query_input.cursor_location = (last_line, last_col)

        # Focus query input - this triggers on_descendant_focus which updates footer bindings
        self.query_input.focus()

    PREFERRED_EDITOR_SETTING = "preferred_editor"

    def action_edit_query_in_editor(self: QueryMixinHost) -> None:
        """Open the current query in the user's terminal editor.

        Pops up an editor picker on first use (or whenever the saved
        preference is missing from PATH). On editor exit, replaces the
        query buffer and — when the connection is saved — saves the
        edited text to history regardless of execution.
        """
        from sqlit.domains.query.app.editor import resolve_editor

        settings = self.services.settings_store.load_all()
        preferred = settings.get(self.PREFERRED_EDITOR_SETTING)
        editor_cmd = resolve_editor(preferred)
        if editor_cmd is None:
            self._open_editor_picker()
            return
        self._run_external_editor(editor_cmd)

    def _open_editor_picker(self: QueryMixinHost) -> None:
        from sqlit.domains.query.app.editor import detect_editors

        from ..screens import EditorPickerScreen

        if not any(e.is_installed for e in detect_editors()):
            self.notify(
                "No terminal editor detected. Install nvim, vim, hx, micro, or "
                "nano and try again (or set $EDITOR).",
                severity="error",
                timeout=8,
            )
            return

        self.notify(
            "No preferred editor set — pick one (saved for next time).",
            severity="information",
        )

        def on_pick(result: str | None) -> None:
            if not result:
                self.notify("Editor selection cancelled.", severity="warning")
                return
            settings = self.services.settings_store.load_all()
            settings[self.PREFERRED_EDITOR_SETTING] = result
            self.services.settings_store.save_all(settings)
            self.notify(f"Editor set to '{result}'. Launching…", severity="information")
            self._run_external_editor(result)

        self.push_screen(EditorPickerScreen(), on_pick)

    def _run_external_editor(self: QueryMixinHost, editor_cmd: str) -> None:
        import os
        import subprocess
        import tempfile
        from datetime import datetime

        from sqlit.domains.query.app.editor import build_editor_argv

        original = self.query_input.text
        timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
        fd, path_str = tempfile.mkstemp(
            prefix=f"sqlit-edit-{timestamp}-",
            suffix=".sql",
        )
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as f:
                f.write(original)
            argv = build_editor_argv(editor_cmd) + [path_str]
            try:
                with self.suspend():
                    subprocess.run(argv, check=False)
            except Exception as exc:
                self.notify(f"Failed to launch editor '{editor_cmd}': {exc}", severity="error")
                return

            try:
                with open(path_str, encoding="utf-8") as f:
                    edited = f.read()
            except OSError as exc:
                self.notify(f"Could not read edited file: {exc}", severity="error")
                return
        finally:
            try:
                os.unlink(path_str)
            except OSError:
                pass

        # Trim a single trailing newline that most editors add on save.
        if edited.endswith("\n"):
            edited = edited[:-1]

        if edited == original:
            self.notify("Editor closed — no changes.", severity="information")
            return

        self._apply_history_query(edited)
        self.notify("Query updated from editor.", severity="information")

        if self.current_config and self._should_save_query_history(self.current_config):
            database = getattr(self, "_active_database", None) or ""
            self._get_history_store().save_query(
                self.current_config.name, edited, database=database
            )

    def action_show_history(self: QueryMixinHost) -> None:
        """Show query history for the current connection."""
        if not self.current_config:
            self.notify("Not connected", severity="warning")
            return

        from ..screens import QueryHistoryScreen

        starred_store = self.services.starred_store
        if self._should_save_query_history(self.current_config):
            history = self._get_history_store().load_for_connection(self.current_config.name)
        else:
            history = self._get_unsaved_history_store().load_for_connection(self.current_config.name)
        starred = starred_store.load_for_connection(self.current_config.name)
        self.push_screen(
            QueryHistoryScreen(history, self.current_config.name, starred),
            self._handle_history_result,
        )

    def _handle_history_result(self: QueryMixinHost, result: Any) -> None:
        """Handle the result from the history screen."""
        if result is None:
            return

        action, data = result
        if action == "select":
            self._apply_history_query(data)
        elif action == "delete":
            self._delete_history_entry(data)
            self.action_show_history()
        elif action == "toggle_star":
            self._toggle_star(data)
            self.action_show_history()

    def action_telescope(self: QueryMixinHost) -> None:
        """Show query history across all connections."""
        self._show_telescope(auto_open_filter=False)

    def action_telescope_filter(self: QueryMixinHost) -> None:
        """Show query history across all connections with filter active."""
        self._show_telescope(auto_open_filter=True)

    def _show_telescope(self: QueryMixinHost, *, auto_open_filter: bool) -> None:
        """Open telescope with optional filter preset."""
        from ..screens import QueryHistoryScreen

        connection_map = self._get_telescope_connection_map()
        available_connections = set(connection_map.keys())

        history_store = self._get_history_store()
        if hasattr(history_store, "load_all"):
            history = history_store.load_all()
        else:
            history = []
            for config in connection_map.values():
                history.extend(history_store.load_for_connection(config.name))
            history.sort(key=lambda entry: entry.timestamp, reverse=True)

        unsaved_store = getattr(self, "_unsaved_history_store", None)
        if unsaved_store is not None and hasattr(unsaved_store, "load_all"):
            history.extend(unsaved_store.load_all())

        if available_connections:
            history = [
                entry for entry in history
                if getattr(entry, "connection_name", None) in available_connections
            ]
        history.sort(key=lambda entry: entry.timestamp, reverse=True)
        connection_labels = {
            name: self._format_telescope_connection_label(config)
            for name, config in connection_map.items()
        }
        starred_by_connection = self._load_starred_by_connection(connection_map)
        self._telescope_auto_filter = auto_open_filter

        self.push_screen(
            QueryHistoryScreen(
                history,
                "All Servers",
                None,
                multi_connection=True,
                connection_labels=connection_labels,
                starred_by_connection=starred_by_connection,
                auto_open_filter=auto_open_filter,
            ),
            self._handle_telescope_result,
        )

    def _handle_telescope_result(self: QueryMixinHost, result: Any) -> None:
        """Handle the result from the telescope screen."""
        if result is None:
            return

        action, data = result
        if action == "select":
            query = data.get("query", "")
            connection_name = data.get("connection_name", "")
            database = data.get("database", "")
            self._run_telescope_query(connection_name, query, database=database)
        elif action == "delete":
            timestamp = data.get("timestamp", "")
            connection_name = data.get("connection_name", "")
            if timestamp and connection_name:
                self._get_history_store().delete_entry(connection_name, timestamp)
            if self._telescope_auto_filter:
                self.action_telescope_filter()
            else:
                self.action_telescope()
        elif action == "toggle_star":
            query = data.get("query", "")
            connection_name = data.get("connection_name", "")
            if query and connection_name:
                is_now_starred = self.services.starred_store.toggle_star(connection_name, query)
                if is_now_starred:
                    self.notify("Query starred")
                else:
                    self.notify("Query unstarred")
            if self._telescope_auto_filter:
                self.action_telescope_filter()
            else:
                self.action_telescope()

    def _run_telescope_query(
        self: QueryMixinHost, connection_name: str, query: str, *, database: str = ""
    ) -> None:
        if not query or not connection_name:
            return

        config = self._get_telescope_connection_map().get(connection_name)
        if config is None:
            self.notify(f"Connection '{connection_name}' not found", severity="warning")
            return

        # Resolve the active database from when the query was originally executed.
        # For old history entries without a database field, try to find one
        # from other entries for the same connection.
        if not database:
            database = self._infer_database_from_history(connection_name)

        self._apply_history_query(query)

        if (
            self.current_connection is not None
            and self.current_config is not None
            and self.current_config.name == connection_name
        ):
            if database:
                self._active_database = database
            self._pending_telescope_query = None
            return

        # Store database in the pending tuple so _maybe_run_pending_telescope_query
        # can restore it AFTER the connection callback resets _active_database.
        self._pending_telescope_query = (connection_name, query, database)
        self._connect_like_explorer(connection_name, config)

    def _get_telescope_connection_map(self: QueryMixinHost) -> dict[str, Any]:
        connection_map = {config.name: config for config in getattr(self, "connections", [])}
        for config in (
            getattr(self, "_direct_connection_config", None),
            getattr(self, "current_config", None),
        ):
            if config and config.name not in connection_map:
                connection_map[config.name] = config
        return connection_map

    def _connect_like_explorer(self: QueryMixinHost, connection_name: str, config: Any) -> None:
        node = None
        object_tree = getattr(self, "object_tree", None)
        if object_tree is not None:
            try:
                stack = [object_tree.root]
                while stack:
                    current = stack.pop()
                    for child in current.children:
                        if getattr(self, "_get_node_kind", None) and self._get_node_kind(child) != "connection":
                            stack.append(child)
                            continue
                        data = getattr(child, "data", None)
                        node_config = getattr(data, "config", None)
                        if node_config and node_config.name == connection_name:
                            node = child
                            stack = []
                            break
                        stack.append(child)
            except Exception:
                node = None

        if node is not None and hasattr(self, "_activate_tree_node"):
            self._activate_tree_node(node)
            return

        self.connect_to_server(config)

    def _infer_database_from_history(self: QueryMixinHost, connection_name: str) -> str:
        """Try to find the most recently used database for a connection from history."""
        try:
            history_store = self._get_history_store()
            entries = history_store.load_for_connection(connection_name)
            for entry in entries:
                db = getattr(entry, "database", "") or ""
                if db:
                    return db
        except Exception:
            pass
        return ""

    def _format_telescope_connection_label(self: QueryMixinHost, config: Any) -> str:
        endpoint = getattr(config, "tcp_endpoint", None)
        if endpoint is None:
            file_endpoint = getattr(config, "file_endpoint", None)
            if file_endpoint and getattr(file_endpoint, "path", ""):
                return file_endpoint.path
            return getattr(config, "name", "")
        database = getattr(config, "database", "")
        return database or getattr(config, "name", "")

    def _load_starred_by_connection(self: QueryMixinHost, connection_map: dict[str, Any]) -> dict[str, set[str]]:
        starred_store = self.services.starred_store
        loader = getattr(starred_store, "load_all", None)
        if callable(loader):
            return loader()
        return {
            name: starred_store.load_for_connection(name)
            for name in connection_map.keys()
        }

    def _delete_history_entry(self: QueryMixinHost, timestamp: str) -> None:
        """Delete a specific history entry by timestamp."""
        if not self.current_config:
            return
        self._get_history_store().delete_entry(self.current_config.name, timestamp)

    def _toggle_star(self: QueryMixinHost, query: str) -> None:
        """Toggle star status for a query."""
        if not self.current_config:
            return

        is_now_starred = self.services.starred_store.toggle_star(self.current_config.name, query)
        if is_now_starred:
            self.notify("Query starred")
        else:
            self.notify("Query unstarred")
