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, pyqtSignal
from PyQt6.QtGui import (
    QColor,
    QFont,
    QIcon,
    QLinearGradient,
    QPainter,
    QPainterPath,
    QPalette,
    QPixmap,
    QTextBlockFormat,
    QTextCharFormat,
    QTextCursor,
    QTextDocument,
    QTextDocumentFragment,
    QTextLength,
    QTextTableCellFormat,
    QTextTableFormat,
)
from PyQt6.QtWidgets import (
    QApplication,
    QHBoxLayout,
    QLabel,
    QLineEdit,
    QMainWindow,
    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".
BG_APP = "#0b0d12"
BG_PANEL = "#12151c"
BG_ELEVATED = "#181c25"
BORDER = "#262b36"
BORDER_STRONG = "#333949"
TEXT_PRIMARY = "#e6e8ee"
TEXT_SECONDARY = "#8b92a5"
ACCENT = "#6366f1"
ACCENT_HOVER = "#7b7ef4"
ACCENT_PRESSED = "#4f52d6"

# (accent, background) per chat role, used for the accent stripe and tinted
# card background of each message.
ROLE_STYLES = {
    "user": ("#60a5fa", "#141b2b"),
    "assistant": ("#34d399", "#0f211c"),
    "system": ("#9aa1b4", "#161922"),
}

_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;
}}

#header {{
    border-bottom: 1px solid {BORDER};
}}

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

#subtitleLabel {{
    font-size: 9pt;
    color: {TEXT_SECONDARY};
}}

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

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

QLineEdit#messageInput {{
    background-color: {BG_ELEVATED};
    border: 1px solid {BORDER};
    border-radius: 19px;
    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: #ffffff;
    border: none;
    border-radius: 19px;
    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};
    border-radius: 5px;
    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("#ffffff"))
    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(ACCENT_HOVER))
    gradient.setColorAt(1.0, QColor(ACCENT_PRESSED))

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

    painter.setPen(QColor("#ffffff"))
    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)
_CODE_BLOCK_FORMAT.setTopMargin(1)
_CODE_BLOCK_FORMAT.setBottomMargin(1)

_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 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)
    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,
            )
        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._history: list[tuple[str, str]] = []
        self._pending_question = ""
        self._ask_thread: AskThread | None = None

        central = QWidget()
        self.setCentralWidget(central)
        layout = QVBoxLayout(central)
        layout.setContentsMargins(20, 16, 20, 16)
        layout.setSpacing(14)

        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")
        subtitle_label = QLabel("Local LLM · tool-augmented")
        subtitle_label.setObjectName("subtitleLabel")
        title_box.addWidget(title_label)
        title_box.addWidget(subtitle_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")
        self.chat_view.setOpenExternalLinks(True)
        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._init_gpu_monitor()

        self.statusBar().showMessage("Loading model...")
        self._set_input_enabled(False)

        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)
        self.gpu_status_label.setText(f"GPU {util.gpu}%  ·  VRAM {used_gb:.1f}/{total_gb:.1f} GB")

    def _set_input_enabled(self, enabled: bool) -> None:
        self.input_box.setEnabled(enabled)
        self.send_button.setEnabled(enabled)
        if enabled:
            self.input_box.setFocus()

    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...")

        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.finished_ok.connect(self._on_answer)
        self._ask_thread.failed.connect(self._on_ask_failed)
        self._ask_thread.start()

    def _on_tool_call(self, query: str) -> None:
        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:
        self._append_assistant(answer)
        self._history.append((self._pending_question, answer.strip()))
        self.statusBar().showMessage("Ready.")
        self._ask_thread = None
        self._set_input_enabled(True)

    def _on_ask_failed(self, message: str) -> None:
        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._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()