from __future__ import annotations

import re
from collections.abc import Callable
from re import Pattern

import gdb

import pwndbg
from pwndbg.color import message
from pwndbg.commands.context import context
from pwndbg.commands.context import context_sections
from pwndbg.commands.context import contextoutput
from pwndbg.commands.context import resetcontextoutput
from pwndbg.gdblib import gdb_version

tui_adjust_height = pwndbg.config.add_param(
    "context-tui-adjust-height",
    True,
    "adjust height of context sections to fit the TUI",
    help_docstring="""
Adjust the height of context sections to fit the TUI window, ignoring section specific limits.
This uses all vertical space available in the TUI windows if the context section supports it
like the "disasm" or "stack" windows.
""",
)


class ContextTUIWindow:
    _tui_window: gdb.TuiWindow
    _section: str
    _lines: list[str]
    _blank_line_lengths: list[int]
    _longest_line: int
    _before_prompt_listener: Callable[[None], object]
    _vscroll_start: int
    _hscroll_start: int
    _old_width: int
    _old_height: int
    _ansi_escape_regex: Pattern[str]
    _enabled: bool

    _static_enabled: bool = True
    _context_windows: list[ContextTUIWindow] = []

    def __init__(self, tui_window: gdb.TuiWindow, section: str) -> None:
        self._tui_window = tui_window
        self._section = section
        self._tui_window.title = section
        self._lines = []
        self._blank_line_lengths = []
        self._longest_line = 0
        self._before_prompt_listener = self._before_prompt
        self._old_width = 0
        self._old_height = 0
        self._vscroll_start = 0
        self._hscroll_start = 0
        self._ansi_escape_regex = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
        self._enabled = False
        self._enable()
        gdb.events.before_prompt.connect(self._before_prompt_listener)
        ContextTUIWindow._context_windows.append(self)

    def close(self) -> None:
        ContextTUIWindow._context_windows.remove(self)
        if self._enabled:
            self._disable()
        gdb.events.before_prompt.disconnect(self._before_prompt_listener)

    def render(self) -> None:
        # render is called again after the TUI was disabled
        self._verify_enabled_state()

        if (
            not self._lines
            and self._section != "legend"
            and self._section not in str(pwndbg.config.context_sections)
        ):
            self._tui_window.write(
                message.warn(
                    f"Section '{self._section}' is not in 'context-sections' and won't be updated automatically."
                ),
                True,
            )
            return

        height = self._tui_window.height
        width = self._tui_window.width
        start = self._vscroll_start
        end = start + height
        lines = self._lines[start:end]
        output = ""
        for idx, line in enumerate(lines):
            if self._blank_line_lengths[start + idx] < width:
                line = self._ansi_substr(line, self._hscroll_start, len(line)) + "\n"
            else:
                line = self._ansi_substr(line, self._hscroll_start, self._hscroll_start + width)
                if self._blank_line_lengths[start + idx] - self._hscroll_start < width:
                    line += "\n"
            output += line
        self._tui_window.write(output, True)

    def hscroll(self, num: int) -> None:
        old_start = self._hscroll_start
        max_start = max(0, self._longest_line - self._tui_window.width)
        self._hscroll_start = min(max(0, self._hscroll_start + num), max_start)
        if old_start != self._hscroll_start:
            self.render()

    def vscroll(self, num: int) -> None:
        old_start = self._vscroll_start
        max_start = max(0, len(self._lines) - self._tui_window.height)
        self._vscroll_start = min(max(0, self._vscroll_start + num), max_start)
        if old_start != self._vscroll_start:
            self.render()

    def click(self, x: int, y: int, button: int) -> None:
        gdb.execute(f"focus pwndbg_{self._section}", to_string=True)

    def _before_prompt(self):
        if not self._verify_enabled_state():
            return
        self._update()
        self.render()

    def _enable(self):
        self._update()
        self._enabled = True

    def _disable(self):
        self._old_width = 0
        self._old_height = 0
        resetcontextoutput(self._section)
        self._enabled = False

    def _update(self):
        if self._old_width != self._tui_window.width or self._old_height != self._tui_window.height:
            self._old_width = self._tui_window.width
            self._old_height = self._tui_window.height
            contextoutput(
                self._section,
                self._receive_context_output,
                clearing=True,
                banner="none",
                width=self._old_width - 1,
                height=self._old_height if tui_adjust_height else None,
            )

    def _receive_context_output(self, data: str):
        if not self._verify_enabled_state():
            return
        self._lines = data.strip("\n").split("\n")
        self._blank_line_lengths = [
            len(self._ansi_escape_regex.sub("", line)) for line in self._lines
        ]
        self._longest_line = max(self._blank_line_lengths)
        self.render()

    def _verify_enabled_state(self) -> bool:
        is_valid = self._tui_window.is_valid()
        if is_valid:
            if not self._enabled:
                for context_window in ContextTUIWindow._context_windows:
                    context_window._enable()
                if not ContextTUIWindow._static_enabled and pwndbg.dbg.selected_inferior().alive():
                    context()
                ContextTUIWindow._static_enabled = True
        else:
            if self._enabled:
                self._disable()
                ContextTUIWindow._static_enabled = False
        return is_valid

    def _ansi_substr(self, line: str, start_char: int, end_char: int) -> str:
        ansi_escape_before_start = ""
        ansi_escape_after_end = ""
        colored_start_idx = 0
        colored_end_idx = 0
        colored_idx = 0
        char_count = 0
        while colored_idx < len(line):
            c = line[colored_idx]
            # collect all ansi escape sequences before the start of the colored substring
            # as well as after the end of the colored substring
            # skip them while counting the characters to slice
            if c == "\x1b":
                m = self._ansi_escape_regex.match(line[colored_idx:])
                if m:
                    colored_idx += m.end()
                    if char_count < start_char:
                        ansi_escape_before_start += m.group()
                        colored_start_idx += m.end()
                    if char_count < end_char:
                        colored_end_idx += m.end()
                    if colored_idx > colored_end_idx:
                        ansi_escape_after_end += m.group()
                    continue

            if char_count < start_char:
                colored_start_idx += 1
            if char_count < end_char:
                colored_end_idx += 1
                char_count += 1
            colored_idx += 1
        return (
            ansi_escape_before_start
            + line[colored_start_idx:colored_end_idx]
            + ansi_escape_after_end
        )

    # Workaround for display bug in GDB TUI ANSI escape sequence translation.
    # Only resetting the foreground or background colors results in
    # generating wrong colors in the TUI.
    # The workaround is to avoid the sequences that only reset the colors
    # and replace them with a full reset sequence.
    # It resets other styling attributes like bold too but it's better
    # than having wrong colors.
    # https://github.com/pwndbg/pwndbg/issues/2654
    if gdb_version < (16, 3):
        ___ansi_substr = _ansi_substr
        _ansi_substr = (
            lambda *a, **kw: ContextTUIWindow.___ansi_substr(*a, **kw)
            .replace("\x1b[39m", "\x1b[0m")
            .replace("\x1b[49m", "\x1b[0m")
            .replace("\x1b[39;49m", "\x1b[0m")
        )


if hasattr(gdb, "register_window_type"):
    sections = ["legend"] + [
        section.__name__.replace("context_", "") for section in context_sections.values()
    ]
    for section_name in sections:
        # https://github.com/python/mypy/issues/12557
        target_func: Callable[..., gdb._Window] = (
            lambda window, section_name=section_name: ContextTUIWindow(window, section_name)
        )
        gdb.register_window_type(
            "pwndbg_" + section_name,
            target_func,
        )
