from __future__ import annotations

from collections.abc import Generator, Iterable
from dataclasses import dataclass
from datetime import datetime
import functools
from functools import partial
import html
from importlib import resources
import itertools
import os
import random
import re
import shutil
import subprocess
import sys
import time
import traceback
from typing import TYPE_CHECKING, Any
from xml.parsers.expat import ExpatError

import click
import prompt_toolkit
from prompt_toolkit.application.current import get_app
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory, ThreadedAutoSuggest
from prompt_toolkit.completion import DynamicCompleter
from prompt_toolkit.enums import DEFAULT_BUFFER, EditingMode
from prompt_toolkit.filters import Condition, has_focus, is_done
from prompt_toolkit.formatted_text import (
    ANSI,
    HTML,
    FormattedText,
    to_formatted_text,
    to_plain_text,
)
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.layout.processors import (
    ConditionalProcessor,
    HighlightMatchingBracketProcessor,
    TabsProcessor,
)
from prompt_toolkit.lexers import PygmentsLexer
from prompt_toolkit.output import ColorDepth
from prompt_toolkit.shortcuts import CompleteStyle, PromptSession
import pymysql
from pymysql.cursors import Cursor

import mycli as mycli_package
from mycli.clibuffer import cli_is_multiline
from mycli.clistyle import style_factory_ptoolkit
from mycli.clitoolbar import create_toolbar_tokens_func, get_vi_mode
from mycli.compat import WIN
from mycli.constants import (
    DEFAULT_HOST,
    DEFAULT_WIDTH,
    ER_MUST_CHANGE_PASSWORD,
    HOME_URL,
    ISSUES_URL,
)
from mycli.key_bindings import mycli_bindings
from mycli.lexer import MyCliLexer
from mycli.packages import special
from mycli.packages.filepaths import dir_path_exists
from mycli.packages.hybrid_redirection import get_redirect_components, is_redirect_command
from mycli.packages.interactive_utils import confirm, confirm_destructive_query
from mycli.packages.key_binding_utils import (
    handle_clip_command,
    handle_editor_command,
)
from mycli.packages.polars_transform import (
    PolarsTransform,
    PolarsTransformError,
    parse_polars_transform,
    prepare_polars_transform,
    run_polars_transform,
)
from mycli.packages.ptoolkit.history import FRECENCY_HISTORY_ENTRIES, FRECENCY_REFRESH_INTERVAL, FileHistoryWithTimestamp
from mycli.packages.special.utils import format_uptime, get_ssl_version, get_uptime, get_warning_count
from mycli.packages.sql_utils import (
    extract_new_password,
    is_dropping_database,
    is_mutating,
    is_password_change,
    is_sandbox_allowed,
    is_select,
    need_completion_refresh,
    need_completion_reset,
)
from mycli.packages.sqlresult import SQLResult
from mycli.packages.string_utils import sanitize_terminal_title
from mycli.sqlexecute import SQLExecute
from mycli.types import Query

if TYPE_CHECKING:
    from prompt_toolkit.formatted_text import AnyFormattedText

    from mycli.client import MyCli


SUPPORT_INFO = f"Home: {HOME_URL}\nBug tracker: {ISSUES_URL}"
MIN_COMPLETION_TRIGGER = 1
_PROMPT_TARGETS: dict[int, 'MyCli'] = {}


@dataclass(slots=True)
class ReplState:
    iterations: int = 0
    mutating: bool = False
    buffer_text: str | None = None


@Condition
def complete_while_typing_filter() -> bool:
    """Whether enough characters have been typed to trigger completion.

    Written in a verbose way, with a string slice, for efficiency."""
    if MIN_COMPLETION_TRIGGER <= 1:
        return True
    app = get_app()
    text = app.current_buffer.text.lstrip()
    if text.startswith('/') and not text.startswith('/*') and ' ' not in text:
        return True
    completing_config_property = bool(re.match(r'^(?:/|\\)config\s+get\s+', text, re.IGNORECASE))
    text_len = len(text)
    if text_len < MIN_COMPLETION_TRIGGER:
        return False
    last_word = text[-MIN_COMPLETION_TRIGGER:]
    if len(last_word) == text_len:
        return text_len >= MIN_COMPLETION_TRIGGER
    # does \. make sense with text[:6] ?
    if text[:6].lower() in ['source', r'\.', '/.']:
        # Different word characters for paths; see comment below.
        # In fact, it might be nice if paths had a different threshold.
        return not bool(re.search(r'[\s!-,:-@\[-^\{\}-]', last_word))
    else:
        # This is "whitespace and all punctuation except underscore and backtick"
        # acting as word breaks, but it would be neat if we could complete differently
        # when inside a backtick, accepting all legal characters towards the trigger
        # limit. We would have to parse the statement, or at least go back more
        # characters, costing performance. This still works within a backtick! So
        # long as there are three trailing non-punctuation characters.
        punctuation_check = last_word.replace('.', '') if completing_config_property else last_word
        return not bool(re.search(r'[\s!-/:-@\[-^\{-~]', punctuation_check))


def _create_history(mycli: 'MyCli') -> FileHistoryWithTimestamp | None:
    history_file = os.path.expanduser(os.environ.get('MYCLI_HISTFILE', mycli.config['main'].get('history_file', '~/.mycli-history')))
    if dir_path_exists(history_file):
        frecency_history_entries = int(mycli.config['main'].get('frecency_history_entries', FRECENCY_HISTORY_ENTRIES) or 0)
        frecency_refresh_interval = int(mycli.config['main'].get('frecency_refresh_interval', FRECENCY_REFRESH_INTERVAL) or 0)
        return FileHistoryWithTimestamp(
            history_file,
            frecency_history_entries=frecency_history_entries,
            frecency_refresh_interval=frecency_refresh_interval,
        )

    mycli.echo(
        f'Error: Unable to open the history file "{history_file}". Your query history will not be saved.',
        err=True,
        fg='red',
    )
    return None


def _show_startup_banner(
    mycli: 'MyCli',
    sqlexecute: SQLExecute,
) -> None:
    if mycli.verbosity < 0:
        return

    if sqlexecute.server_info is not None:
        print(sqlexecute.server_info)
    print('mycli', mycli_package.__version__)
    print(SUPPORT_INFO)
    if random.random() <= 0.25:
        print('Thanks to the sponsor —', _sponsors_picker())
    elif random.random() <= 0.5:
        print('Thanks to the contributor —', _contributors_picker())
    else:
        print('Tip —', _tips_picker())


def set_all_external_titles(mycli: 'MyCli') -> None:
    set_external_terminal_tab_title(mycli)
    set_external_terminal_window_title(mycli)
    set_external_multiplex_window_title(mycli)
    set_external_multiplex_pane_title(mycli)


def set_external_terminal_tab_title(mycli: 'MyCli') -> None:
    if not mycli.terminal_tab_title_format:
        return
    if not mycli.prompt_session:
        return
    if not sys.stderr.isatty():
        return
    title = sanitize_terminal_title(
        render_prompt_string(
            mycli,
            mycli.terminal_tab_title_format,
            mycli.prompt_session.app.render_counter,
        )
    )
    print(f'\x1b]1;{title}\a', file=sys.stderr, end='')
    sys.stderr.flush()


def set_external_terminal_window_title(mycli: 'MyCli') -> None:
    if not mycli.terminal_window_title_format:
        return
    if not mycli.prompt_session:
        return
    if not sys.stderr.isatty():
        return
    title = sanitize_terminal_title(
        render_prompt_string(
            mycli,
            mycli.terminal_window_title_format,
            mycli.prompt_session.app.render_counter,
        )
    )
    print(f'\x1b]2;{title}\a', file=sys.stderr, end='')
    sys.stderr.flush()


def set_external_multiplex_window_title(mycli: 'MyCli') -> None:
    if not mycli.multiplex_window_title_format:
        return
    if not os.getenv('TMUX'):
        return
    if not mycli.prompt_session:
        return
    title = sanitize_terminal_title(
        render_prompt_string(
            mycli,
            mycli.multiplex_window_title_format,
            mycli.prompt_session.app.render_counter,
        )
    )
    try:
        subprocess.run(
            ['tmux', 'rename-window', title],
            check=False,
            stdin=subprocess.DEVNULL,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
    except FileNotFoundError:
        pass


def set_external_multiplex_pane_title(mycli: 'MyCli') -> None:
    if not mycli.multiplex_pane_title_format:
        return
    if not os.getenv('TMUX'):
        return
    if not mycli.prompt_session:
        return
    if not sys.stderr.isatty():
        return
    title = sanitize_terminal_title(
        render_prompt_string(
            mycli,
            mycli.multiplex_pane_title_format,
            mycli.prompt_session.app.render_counter,
        )
    )
    print(f'\x1b]2;{title}\x1b\\', file=sys.stderr, end='')
    sys.stderr.flush()


def get_custom_toolbar(
    mycli: 'MyCli',
    toolbar_format: str,
) -> FormattedText:
    if not mycli.prompt_session:
        return to_formatted_text('')
    if not mycli.prompt_session.app:
        return to_formatted_text('')
    if mycli.prompt_session.app.current_buffer.text:
        return mycli.last_custom_toolbar_message
    mycli.last_custom_toolbar_message = render_prompt_string(
        mycli,
        toolbar_format,
        mycli.prompt_session.app.render_counter,
    )
    return mycli.last_custom_toolbar_message


def maybe_html_escape(string: str, is_html: bool) -> str:
    if is_html:
        return html.escape(string, quote=False)
    return string


@functools.lru_cache(maxsize=256)
def render_prompt_string(
    mycli: 'MyCli',
    string: str,
    _render_counter: int,
) -> FormattedText:
    sqlexecute = mycli.sqlexecute
    assert sqlexecute is not None
    if mycli.login_path and mycli.login_path_as_host:
        prompt_host = mycli.login_path
    elif sqlexecute.host is not None:
        prompt_host = sqlexecute.host
    else:
        prompt_host = DEFAULT_HOST
    short_prompt_host, _, _ = prompt_host.partition('.')
    if re.match(r'^[\d\.]+$', short_prompt_host):
        short_prompt_host = prompt_host
    now = datetime.now()
    species_name = sqlexecute.server_info.species.name if sqlexecute.server_info and sqlexecute.server_info.species else 'MySQL'
    strings = string.split('\\\\')
    is_html = strings[0].startswith('\\<html>')
    strings = [x.replace('\\u', maybe_html_escape(sqlexecute.user or '(none)', is_html)) for x in strings]
    strings = [x.replace('\\h', maybe_html_escape(prompt_host or '(none)', is_html)) for x in strings]
    strings = [x.replace('\\H', maybe_html_escape(short_prompt_host or '(none)', is_html)) for x in strings]
    strings = [x.replace('\\d', maybe_html_escape(sqlexecute.dbname or '(none)', is_html)) for x in strings]
    strings = [x.replace('\\t', maybe_html_escape(species_name, is_html)) for x in strings]
    strings = [x.replace('\\n', '\n') for x in strings]
    strings = [x.replace('\\D', maybe_html_escape(now.strftime('%a %b %d %H:%M:%S %Y'), is_html)) for x in strings]
    strings = [x.replace('\\m', maybe_html_escape(now.strftime('%M'), is_html)) for x in strings]
    strings = [x.replace('\\P', maybe_html_escape(now.strftime('%p'), is_html)) for x in strings]
    strings = [x.replace('\\R', maybe_html_escape(now.strftime('%H'), is_html)) for x in strings]
    strings = [x.replace('\\r', maybe_html_escape(now.strftime('%I'), is_html)) for x in strings]
    strings = [x.replace('\\s', maybe_html_escape(now.strftime('%S'), is_html)) for x in strings]
    strings = [x.replace('\\p', maybe_html_escape(str(sqlexecute.port), is_html)) for x in strings]
    strings = [
        x.replace('\\j', maybe_html_escape(os.path.basename(sqlexecute.socket or '(none)').replace('\\', '/'), is_html)) for x in strings
    ]
    strings = [x.replace('\\J', maybe_html_escape((sqlexecute.socket or '(none)').replace('\\', '/'), is_html)) for x in strings]
    strings = [
        x.replace('\\k', maybe_html_escape(os.path.basename(sqlexecute.socket or str(sqlexecute.port)).replace('\\', '/'), is_html))
        for x in strings
    ]
    strings = [
        x.replace('\\K', maybe_html_escape((sqlexecute.socket or str(sqlexecute.port)).replace('\\', '/'), is_html)) for x in strings
    ]
    strings = [x.replace('\\A', maybe_html_escape(mycli.dsn_alias or '(none)', is_html)) for x in strings]
    strings = [x.replace('\\_', ' ') for x in strings]

    checker_string = ' '.join(strings)
    if r'\e' in checker_string:
        if mycli.prompt_session:
            edit_mode = mycli.prompt_session.editing_mode.value.lower()
            if edit_mode == 'vi':
                edit_mode += ':' + get_vi_mode()
        else:
            edit_mode = mycli.key_bindings.lower()
        strings = [x.replace(r'\e', maybe_html_escape(edit_mode, is_html)) for x in strings]
    if hasattr(sqlexecute, 'conn') and sqlexecute.conn is not None:
        if '\\y' in checker_string:
            with sqlexecute.conn.cursor() as cur:
                strings = [x.replace('\\y', maybe_html_escape(str(get_uptime(cur)) or '(none)', is_html)) for x in strings]
        if '\\Y' in checker_string:
            with sqlexecute.conn.cursor() as cur:
                strings = [x.replace('\\Y', maybe_html_escape(format_uptime(str(get_uptime(cur))) or '(none)', is_html)) for x in strings]
    else:
        strings = [x.replace('\\y', '(none)') for x in strings]
        strings = [x.replace('\\Y', '(none)') for x in strings]

    if hasattr(sqlexecute, 'conn') and sqlexecute.conn is not None:
        if '\\T' in checker_string:
            with sqlexecute.conn.cursor() as cur:
                strings = [x.replace('\\T', maybe_html_escape(get_ssl_version(cur) or '(none)', is_html)) for x in strings]
    else:
        strings = [x.replace('\\T', '(none)') for x in strings]

    if hasattr(sqlexecute, 'conn') and sqlexecute.conn is not None:
        if '\\w' in checker_string:
            with sqlexecute.conn.cursor() as cur:
                strings = [x.replace('\\w', maybe_html_escape(str(get_warning_count(cur) or '(none)'), is_html)) for x in strings]
    else:
        strings = [x.replace('\\w', '(none)') for x in strings]
    if hasattr(sqlexecute, 'conn') and sqlexecute.conn is not None:
        if '\\W' in checker_string:
            with sqlexecute.conn.cursor() as cur:
                strings = [x.replace('\\W', maybe_html_escape(str(get_warning_count(cur) or ''), is_html)) for x in strings]
    else:
        strings = [x.replace('\\W', '') for x in strings]

    if is_html:
        strings[0] = strings[0].removeprefix('\\<html>')
        strings[-1] = strings[-1].removesuffix('\\</html>')
    elif '\\x1b' in checker_string:
        strings = [x.replace('\\x1b', '\x1b') for x in strings]

    strings = [re.sub(r'\\(.)', r'(unknown prompt format string: \\\1)', x) for x in strings]

    string = '\\'.join(strings)

    if is_html:
        try:
            formatted_string = to_formatted_text(HTML(string))
        except (ExpatError, ValueError):
            formatted_string = to_formatted_text(HTML('(cannot parse HTML prompt string)'))
    else:
        formatted_string = to_formatted_text(ANSI(string))

    return formatted_string


def _get_prompt_message(
    mycli: 'MyCli',
    app: prompt_toolkit.application.application.Application,
) -> FormattedText:
    if app.current_buffer.text:
        return mycli.last_prompt_message

    prompt = render_prompt_string(mycli, mycli.prompt_format, app.render_counter)
    prompt_plain = to_plain_text(prompt)
    if mycli.prompt_format == mycli.default_prompt and len(prompt_plain) > mycli.max_len_prompt:
        prompt = render_prompt_string(mycli, mycli.default_prompt_splitln, app.render_counter)
        prompt_plain = to_plain_text(prompt)
        mycli.prompt_lines = prompt_plain.count('\n') + 1
    if not mycli.prompt_lines:
        mycli.prompt_lines = prompt_plain.count('\n') + 1

    mycli.last_prompt_message = prompt
    return mycli.last_prompt_message


def _get_continuation(
    mycli: 'MyCli',
    width: int,
    _two: int,
    _three: int,
) -> AnyFormattedText:
    if mycli.multiline_continuation_char == '':
        continuation = ''
    elif mycli.multiline_continuation_char:
        left_padding = width - len(mycli.multiline_continuation_char)
        continuation = ' ' * max((left_padding - 1), 0) + mycli.multiline_continuation_char + ' '
    else:
        continuation = ' '
    return [('class:continuation', continuation)]


def _output_results(
    mycli: 'MyCli',
    state: ReplState,
    results: Iterable[SQLResult],
    start: float,
) -> None:
    sqlexecute = mycli.sqlexecute
    assert sqlexecute is not None

    result_iterator = iter(results)
    try:
        first_result = next(result_iterator)
    except StopIteration:
        return

    if first_result.command is not None and first_result.command['name'] == 'source_page':
        if special.is_redirected():
            _output_results(mycli, state, result_iterator, start)
            return
        paged_output = _single_paged_output_results(mycli, state, result_iterator, start)
        try:
            click.echo_via_pager(paged_output)
        except KeyboardInterrupt:
            pass
        finally:
            paged_output.close()
        return

    results = itertools.chain([first_result], result_iterator)

    result_count = 0
    watch_count = 0
    for result in results:
        mycli.logger.debug('preamble: %r', result.preamble)
        mycli.logger.debug('header: %r', result.header)
        mycli.logger.debug('rows: %r', result.rows)
        mycli.logger.debug('status: %r', result.status)
        mycli.logger.debug('command: %r', result.command)
        threshold = 1000
        if result.command is not None and result.command['name'] == 'source_show':
            click.secho(f"> {result.command['text']}")
            continue
        if result.command is not None and result.command['name'] == 'set_buffer':
            state.buffer_text = str(result.command['text'])
            continue
        if result.command is not None and result.command['name'] == 'watch':
            if watch_count > 0:
                try:
                    watch_seconds = float(result.command['seconds'])
                    start += watch_seconds
                except ValueError as e:
                    mycli.echo(f'Invalid watch sleep time provided ({e}).', err=True, fg='red')
                    sys.exit(1)
            else:
                watch_count += 1

        if is_select(result.status_plain) and isinstance(result.rows, Cursor) and result.rows.rowcount > threshold:
            mycli.echo(
                f'The result set has more than {threshold} rows.',
                fg='red',
            )
            if not confirm('Do you want to continue?'):
                mycli.echo('Aborted!', err=True, fg='red')
                break

        if mycli.auto_vertical_output:
            if mycli.prompt_session is not None:
                max_width = mycli.prompt_session.output.get_size().columns
            else:
                max_width = DEFAULT_WIDTH
        else:
            max_width = None

        formatted = mycli.format_sqlresult(
            result,
            is_expanded=special.is_expanded_output(),
            is_redirected=special.is_redirected(),
            null_string=mycli.null_string,
            numeric_alignment=mycli.numeric_alignment,
            binary_display=mycli.binary_display,
            max_width=max_width,
        )

        duration = time.time() - start
        try:
            if result_count > 0:
                mycli.echo('')
            try:
                mycli.output(formatted, result)
            except KeyboardInterrupt:
                pass
            if mycli.beep_after_seconds > 0 and duration >= mycli.beep_after_seconds:
                assert mycli.prompt_session is not None
                mycli.prompt_session.output.bell()
            if special.is_timing_enabled():
                mycli.output_timing(f'Time: {duration:0.03f}s')
        except KeyboardInterrupt:
            pass

        start = time.time()
        result_count += 1
        state.mutating = state.mutating or is_mutating(result.status_plain)

        if special.is_show_warnings_enabled() and isinstance(result.rows, Cursor) and result.rows.warning_count > 0:
            warnings = sqlexecute.run('SHOW WARNINGS')
            warnings_duration = time.time() - start
            saw_warning = False
            for warning in warnings:
                saw_warning = True
                formatted = mycli.format_sqlresult(
                    warning,
                    is_expanded=special.is_expanded_output(),
                    is_redirected=special.is_redirected(),
                    null_string=mycli.null_string,
                    numeric_alignment=mycli.numeric_alignment,
                    binary_display=mycli.binary_display,
                    max_width=max_width,
                    is_warnings_style=True,
                )
                mycli.echo('')
                mycli.output(formatted, warning, is_warnings_style=True)

            if saw_warning and special.is_timing_enabled():
                mycli.output_timing(f'Time: {warnings_duration:0.03f}s', is_warnings_style=True)


def _single_paged_output_results(
    mycli: 'MyCli',
    state: ReplState,
    results: Iterable[SQLResult],
    start: float,
) -> Generator[str, None, None]:
    """Render results lazily through one pager session."""
    sqlexecute = mycli.sqlexecute
    assert sqlexecute is not None
    result_iterator = iter(results)
    result_count = 0
    watch_count = 0
    try:
        for result in result_iterator:
            mycli.logger.debug('preamble: %r', result.preamble)
            mycli.logger.debug('header: %r', result.header)
            mycli.logger.debug('rows: %r', result.rows)
            mycli.logger.debug('status: %r', result.status)
            mycli.logger.debug('command: %r', result.command)

            if result.command is not None and result.command['name'] == 'source_show':
                yield f"> {result.command['text']}\n"
                continue
            if result.command is not None and result.command['name'] == 'set_buffer':
                state.buffer_text = str(result.command['text'])
                continue
            if result.command is not None and result.command['name'] == 'watch':
                if watch_count > 0:
                    try:
                        start += float(result.command['seconds'])
                    except ValueError as error:
                        message = f'Invalid watch sleep time provided ({error}).'
                        mycli.log_output(message)
                        yield f'{message}\n'
                        return
                else:
                    watch_count += 1

            if mycli.auto_vertical_output:
                if mycli.prompt_session is not None:
                    max_width = mycli.prompt_session.output.get_size().columns
                else:
                    max_width = DEFAULT_WIDTH
            else:
                max_width = None

            formatted = mycli.format_sqlresult(
                result,
                is_expanded=special.is_expanded_output(),
                is_redirected=False,
                null_string=mycli.null_string,
                numeric_alignment=mycli.numeric_alignment,
                binary_display=mycli.binary_display,
                max_width=max_width,
            )
            duration = time.time() - start

            if result_count > 0:
                mycli.log_output('')
                yield '\n'
            for line in formatted:
                mycli.log_output(line)
                special.write_tee(line)
                special.write_once(line)
                special.write_pipe_once(line)
                yield f'{line}\n'
            if result.status:
                mycli.log_output(result.status_plain)
                yield f'{result.status_plain}\n'

            if mycli.beep_after_seconds > 0 and duration >= mycli.beep_after_seconds:
                assert mycli.prompt_session is not None
                mycli.prompt_session.output.bell()
            if special.is_timing_enabled():
                timing = f'Time: {duration:0.03f}s'
                mycli.log_output(timing)
                yield f'{timing}\n'

            start = time.time()
            result_count += 1
            state.mutating = state.mutating or is_mutating(result.status_plain)

            if special.is_show_warnings_enabled() and isinstance(result.rows, Cursor) and result.rows.warning_count > 0:
                warnings = sqlexecute.run('SHOW WARNINGS')
                warnings_duration = time.time() - start
                saw_warning = False
                for warning in warnings:
                    saw_warning = True
                    warning_output = mycli.format_sqlresult(
                        warning,
                        is_expanded=special.is_expanded_output(),
                        is_redirected=False,
                        null_string=mycli.null_string,
                        numeric_alignment=mycli.numeric_alignment,
                        binary_display=mycli.binary_display,
                        max_width=max_width,
                        is_warnings_style=True,
                    )
                    mycli.log_output('')
                    yield '\n'
                    for line in warning_output:
                        mycli.log_output(line)
                        special.write_tee(line)
                        special.write_once(line)
                        special.write_pipe_once(line)
                        yield f'{line}\n'
                    if warning.status:
                        mycli.log_output(warning.status_plain)
                        yield f'{warning.status_plain}\n'

                if saw_warning and special.is_timing_enabled():
                    timing = f'Time: {warnings_duration:0.03f}s'
                    mycli.log_output(timing)
                    yield f'{timing}\n'
    finally:
        close = getattr(result_iterator, 'close', None)
        if close is not None:
            close()


def _keepalive_hook(
    mycli: 'MyCli',
    _context: Any,
) -> None:
    if mycli.keepalive_ticks is None:
        return
    if mycli.keepalive_ticks < 1:
        return

    mycli._keepalive_counter += 1
    if mycli._keepalive_counter > mycli.keepalive_ticks:
        mycli._keepalive_counter = 0
        mycli.logger.debug('keepalive ping')
        try:
            assert mycli.sqlexecute is not None
            assert mycli.sqlexecute.conn is not None
            mycli.sqlexecute.conn.ping(reconnect=False)
        except Exception as e:
            mycli.logger.debug('keepalive ping error %r', e)


def _build_prompt_session(
    mycli: 'MyCli',
    state: ReplState,
    history: FileHistoryWithTimestamp | None,
    key_bindings: KeyBindings,
) -> None:
    if mycli.toolbar_format.lower() == 'none':
        get_toolbar_tokens = None
    else:
        get_toolbar_tokens = create_toolbar_tokens_func(
            mycli,
            lambda: state.iterations == 0,
            mycli.toolbar_format,
            partial(get_custom_toolbar, mycli),
        )

    if mycli.wider_completion_menu:
        complete_style = CompleteStyle.MULTI_COLUMN
    else:
        complete_style = CompleteStyle.COLUMN

    with mycli._completer_lock:
        if mycli.key_bindings == 'vi':
            editing_mode = EditingMode.VI
        else:
            editing_mode = EditingMode.EMACS

        mycli.prompt_session = PromptSession(
            color_depth=ColorDepth.DEPTH_24_BIT if 'truecolor' in os.getenv('COLORTERM', '').lower() else None,
            lexer=PygmentsLexer(MyCliLexer),
            reserve_space_for_menu=mycli.get_reserved_space(),
            prompt_continuation=lambda width, two, three: _get_continuation(mycli, width, two, three),
            bottom_toolbar=get_toolbar_tokens,  # type: ignore[arg-type]
            complete_style=complete_style,
            input_processors=[
                ConditionalProcessor(
                    processor=HighlightMatchingBracketProcessor(chars='[](){}'),
                    filter=has_focus(DEFAULT_BUFFER) & ~is_done,
                ),
                # renders tab as four spaces rather than ^I control character
                TabsProcessor(char1=' ', char2=' '),
            ],
            tempfile_suffix='.sql',
            completer=DynamicCompleter(lambda: mycli.completer),
            complete_in_thread=True,
            history=history,
            auto_suggest=ThreadedAutoSuggest(AutoSuggestFromHistory()),
            complete_while_typing=complete_while_typing_filter,
            multiline=cli_is_multiline(mycli),
            style=style_factory_ptoolkit(mycli.syntax_style, mycli.cli_style),
            include_default_pygments_style=False,
            key_bindings=key_bindings,
            enable_open_in_editor=True,
            enable_system_prompt=True,
            enable_suspend=True,
            editing_mode=editing_mode,
            search_ignore_case=True,
        )

        if mycli.key_bindings == 'vi':
            mycli.prompt_session.app.ttimeoutlen = mycli.vi_ttimeoutlen
        else:
            mycli.prompt_session.app.ttimeoutlen = mycli.emacs_ttimeoutlen


def _one_iteration(
    mycli: 'MyCli',
    state: ReplState,
    text: str | None = None,
) -> None:
    sqlexecute = mycli.sqlexecute
    assert sqlexecute is not None

    inputhook = partial(_keepalive_hook, mycli) if mycli.keepalive_ticks and mycli.keepalive_ticks >= 1 else None

    if text is None:
        try:
            assert mycli.prompt_session is not None
            loaded_message_fn = partial(_get_prompt_message, mycli, mycli.prompt_session.app)
            if state.buffer_text is None:
                text = mycli.prompt_session.prompt(
                    inputhook=inputhook,
                    message=loaded_message_fn,
                )
            else:
                default = state.buffer_text
                state.buffer_text = None
                text = mycli.prompt_session.prompt(
                    default=default,
                    inputhook=inputhook,
                    message=loaded_message_fn,
                )
        except KeyboardInterrupt:
            return

        special.set_expanded_output(False)
        special.set_forced_horizontal_output(False)
        special.set_explorer_output(False)

        try:
            text = handle_editor_command(
                mycli,
                text,
                inputhook,
                loaded_message_fn,
            )
        except RuntimeError as e:
            mycli.logger.error('sql: %r, error: %r', text, e)
            mycli.logger.error('traceback: %r', traceback.format_exc())
            mycli.echo(str(e), err=True, fg='red')
            return

        try:
            if handle_clip_command(mycli, text):
                return
        except RuntimeError as e:
            mycli.logger.error('sql: %r, error: %r', text, e)
            mycli.logger.error('traceback: %r', traceback.format_exc())
            mycli.echo(str(e), err=True, fg='red')
            return

        while special.is_llm_command(text):
            start = time.time()
            try:
                assert sqlexecute.conn is not None
                cur = sqlexecute.conn.cursor()
                context, sql, duration = special.handle_llm(
                    text,
                    cur,
                    sqlexecute.dbname or '',
                    mycli.llm_prompt_field_truncate,
                    mycli.llm_prompt_section_truncate,
                )
                if context:
                    click.echo('LLM Response:')
                    click.echo(context)
                    click.echo('---')
                if special.is_timing_enabled():
                    mycli.output_timing(f'Time: {duration:0.03f}s')
                assert mycli.prompt_session is not None
                text = mycli.prompt_session.prompt(
                    default=sql or '',
                    inputhook=inputhook,
                    message=loaded_message_fn,
                )
            except KeyboardInterrupt:
                return
            except special.FinishIteration as e:
                if e.results:
                    _output_results(mycli, state, e.results, start)
                return
            except RuntimeError as e:
                mycli.logger.error('sql: %r, error: %r', text, e)
                mycli.logger.error('traceback: %r', traceback.format_exc())
                mycli.echo(str(e), err=True, fg='red')
                return

    text = text.strip()
    if not text:
        return

    original_text = text
    try:
        polars_pipeline = parse_polars_transform(text)
    except PolarsTransformError as exc:
        mycli.echo(str(exc), err=True, fg='red')
        return
    if polars_pipeline is not None:
        text = polars_pipeline.sql
    elif is_redirect_command(text):
        sql_part, command_part, file_operator_part, file_part = get_redirect_components(text)
        text = sql_part or ''
        try:
            special.set_redirect(command_part, file_operator_part, file_part)
        except (FileNotFoundError, OSError, RuntimeError) as e:
            mycli.logger.error('sql: %r, error: %r', text, e)
            mycli.logger.error('traceback: %r', traceback.format_exc())
            mycli.echo(str(e), err=True, fg='red')
            return

    if mycli.sandbox_mode and not is_sandbox_allowed(text):
        mycli.echo(
            "ERROR 1820: You must reset your password using ALTER USER or SET PASSWORD before executing this statement.",
            err=True,
            fg='red',
        )
        return

    if mycli.destructive_warning:
        destroy = confirm_destructive_query(mycli.destructive_keywords, text)
        if destroy is None:
            pass
        elif destroy is True:
            mycli.echo('Your call!')
        else:
            mycli.echo('Wise choice!')
            return

    dropping_active_database = is_dropping_database(text, sqlexecute.dbname)
    if dropping_active_database:
        mycli.completion_refresher.stop()

    successful = False
    try:
        mycli.logger.debug('sql: %r', text)
        polars_transform: PolarsTransform | None = (
            prepare_polars_transform(text, polars_pipeline.expression) if polars_pipeline is not None else None
        )
        query_history_text = original_text if polars_transform is not None else text
        special.write_tee(mycli.last_prompt_message, nl=False)
        special.write_tee(query_history_text)
        mycli.log_query(query_history_text)

        start = time.time()
        results = sqlexecute.run(text)
        mycli.main_formatter.query = text
        mycli.redirect_formatter.query = text
        mycli.explorer_formatter.query = text
        if polars_transform is not None:
            assert polars_pipeline is not None
            if polars_pipeline.output_path is None:
                polars_result = run_polars_transform(
                    polars_transform,
                    results,
                    image_protocol=mycli.image_protocol,
                    plot_scale_factor=mycli.plot_scale_factor,
                    plot_ppi=mycli.plot_ppi,
                    plot_theme=mycli.plot_theme,
                )
            else:
                polars_result = run_polars_transform(
                    polars_transform,
                    results,
                    polars_pipeline.output_path,
                    image_protocol=mycli.image_protocol,
                    plot_scale_factor=mycli.plot_scale_factor,
                    plot_ppi=mycli.plot_ppi,
                    plot_theme=mycli.plot_theme,
                )
            if polars_pipeline.output_path is None:
                if polars_pipeline.output_mode == 'explorer':
                    special.set_explorer_output(True)
                elif polars_pipeline.output_mode == 'expanded':
                    special.set_expanded_output(True)
            _output_results(mycli, state, iter([polars_result]), start)
            if polars_pipeline.output_path is not None:
                special.run_post_redirect_hook(
                    mycli.post_redirect_command,
                    polars_pipeline.output_path,
                )
        else:
            _output_results(mycli, state, results, start)
            special.unset_once_if_written(mycli.post_redirect_command)
            special.flush_pipe_once_if_written(mycli.post_redirect_command)
        successful = True
    except pymysql.err.InterfaceError:
        if not mycli.reconnect():
            return
        _one_iteration(mycli, state, text)
        return
    except EOFError as e:
        raise e
    except KeyboardInterrupt:
        connection_id_to_kill = sqlexecute.connection_id or 0
        if connection_id_to_kill > 0:
            mycli.logger.debug('connection id to kill: %r', connection_id_to_kill)
            try:
                sqlexecute.connect()
                for kill_result in sqlexecute.run(f'kill {connection_id_to_kill}'):
                    status_str = str(kill_result.status_plain).lower()
                    if status_str.find('ok') > -1:
                        mycli.logger.debug('cancelled query, connection id: %r, sql: %r', connection_id_to_kill, text)
                        mycli.echo(f'Cancelled query id: {connection_id_to_kill}', err=True, fg='blue')
                    else:
                        mycli.logger.debug(
                            'Failed to confirm query cancellation, connection id: %r, sql: %r',
                            connection_id_to_kill,
                            text,
                        )
                        mycli.echo(f'Failed to confirm query cancellation, id: {connection_id_to_kill}', err=True, fg='red')
            except Exception as e2:
                mycli.echo(f'Encountered error while cancelling query: {e2}', err=True, fg='red')
        else:
            mycli.logger.debug('Did not get a connection id, skip cancelling query')
            mycli.echo('Did not get a connection id, skip cancelling query', err=True, fg='red')
    except NotImplementedError:
        mycli.echo('Not Yet Implemented.', fg='yellow')
    except PolarsTransformError as exc:
        mycli.logger.error('sql: %r, error: %r', text, exc)
        mycli.echo(str(exc), err=True, fg='red')
    except pymysql.OperationalError as e1:
        mycli.logger.debug('Exception: %r', e1)
        if e1.args[0] == ER_MUST_CHANGE_PASSWORD:
            mycli.sandbox_mode = True
            mycli.echo(
                "ERROR 1820: You must reset your password using ALTER USER or SET PASSWORD before executing this statement.",
                err=True,
                fg='red',
            )
        elif e1.args[0] in (2003, 2006, 2013):
            if not mycli.reconnect():
                return
            _one_iteration(mycli, state, text)
            return
        else:
            mycli.logger.error('sql: %r, error: %r', text, e1)
            mycli.logger.error('traceback: %r', traceback.format_exc())
            mycli.echo(str(e1), err=True, fg='red')
    except Exception as e:
        mycli.logger.error('sql: %r, error: %r', text, e)
        mycli.logger.error('traceback: %r', traceback.format_exc())
        mycli.echo(str(e), err=True, fg='red')
    else:
        if mycli.sandbox_mode and is_password_change(text):
            new_password = extract_new_password(text)
            if new_password is not None:
                sqlexecute.password = new_password
            try:
                sqlexecute.connect()
                mycli.sandbox_mode = False
                mycli.echo("Password changed successfully. Reconnected.", err=True, fg='green')
                mycli.refresh_completions()
            except Exception as e:
                mycli.sandbox_mode = False
                mycli.echo(
                    f"Password changed but reconnection failed: {e}\nPlease restart mycli with your new password.",
                    err=True,
                    fg='yellow',
                )

        if dropping_active_database:
            sqlexecute.dbname = None
            sqlexecute.connect()

        if need_completion_refresh(text):
            mycli.refresh_completions(reset=dropping_active_database or need_completion_reset(text))
    finally:
        if dropping_active_database and not successful:
            mycli.refresh_completions()
        if mycli.logfile is False:
            mycli.echo('Warning: This query was not logged.', err=True, fg='red')

    query = Query(original_text if polars_pipeline is not None else text, successful, state.mutating)
    mycli.query_history.append(query)


def _contributors_picker() -> str:
    lines: str = ""

    try:
        with resources.files(mycli_package).joinpath("AUTHORS").open('r') as f:
            lines += f.read()
    except FileNotFoundError:
        pass

    contents = []
    for line in lines.split("\n"):
        if m := re.match(r"^ *\* (.*)", line):
            contents.append(m.group(1))
    return random.choice(contents) if contents else 'our contributors'


def _sponsors_picker() -> str:
    lines: str = ""

    try:
        with resources.files(mycli_package).joinpath("SPONSORS").open('r') as f:
            lines += f.read()
    except FileNotFoundError:
        pass

    contents = []
    for line in lines.split("\n"):
        if m := re.match(r"^ *\* (.*)", line):
            contents.append(m.group(1))
    return random.choice(contents) if contents else 'our sponsors'


def _tips_picker() -> str:
    tips = []

    try:
        with resources.files(mycli_package).joinpath('TIPS').open('r') as f:
            for line in f:
                if line.startswith("#"):
                    continue
                if tip := line.strip():
                    tips.append(tip)
    except FileNotFoundError:
        pass

    return random.choice(tips) if tips else r'\? or "help" for help!'


def _configure_editor(mycli: 'MyCli') -> None:
    if configured_editor := mycli.config['editor'].get('editor_command'):
        os.environ['VISUAL'] = configured_editor
    elif not os.environ.get('VISUAL') and not os.environ.get('EDITOR'):
        if shutil.which('code'):
            os.environ['VISUAL'] = 'code --wait'
        elif WIN:
            os.environ['VISUAL'] = 'edit'
        else:
            os.environ['VISUAL'] = 'vi'


def main_repl(mycli: 'MyCli') -> None:
    sqlexecute = mycli.sqlexecute
    assert sqlexecute is not None
    state = ReplState()

    mycli.configure_pager()
    _configure_editor(mycli)
    history = _create_history(mycli)
    if history is not None:
        mycli.completer.frecency_provider = lambda: history.frecency
    if mycli.smart_completion and not mycli.sandbox_mode:
        mycli.refresh_completions()

    key_bindings = mycli_bindings(mycli)
    _show_startup_banner(mycli, sqlexecute)
    _build_prompt_session(mycli, state, history, key_bindings)
    set_all_external_titles(mycli)

    try:
        while True:
            _one_iteration(mycli, state)
            state.iterations += 1
    except EOFError:
        special.close_tee()
        if mycli.verbosity >= 0:
            mycli.echo('Goodbye!')
