import re
import sys

import pynvml
from pygments import highlight
from pygments.formatters import HtmlFormatter
from pygments.lexers import TextLexer, get_lexer_by_name, guess_lexer
from pygments.util import ClassNotFound
from PyQt6.QtCore import Qt, QThread, QTimer, QUrl, pyqtSignal
from PyQt6.QtGui import (
    QColor,
    QDesktopServices,
    QFont,
    QIcon,
    QLinearGradient,
    QPainter,
    QPainterPath,
    QPalette,
    QPixmap,
    QTextBlockFormat,
    QTextCharFormat,
    QTextCursor,
    QTextDocument,
    QTextDocumentFragment,
    QTextLength,
    QTextTableCellFormat,
    QTextTableFormat,
)
from PyQt6.QtWidgets import (
    QApplication,
    QHBoxLayout,
    QLabel,
    QLineEdit,
    QListWidget,
    QListWidgetItem,
    QMainWindow,
    QMenu,
    QMessageBox,
    QPushButton,
    QTextBrowser,
    QVBoxLayout,
    QWidget,
)

import chatbot

# How often the bottom-right GPU/VRAM readout refreshes.
GPU_STATUS_INTERVAL_MS = 1000

# --- Visual theme ------------------------------------------------------
# An explicit dark palette rather than relying on the OS theme, so the app
# looks the same everywhere instead of "Qt follows Windows' dark app theme".
# Brand palette: jet black / blue slate / silver / white / coral glow.
JET_BLACK = "#2d3142"
BLUE_SLATE = "#4f5d75"
SILVER = "#bfc0c0"
WHITE = "#ffffff"
CORAL_GLOW = "#ef8354"

BG_APP = "#242735"  # jet black, darkened further for the outer chrome
BG_PANEL = JET_BLACK
BG_ELEVATED = BLUE_SLATE
BORDER = "rgba(191, 192, 192, 0.35)"  # silver, translucent
BORDER_STRONG = "rgba(191, 192, 192, 0.55)"
TEXT_PRIMARY = WHITE
TEXT_SECONDARY = SILVER
ACCENT = CORAL_GLOW
ACCENT_HOVER = "#ff9768"
ACCENT_PRESSED = "#cb6f47"

# Green -> yellow -> red ramp used to color the GPU/VRAM load readout.
LOAD_LOW = (0x22, 0xc5, 0x5e)
LOAD_MID = (0xea, 0xb3, 0x08)
LOAD_HIGH = (0xef, 0x44, 0x44)


def _load_color(pct: float) -> str:
    """Hex color for a 0-100 load percentage, interpolated low->mid->high."""
    pct = max(0.0, min(100.0, pct))
    if pct <= 50:
        c1, c2, t = LOAD_LOW, LOAD_MID, pct / 50.0
    else:
        c1, c2, t = LOAD_MID, LOAD_HIGH, (pct - 50) / 50.0
    r = round(c1[0] + (c2[0] - c1[0]) * t)
    g = round(c1[1] + (c2[1] - c1[1]) * t)
    b = round(c1[2] + (c2[2] - c1[2]) * t)
    return f"#{r:02x}{g:02x}{b:02x}"

# (accent, background) per chat role, used for the accent stripe and tinted
# card background of each message.
ROLE_STYLES = {
    "user": (CORAL_GLOW, "#3a2c26"),
    "assistant": (SILVER, "#2f3342"),
    "system": (BLUE_SLATE, "#262a36"),
    "thinking": (TEXT_SECONDARY, "#282b38"),
}

# Prefix for the fake anchor href used to make a thinking block's header
# clickable (e.g. "think:think-3"), intercepted in _on_chat_anchor_clicked
# rather than let QTextBrowser try to navigate to it.
_THINK_ANCHOR_PREFIX = "think:"

_STYLESHEET = f"""
QMainWindow {{
    background-color: {BG_APP};
}}

QWidget {{
    color: {TEXT_PRIMARY};
    font-family: "Segoe UI Variable Text", "Segoe UI", sans-serif;
    font-size: 10.5pt;
}}

#titleLabel {{
    font-size: 15pt;
    font-weight: 600;
    color: {TEXT_PRIMARY};
}}

#sidebar {{
    background-color: {BG_APP};
    border-right: 1px solid {BORDER};
}}

QPushButton#newChatButton {{
    background-color: {BG_ELEVATED};
    border: 1px solid {BORDER};
    padding: 8px 12px;
    color: {TEXT_PRIMARY};
    font-weight: 600;
    text-align: left;
}}

QPushButton#newChatButton:hover {{
    border: 1px solid {ACCENT};
}}

QListWidget#sessionList {{
    background-color: transparent;
    border: none;
    outline: none;
}}

QListWidget#sessionList::item {{
    padding: 8px 10px;
    color: {TEXT_SECONDARY};
}}

QListWidget#sessionList::item:selected {{
    background-color: {BG_ELEVATED};
    color: {TEXT_PRIMARY};
}}

QListWidget#sessionList::item:hover {{
    background-color: rgba(191, 192, 192, 0.08);
}}

#badge {{
    background-color: {BG_ELEVATED};
    border: 1px solid {BORDER};
    padding: 4px 12px;
    color: {TEXT_SECONDARY};
    font-size: 9pt;
}}

QTextBrowser#chatView {{
    background-color: {BG_PANEL};
    border: 1px solid {BORDER};
    padding: 10px;
}}

QLineEdit#messageInput {{
    background-color: {BG_ELEVATED};
    border: 1px solid {BORDER};
    padding: 0 16px;
    color: {TEXT_PRIMARY};
    selection-background-color: {ACCENT};
}}

QLineEdit#messageInput:focus {{
    border: 1px solid {ACCENT};
}}

QLineEdit#messageInput:disabled {{
    color: {TEXT_SECONDARY};
}}

QPushButton#sendButton {{
    background-color: {ACCENT};
    color: {WHITE};
    border: none;
    padding: 0 22px;
    font-weight: 600;
}}

QPushButton#sendButton:hover {{
    background-color: {ACCENT_HOVER};
}}

QPushButton#sendButton:pressed {{
    background-color: {ACCENT_PRESSED};
}}

QPushButton#sendButton:disabled {{
    background-color: {BG_ELEVATED};
    color: {TEXT_SECONDARY};
    border: 1px solid {BORDER};
}}

QStatusBar {{
    background-color: {BG_APP};
    border-top: 1px solid {BORDER};
    color: {TEXT_SECONDARY};
    font-size: 9pt;
}}

QStatusBar::item {{
    border: none;
}}

QScrollBar:vertical {{
    background: transparent;
    width: 10px;
    margin: 4px 2px 4px 0;
}}

QScrollBar::handle:vertical {{
    background: {BORDER_STRONG};
    min-height: 24px;
}}

QScrollBar::handle:vertical:hover {{
    background: {TEXT_SECONDARY};
}}

QScrollBar::add-line:vertical, QScrollBar::sub-line:vertical {{
    height: 0;
}}

QScrollBar::add-page:vertical, QScrollBar::sub-page:vertical {{
    background: transparent;
}}

QToolTip {{
    background-color: {BG_ELEVATED};
    color: {TEXT_PRIMARY};
    border: 1px solid {BORDER};
    padding: 4px 8px;
}}
"""


def _dark_palette() -> QPalette:
    """Baseline QPalette for the Fusion style, so anything the stylesheet
    doesn't reach (native dialogs, menus) still matches the dark theme."""
    palette = QPalette()
    palette.setColor(QPalette.ColorRole.Window, QColor(BG_APP))
    palette.setColor(QPalette.ColorRole.WindowText, QColor(TEXT_PRIMARY))
    palette.setColor(QPalette.ColorRole.Base, QColor(BG_ELEVATED))
    palette.setColor(QPalette.ColorRole.AlternateBase, QColor(BG_PANEL))
    palette.setColor(QPalette.ColorRole.ToolTipBase, QColor(BG_ELEVATED))
    palette.setColor(QPalette.ColorRole.ToolTipText, QColor(TEXT_PRIMARY))
    palette.setColor(QPalette.ColorRole.Text, QColor(TEXT_PRIMARY))
    palette.setColor(QPalette.ColorRole.PlaceholderText, QColor(TEXT_SECONDARY))
    palette.setColor(QPalette.ColorRole.Button, QColor(BG_ELEVATED))
    palette.setColor(QPalette.ColorRole.ButtonText, QColor(TEXT_PRIMARY))
    palette.setColor(QPalette.ColorRole.Highlight, QColor(ACCENT))
    palette.setColor(QPalette.ColorRole.HighlightedText, QColor(WHITE))
    return palette


def _make_app_icon() -> QIcon:
    """Small generated glyph for the window/taskbar icon, since the project
    ships no icon asset."""
    size = 64
    pixmap = QPixmap(size, size)
    pixmap.fill(QColor(0, 0, 0, 0))

    painter = QPainter(pixmap)
    painter.setRenderHint(QPainter.RenderHint.Antialiasing)

    gradient = QLinearGradient(0, 0, size, size)
    gradient.setColorAt(0.0, QColor(JET_BLACK))
    gradient.setColorAt(1.0, QColor(CORAL_GLOW))

    path = QPainterPath()
    path.addRect(2, 2, size - 4, size - 4)
    painter.fillPath(path, gradient)

    painter.setPen(QColor(WHITE))
    painter.setFont(QFont("Segoe UI", 28, QFont.Weight.Bold))
    painter.drawText(pixmap.rect(), Qt.AlignmentFlag.AlignCenter, "D")
    painter.end()

    return QIcon(pixmap)

# GitHub dialect adds fenced code blocks with language tags, tables, and
# strikethrough on top of CommonMark; NoHTML stops raw HTML in the model's
# output from being interpreted as literal markup in the chat view.
_MARKDOWN_FEATURES = (
    QTextDocument.MarkdownFeature.MarkdownDialectGitHub
    | QTextDocument.MarkdownFeature.MarkdownNoHTML
)

# Matches fenced code blocks (```lang\ncode\n```) so they can be pulled out
# and rendered through Pygments instead of QTextDocumentFragment.fromMarkdown,
# which only gives code a monospace font and no syntax colors.
_CODE_FENCE_RE = re.compile(r"```([\w+-]*)\n(.*?)```", re.DOTALL)

# noclasses=True inlines every color as a style="..." attribute rather than
# relying on a <style> stylesheet, which QTextBrowser's limited HTML/CSS
# support won't pick up. nowrap=True drops the surrounding <div>/<pre>, since
# a QTextCursor only applies that container's block styling (background,
# monospace font) to whichever line happens to share the cursor's existing
# block - every other line silently falls back to plain text. Instead each
# source line is inserted as its own block, below, with that styling applied
# explicitly and identically to all of them.
#
# "monokai" (dark) rather than a light Pygments style: the chat view itself
# renders with a dark palette (Qt follows Windows' dark app theme here), so a
# light code background would show as a glaring white box, and - worse -
# Pygments only emits a color for tokens it specifically highlights, leaving
# plain identifiers/punctuation uncolored to inherit whatever text color is
# ambient; against this app's dark theme that meant near-invisible white-on-
# near-white text. _CODE_BG/_CODE_FG give the wrapping <span> in
# _highlight_code_lines() an explicit default color so uncolored tokens
# inherit a color that's actually legible on a dark block, via normal CSS
# inheritance, while Pygments' own colored spans still override it.
_CODE_FORMATTER = HtmlFormatter(style="monokai", noclasses=True, nowrap=True)
_CODE_BG = "#272822"
_CODE_FG = "#f8f8f2"

_CODE_BLOCK_FORMAT = QTextBlockFormat()
_CODE_BLOCK_FORMAT.setBackground(QColor(_CODE_BG))
_CODE_BLOCK_FORMAT.setLeftMargin(10)
# Each source line is its own QTextBlock (see _append_markdown_block); Qt's
# block background doesn't paint into top/bottom margins, so any non-zero
# margin here left a sliver of the message card's own background showing
# between every line, reading as horizontal gaps through the code block.

_CODE_CHAR_FORMAT = QTextCharFormat()
_CODE_CHAR_FORMAT.setFontFamilies(["Consolas", "Courier New"])


def _highlight_code_lines(code: str, lang: str) -> list[str]:
    lexer = None
    if lang:
        try:
            lexer = get_lexer_by_name(lang)
        except ClassNotFound:
            lexer = None
    if lexer is None:
        try:
            lexer = guess_lexer(code)
        except ClassNotFound:
            lexer = TextLexer()
    highlighted = highlight(code, lexer, _CODE_FORMATTER)
    lines = highlighted.rstrip("\n").split("\n")
    return [f'<span style="color:{_CODE_FG}; white-space:pre;">{line}</span>' for line in lines]


def _reset_cursor_format(cursor: QTextCursor) -> None:
    """Start a new block with plain formatting, so a preceding code block's
    background/monospace font can't bleed into whatever is inserted next -
    insertBlock() otherwise inherits the current char format."""
    cursor.insertBlock(QTextBlockFormat(), QTextCharFormat())


def _insert_message_card(cursor: QTextCursor, accent: str, bg: str) -> QTextCursor:
    """Insert a flat card with a colored accent stripe for one chat message,
    and return a cursor positioned inside it, ready for content.

    QTextBrowser's HTML subset has no border-radius, so a rounded chat bubble
    isn't achievable directly; a two-column table (a thin accent-colored
    strip beside a tinted message body) stands in for one. Each message gets
    its own table/cell, so - unlike the previous shared-cursor approach -
    there's no risk of one message's formatting bleeding into the next.
    """
    table_format = QTextTableFormat()
    table_format.setCellPadding(0)
    table_format.setCellSpacing(0)
    table_format.setBorder(0)
    table_format.setWidth(QTextLength(QTextLength.Type.PercentageLength, 100))
    table_format.setColumnWidthConstraints(
        [
            QTextLength(QTextLength.Type.FixedLength, 3),
            QTextLength(QTextLength.Type.PercentageLength, 100),
        ]
    )
    table_format.setTopMargin(2)
    table_format.setBottomMargin(10)

    table = cursor.insertTable(1, 2, table_format)

    strip_format = QTextTableCellFormat()
    strip_format.setBackground(QColor(accent))
    table.cellAt(0, 0).setFormat(strip_format)

    body_format = QTextTableCellFormat()
    body_format.setBackground(QColor(bg))
    body_format.setTopPadding(10)
    body_format.setBottomPadding(10)
    body_format.setLeftPadding(14)
    body_format.setRightPadding(14)
    body_cell = table.cellAt(0, 1)
    body_cell.setFormat(body_format)

    return body_cell.firstCursorPosition()


class ModelDownloadThread(QThread):
    # qint64, not int - int maps to a 32-bit C int for this signal's
    # cross-thread queued connection, which silently wraps a ~15GB byte
    # count (e.g. 14_977_484_704 truncates to 2_092_582_816).
    progress = pyqtSignal("qint64", "qint64")
    done = pyqtSignal()
    failed = pyqtSignal(str)

    def run(self) -> None:
        try:
            chatbot.download_model(on_progress=self.progress.emit)
        except Exception as e:
            self.failed.emit(str(e))
            return
        self.done.emit()


class ModelLoaderThread(QThread):
    loaded = pyqtSignal(object)
    failed = pyqtSignal(str)

    def run(self) -> None:
        try:
            llm = chatbot.load_model()
        except Exception as e:
            self.failed.emit(str(e))
            return
        self.loaded.emit(llm)


class AskThread(QThread):
    tool_call = pyqtSignal(str)
    usage = pyqtSignal(dict)
    token = pyqtSignal(str)
    thinking_token = pyqtSignal(str)
    finished_ok = pyqtSignal(str)
    failed = pyqtSignal(str)

    def __init__(self, llm, messages, parent=None):
        super().__init__(parent)
        self._llm = llm
        self._messages = messages

    def run(self) -> None:
        try:
            answer = chatbot.ask(
                self._llm,
                self._messages,
                on_tool_call=self.tool_call.emit,
                on_usage=self.usage.emit,
                on_token=self.token.emit,
                on_thinking_token=self.thinking_token.emit,
            )
        except Exception as e:
            self.failed.emit(str(e))
            return
        self.finished_ok.emit(answer)


class ChatWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Dev Assist Chatbot")
        self.resize(900, 700)

        self._llm = None
        self._sessions: list[chatbot.ChatSession] = chatbot.list_sessions()
        self._current_session: chatbot.ChatSession = (
            self._sessions[0] if self._sessions else chatbot.new_session()
        )
        self._history: list[tuple[str, str]] = list(self._current_session.history)
        self._pending_question = ""
        self._ask_thread: AskThread | None = None
        # Position in chat_view's document where the in-progress streamed
        # answer card starts, and a cursor into its body cell for appending
        # incoming tokens; both None when no answer is currently streaming.
        self._stream_start_pos: int | None = None
        self._stream_body_cursor: QTextCursor | None = None
        # Whether a real (non-newline) answer character has been inserted
        # yet for the in-progress card - the model's text after </think>
        # typically opens with a blank line or two, which _on_token trims
        # rather than showing as a gap at the top of the answer bubble.
        self._answer_started = False

        # Thinking blocks (one per model reasoning round), keyed by an id
        # referenced from the clickable "think:<id>" anchor in their header.
        # Each holds the full reasoning text, whether it's currently shown
        # expanded, a fixed cursor marking where its content region starts,
        # and a live cursor tracking where that region currently ends - see
        # _render_thinking(). self._stream_thinking_id names whichever block
        # (if any) is still actively streaming in the current round.
        self._thinking_blocks: dict[str, dict] = {}
        self._thinking_seq = 0
        self._stream_thinking_id: str | None = None

        central = QWidget()
        self.setCentralWidget(central)
        root_layout = QHBoxLayout(central)
        root_layout.setContentsMargins(0, 0, 0, 0)
        root_layout.setSpacing(0)

        sidebar = QWidget()
        sidebar.setObjectName("sidebar")
        sidebar.setFixedWidth(220)
        sidebar_layout = QVBoxLayout(sidebar)
        sidebar_layout.setContentsMargins(12, 16, 12, 16)
        sidebar_layout.setSpacing(10)

        self.new_chat_button = QPushButton("+ New Chat")
        self.new_chat_button.setObjectName("newChatButton")
        self.new_chat_button.setMinimumHeight(36)
        self.new_chat_button.setCursor(Qt.CursorShape.PointingHandCursor)
        self.new_chat_button.clicked.connect(self._on_new_chat)
        sidebar_layout.addWidget(self.new_chat_button)

        self.session_list = QListWidget()
        self.session_list.setObjectName("sessionList")
        self.session_list.itemClicked.connect(self._on_session_clicked)
        self.session_list.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self.session_list.customContextMenuRequested.connect(self._on_session_context_menu)
        sidebar_layout.addWidget(self.session_list, 1)

        root_layout.addWidget(sidebar)

        content = QWidget()
        layout = QVBoxLayout(content)
        layout.setContentsMargins(20, 16, 20, 16)
        layout.setSpacing(14)
        root_layout.addWidget(content, 1)

        header = QWidget()
        header.setObjectName("header")
        header_layout = QHBoxLayout(header)
        header_layout.setContentsMargins(4, 0, 4, 12)
        header_layout.setSpacing(12)

        title_box = QVBoxLayout()
        title_box.setSpacing(0)
        title_label = QLabel("Dev Assist")
        title_label.setObjectName("titleLabel")
        title_box.addWidget(title_label)
        header_layout.addLayout(title_box)
        header_layout.addStretch(1)

        self._n_ctx = 0
        self.context_status_label = QLabel("Context: --")
        self.context_status_label.setObjectName("badge")
        header_layout.addWidget(self.context_status_label)

        self.gpu_status_label = QLabel("GPU: --")
        self.gpu_status_label.setObjectName("badge")
        header_layout.addWidget(self.gpu_status_label)

        layout.addWidget(header)

        self.chat_view = QTextBrowser()
        self.chat_view.setObjectName("chatView")
        # Links are handled by hand in _on_chat_anchor_clicked (real URLs get
        # opened externally there too) rather than left to QTextBrowser,
        # since its own link handling doesn't know about the fake
        # "think:<id>" anchors used to toggle a thinking block.
        self.chat_view.setOpenLinks(False)
        self.chat_view.anchorClicked.connect(self._on_chat_anchor_clicked)
        self.chat_view.setFont(QFont("Segoe UI Variable Text", 10))
        layout.addWidget(self.chat_view)

        input_row = QHBoxLayout()
        input_row.setSpacing(10)

        self.input_box = QLineEdit()
        self.input_box.setObjectName("messageInput")
        self.input_box.setPlaceholderText("Ask a development question...")
        self.input_box.setMinimumHeight(40)
        input_palette = self.input_box.palette()
        input_palette.setColor(QPalette.ColorRole.PlaceholderText, QColor(TEXT_SECONDARY))
        self.input_box.setPalette(input_palette)
        self.input_box.returnPressed.connect(self._on_send)
        input_row.addWidget(self.input_box)

        self.send_button = QPushButton("Send  →")
        self.send_button.setObjectName("sendButton")
        self.send_button.setMinimumHeight(40)
        self.send_button.setMinimumWidth(96)
        self.send_button.setCursor(Qt.CursorShape.PointingHandCursor)
        self.send_button.clicked.connect(self._on_send)
        input_row.addWidget(self.send_button)

        layout.addLayout(input_row)

        self._refresh_session_list()
        if self._current_session.history:
            self._render_session_history(self._current_session)

        self._init_gpu_monitor()

        self._set_input_enabled(False)
        self._download_thread: ModelDownloadThread | None = None
        self._loader_thread: ModelLoaderThread | None = None

        if chatbot.MODEL_PATH.exists():
            self._start_model_load()
        else:
            self._start_model_download()

    def _start_model_download(self) -> None:
        self.statusBar().showMessage("Downloading model (first run only)... 0%")
        self._append_system(
            "Downloading the language model for first-time setup. This is a "
            "one-time, multi-gigabyte download and may take a while."
        )
        self._download_thread = ModelDownloadThread()
        self._download_thread.progress.connect(self._on_download_progress)
        self._download_thread.done.connect(self._on_download_done)
        self._download_thread.failed.connect(self._on_model_failed)
        self._download_thread.start()

    def _on_download_progress(self, downloaded: int, total: int) -> None:
        downloaded_gb = downloaded / (1024**3)
        if total:
            pct = 100 * downloaded / total
            total_gb = total / (1024**3)
            self.statusBar().showMessage(
                f"Downloading model... {downloaded_gb:.2f}/{total_gb:.2f} GB ({pct:.0f}%)"
            )
        else:
            self.statusBar().showMessage(f"Downloading model... {downloaded_gb:.2f} GB")

    def _on_download_done(self) -> None:
        self._download_thread = None
        self._start_model_load()

    def _start_model_load(self) -> None:
        self.statusBar().showMessage("Loading model...")
        self._loader_thread = ModelLoaderThread()
        self._loader_thread.loaded.connect(self._on_model_loaded)
        self._loader_thread.failed.connect(self._on_model_failed)
        self._loader_thread.start()

    def _init_gpu_monitor(self) -> None:
        # llama_cpp's n_gpu_layers=-1 offloads to the first CUDA device, so
        # that's the one whose load/VRAM is worth showing here.
        try:
            pynvml.nvmlInit()
            self._gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(0)
        except pynvml.NVMLError as e:
            self._gpu_handle = None
            self.gpu_status_label.setText("GPU: unavailable")
            self.gpu_status_label.setToolTip(str(e))
            return

        self._gpu_timer = QTimer(self)
        self._gpu_timer.timeout.connect(self._update_gpu_status)
        self._gpu_timer.start(GPU_STATUS_INTERVAL_MS)
        self._update_gpu_status()

    def _update_gpu_status(self) -> None:
        try:
            util = pynvml.nvmlDeviceGetUtilizationRates(self._gpu_handle)
            mem = pynvml.nvmlDeviceGetMemoryInfo(self._gpu_handle)
        except pynvml.NVMLError as e:
            self._gpu_timer.stop()
            self.gpu_status_label.setText("GPU: unavailable")
            self.gpu_status_label.setToolTip(str(e))
            return

        used_gb = mem.used / (1024**3)
        total_gb = mem.total / (1024**3)
        vram_pct = (mem.used / mem.total * 100) if mem.total else 0.0
        gpu_color = _load_color(util.gpu)
        vram_color = _load_color(vram_pct)
        self.gpu_status_label.setText(
            f'GPU <span style="color:{gpu_color};">{util.gpu}%</span>  ·  '
            f'VRAM <span style="color:{vram_color};">{used_gb:.1f}/{total_gb:.1f} GB</span>'
        )

    def _set_input_enabled(self, enabled: bool) -> None:
        self.input_box.setEnabled(enabled)
        self.send_button.setEnabled(enabled)
        # Switching/creating/deleting a chat while a generation is streaming
        # into chat_view would race the in-progress cursors, so the sidebar
        # is only usable while idle, same as the input box.
        self.new_chat_button.setEnabled(enabled)
        self.session_list.setEnabled(enabled)
        if enabled:
            self.input_box.setFocus()

    def _refresh_session_list(self) -> None:
        self.session_list.blockSignals(True)
        self.session_list.clear()
        for session in self._sessions:
            item = QListWidgetItem(session.title)
            item.setData(Qt.ItemDataRole.UserRole, session.id)
            self.session_list.addItem(item)
            if session.id == self._current_session.id:
                self.session_list.setCurrentItem(item)
        self.session_list.blockSignals(False)

    def _render_session_history(self, session: "chatbot.ChatSession") -> None:
        self.chat_view.clear()
        for question, answer in session.history:
            self._append_user(question)
            self._append_assistant(answer)

    def _on_new_chat(self) -> None:
        if self._ask_thread is not None:
            return
        if not self._current_session.history:
            return  # already on an empty, unsaved chat
        self._current_session = chatbot.new_session()
        self._history = []
        self.chat_view.clear()
        self._append_system("Dev assist chatbot ready. Ask a development question.")
        self.session_list.clearSelection()
        self.statusBar().showMessage("Ready.")

    def _on_session_clicked(self, item: QListWidgetItem) -> None:
        if self._ask_thread is not None:
            return
        session_id = item.data(Qt.ItemDataRole.UserRole)
        if session_id == self._current_session.id:
            return
        session = next((s for s in self._sessions if s.id == session_id), None)
        if session is None:
            return
        self._current_session = session
        self._history = list(session.history)
        self._render_session_history(session)
        self.statusBar().showMessage("Ready.")

    def _on_session_context_menu(self, pos) -> None:
        item = self.session_list.itemAt(pos)
        if item is None or self._ask_thread is not None:
            return
        session_id = item.data(Qt.ItemDataRole.UserRole)
        menu = QMenu(self)
        delete_action = menu.addAction("Delete chat")
        action = menu.exec(self.session_list.mapToGlobal(pos))
        if action == delete_action:
            self._delete_session(session_id)

    def _delete_session(self, session_id: str) -> None:
        reply = QMessageBox.question(
            self,
            "Delete chat",
            "Delete this chat? This can't be undone.",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No,
        )
        if reply != QMessageBox.StandardButton.Yes:
            return
        chatbot.delete_session(session_id)
        self._sessions = [s for s in self._sessions if s.id != session_id]
        if self._current_session.id == session_id:
            if self._sessions:
                self._current_session = self._sessions[0]
            else:
                self._current_session = chatbot.new_session()
            self._history = list(self._current_session.history)
            self.chat_view.clear()
            if self._current_session.history:
                self._render_session_history(self._current_session)
            else:
                self._append_system("Dev assist chatbot ready. Ask a development question.")
        self._refresh_session_list()

    def _on_model_loaded(self, llm) -> None:
        self._llm = llm
        self._n_ctx = llm.n_ctx()
        self.context_status_label.setText(f"Context: 0 / {self._n_ctx}")
        self.statusBar().showMessage("Ready.")
        self._append_system("Dev assist chatbot ready. Ask a development question.")
        self._set_input_enabled(True)

    def _on_model_failed(self, message: str) -> None:
        self.statusBar().showMessage("Failed to load model.")
        self._append_system(f"Error loading model: {message}")

    def _on_send(self) -> None:
        question = self.input_box.text().strip()
        if not question or self._llm is None or self._ask_thread is not None:
            return

        self.input_box.clear()
        self._append_user(question)
        self._set_input_enabled(False)
        self.statusBar().showMessage("Thinking...")
        self._stream_start_pos = None
        self._stream_body_cursor = None
        self._answer_started = False
        self._stream_thinking_id = None

        messages = chatbot.build_messages(self._history, question)
        self._pending_question = question

        self._ask_thread = AskThread(self._llm, messages)
        self._ask_thread.tool_call.connect(self._on_tool_call)
        self._ask_thread.usage.connect(self._on_usage)
        self._ask_thread.token.connect(self._on_token)
        self._ask_thread.thinking_token.connect(self._on_thinking_token)
        self._ask_thread.finished_ok.connect(self._on_answer)
        self._ask_thread.failed.connect(self._on_ask_failed)
        self._ask_thread.start()

    def _on_token(self, piece: str) -> None:
        if self._stream_body_cursor is None:
            self._collapse_open_thinking()
            self.statusBar().showMessage("Generating...")
            cursor = self.chat_view.textCursor()
            cursor.movePosition(QTextCursor.MoveOperation.End)
            self._stream_start_pos = cursor.position()
            accent, bg = ROLE_STYLES["assistant"]
            body = _insert_message_card(cursor, accent, bg)
            body.insertHtml(f'<span style="color:{accent}; font-weight:600; font-size:8.5pt;">ASSISTANT</span>')
            body.insertBlock(QTextBlockFormat(), QTextCharFormat())
            self._stream_body_cursor = body
            self._answer_started = False

        if not self._answer_started:
            piece = piece.lstrip("\n")
            if not piece:
                return
            self._answer_started = True

        char_format = QTextCharFormat()
        char_format.setForeground(QColor(TEXT_PRIMARY))
        self._stream_body_cursor.insertText(piece, char_format)
        self.chat_view.moveCursor(QTextCursor.MoveOperation.End)

    def _discard_stream_card(self) -> None:
        """Remove the in-progress streamed answer card, if any, so a
        tool-call round's raw output or a failed generation doesn't linger
        in the chat view."""
        if self._stream_start_pos is None:
            return
        cursor = self.chat_view.textCursor()
        cursor.setPosition(self._stream_start_pos)
        cursor.movePosition(QTextCursor.MoveOperation.End, QTextCursor.MoveMode.KeepAnchor)
        cursor.removeSelectedText()
        self._stream_start_pos = None
        self._stream_body_cursor = None

    def _on_thinking_token(self, piece: str) -> None:
        if self._stream_thinking_id is None:
            cursor = self.chat_view.textCursor()
            cursor.movePosition(QTextCursor.MoveOperation.End)
            accent, bg = ROLE_STYLES["thinking"]
            body = _insert_message_card(cursor, accent, bg)
            body.insertHtml(f'<span style="color:{accent}; font-weight:600; font-size:8.5pt;">THINKING</span>')
            body.insertBlock(QTextBlockFormat(), QTextCharFormat())

            self._thinking_seq += 1
            block_id = f"think-{self._thinking_seq}"
            content_start = QTextCursor(body)
            # Without this, inserting the header right at content_start's own
            # position (both here and on every later toggle) would drag it
            # forward along with that insertion, per QTextCursor's default -
            # it needs to stay pinned at the true start of the region instead.
            content_start.setKeepPositionOnInsert(True)
            self._thinking_blocks[block_id] = {
                "text": "",
                "expanded": True,
                "content_start": content_start,
                "cursor": body,
            }
            self._insert_thinking_header(body, block_id, expanded=True)
            body.insertBlock(QTextBlockFormat(), QTextCharFormat())
            self._stream_thinking_id = block_id

        block = self._thinking_blocks[self._stream_thinking_id]
        block["text"] += piece
        italic_format = QTextCharFormat()
        italic_format.setForeground(QColor(TEXT_SECONDARY))
        italic_format.setFontItalic(True)
        block["cursor"].insertText(piece, italic_format)
        self.chat_view.moveCursor(QTextCursor.MoveOperation.End)

    def _insert_thinking_header(self, cursor: QTextCursor, block_id: str, expanded: bool) -> None:
        arrow = "▼" if expanded else "▶"
        cursor.insertHtml(
            f'<a href="{_THINK_ANCHOR_PREFIX}{block_id}" '
            f'style="color:{TEXT_SECONDARY}; text-decoration:none;">{arrow} Thinking</a>'
        )

    def _render_thinking(self, block_id: str, expanded: bool) -> None:
        """Rewrite a thinking block's header and body to match expanded,
        replacing everything from its fixed content_start cursor through
        wherever its content currently ends (tracked by its own cursor,
        which - like any QTextCursor - keeps up with edits made elsewhere in
        the document, so this stays correct even if other messages have
        since been appended below it)."""
        block = self._thinking_blocks[block_id]
        cursor = block["cursor"]
        old_end = cursor.position()
        cursor.setPosition(block["content_start"].position())
        cursor.setPosition(old_end, QTextCursor.MoveMode.KeepAnchor)
        cursor.removeSelectedText()

        self._insert_thinking_header(cursor, block_id, expanded)
        if expanded:
            cursor.insertBlock(QTextBlockFormat(), QTextCharFormat())
            italic_format = QTextCharFormat()
            italic_format.setForeground(QColor(TEXT_SECONDARY))
            italic_format.setFontItalic(True)
            cursor.insertText(block["text"], italic_format)
        block["expanded"] = expanded

    def _collapse_open_thinking(self) -> None:
        if self._stream_thinking_id is not None:
            self._render_thinking(self._stream_thinking_id, expanded=False)
            self._stream_thinking_id = None
            self.chat_view.moveCursor(QTextCursor.MoveOperation.End)

    def _on_chat_anchor_clicked(self, url: QUrl) -> None:
        href = url.toString()
        if not href.startswith(_THINK_ANCHOR_PREFIX):
            QDesktopServices.openUrl(url)
            return
        # A generation currently in flight may be inserting content through
        # this same block's cursor further down; leave the block alone
        # rather than risk racing that live cursor.
        if self._ask_thread is not None:
            return
        block_id = href[len(_THINK_ANCHOR_PREFIX) :]
        block = self._thinking_blocks.get(block_id)
        if block is None:
            return
        self._render_thinking(block_id, expanded=not block["expanded"])

    def _on_tool_call(self, query: str) -> None:
        self._collapse_open_thinking()
        self._discard_stream_card()
        self.statusBar().showMessage(f"Searching: {query}")
        self._append_system(f"Searching the web for: {query}")

    def _on_usage(self, usage: dict) -> None:
        total = usage["total_tokens"]
        pct = 100 * total / self._n_ctx if self._n_ctx else 0
        self.context_status_label.setText(f"Context: {total} / {self._n_ctx} ({pct:.0f}%)")

    def _on_answer(self, answer: str) -> None:
        # Normally already collapsed by the first answer token (_on_token),
        # but a response that hit max_tokens mid-<think> would leave this
        # block open with nothing to trigger the collapse - catch that here.
        self._collapse_open_thinking()
        # The streamed card holds the raw, plain-text answer; swap it out for
        # a fully markdown-rendered and syntax-highlighted one now that the
        # complete text is available.
        self._discard_stream_card()
        self._append_assistant(answer)
        self._history.append((self._pending_question, answer.strip()))
        self._current_session.history = self._history
        if self._current_session.title == chatbot.DEFAULT_TITLE:
            self._current_session.title = chatbot.derive_title(self._pending_question)
        chatbot.save_session(self._current_session)
        self._sessions = [s for s in self._sessions if s.id != self._current_session.id]
        self._sessions.insert(0, self._current_session)
        self._refresh_session_list()
        self.statusBar().showMessage("Ready.")
        self._ask_thread = None
        self._set_input_enabled(True)

    def _on_ask_failed(self, message: str) -> None:
        self._collapse_open_thinking()
        self._discard_stream_card()
        self._append_system(f"Error: {message}")
        self.statusBar().showMessage("Ready.")
        self._ask_thread = None
        self._set_input_enabled(True)

    def _append_user(self, text: str) -> None:
        accent, bg = ROLE_STYLES["user"]
        self._append_block("You", text, accent, bg)

    def _append_assistant(self, text: str) -> None:
        accent, bg = ROLE_STYLES["assistant"]
        self._append_markdown_block("Assistant", text, accent, bg)

    def _append_system(self, text: str) -> None:
        accent, bg = ROLE_STYLES["system"]
        self._append_block("System", text, accent, bg)

    def _append_block(self, label: str, text: str, accent: str, bg: str) -> None:
        escaped = (
            text.replace("&", "&amp;")
            .replace("<", "&lt;")
            .replace(">", "&gt;")
            .replace("\n", "<br>")
        )
        cursor = self.chat_view.textCursor()
        cursor.movePosition(QTextCursor.MoveOperation.End)
        body = _insert_message_card(cursor, accent, bg)
        body.insertHtml(
            f'<span style="color:{accent}; font-weight:600; font-size:8.5pt;">{label.upper()}</span>'
            f'<br><span style="color:{TEXT_PRIMARY};">{escaped}</span>'
        )
        self.chat_view.moveCursor(QTextCursor.MoveOperation.End)

    def _append_markdown_block(self, label: str, markdown_text: str, accent: str, bg: str) -> None:
        cursor = self.chat_view.textCursor()
        cursor.movePosition(QTextCursor.MoveOperation.End)
        # Each message gets its own table cell (see _insert_message_card), so
        # unlike the previous shared-cursor version, there's no need to reset
        # formatting before the label - a fresh cell starts with plain format.
        body = _insert_message_card(cursor, accent, bg)
        body.insertHtml(f'<span style="color:{accent}; font-weight:600; font-size:8.5pt;">{label.upper()}</span>')

        pos = 0
        for match in _CODE_FENCE_RE.finditer(markdown_text):
            before = markdown_text[pos : match.start()]
            if before.strip():
                _reset_cursor_format(body)
                fragment = QTextDocumentFragment.fromMarkdown(before, _MARKDOWN_FEATURES)
                body.insertFragment(fragment)
            lang, code = match.group(1), match.group(2)
            for line in _highlight_code_lines(code, lang):
                body.insertBlock(_CODE_BLOCK_FORMAT)
                start = body.position()
                body.insertHtml(line)
                # insertHtml() builds its own char formats from the parsed
                # HTML and ignores the cursor's format, so the monospace
                # font has to be merged on afterwards; merging (rather than
                # replacing) keeps Pygments' per-token colors and bold.
                line_cursor = QTextCursor(body)
                line_cursor.setPosition(start)
                line_cursor.setPosition(body.position(), QTextCursor.MoveMode.KeepAnchor)
                line_cursor.mergeCharFormat(_CODE_CHAR_FORMAT)
            pos = match.end()
        remainder = markdown_text[pos:]
        if remainder.strip():
            _reset_cursor_format(body)
            fragment = QTextDocumentFragment.fromMarkdown(remainder, _MARKDOWN_FEATURES)
            body.insertFragment(fragment)

        self.chat_view.moveCursor(QTextCursor.MoveOperation.End)

    def closeEvent(self, event) -> None:
        # llama.cpp calls run in native code and can't be cancelled mid-call,
        # so wait for any in-flight load/generation to finish rather than
        # tearing down the thread while it's still running (which crashes Qt).
        for thread in (self._download_thread, self._loader_thread, self._ask_thread):
            if thread is not None and thread.isRunning():
                thread.wait()
        if self._gpu_handle is not None:
            pynvml.nvmlShutdown()
        event.accept()


def main() -> None:
    app = QApplication(sys.argv)
    app.setStyle("Fusion")
    app.setPalette(_dark_palette())
    app.setStyleSheet(_STYLESHEET)
    app.setWindowIcon(_make_app_icon())
    window = ChatWindow()
    window.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    main()