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 QThread, QTimer, pyqtSignal
from PyQt6.QtGui import (
    QColor,
    QFont,
    QTextBlockFormat,
    QTextCharFormat,
    QTextCursor,
    QTextDocument,
    QTextDocumentFragment,
)
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

# 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};">{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())


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)

        self.chat_view = QTextBrowser()
        self.chat_view.setOpenExternalLinks(True)
        self.chat_view.setFont(QFont("Segoe UI", 10))
        layout.addWidget(self.chat_view)

        input_row = QHBoxLayout()
        self.input_box = QLineEdit()
        self.input_box.setPlaceholderText("Ask a development question...")
        self.input_box.returnPressed.connect(self._on_send)
        input_row.addWidget(self.input_box)

        self.send_button = QPushButton("Send")
        self.send_button.clicked.connect(self._on_send)
        input_row.addWidget(self.send_button)

        layout.addLayout(input_row)

        self._n_ctx = 0
        self.context_status_label = QLabel("Context: --")
        self.statusBar().addPermanentWidget(self.context_status_label)

        self.gpu_status_label = QLabel()
        self.statusBar().addPermanentWidget(self.gpu_status_label)
        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:
        self._append_block("You", text, "#2b6cb0")

    def _append_assistant(self, text: str) -> None:
        self._append_markdown_block("Assistant", text, "#2f855a")

    def _append_system(self, text: str) -> None:
        self._append_block("System", text, "#718096")

    def _append_block(self, label: str, text: str, color: str) -> None:
        escaped = (
            text.replace("&", "&amp;")
            .replace("<", "&lt;")
            .replace(">", "&gt;")
            .replace("\n", "<br>")
        )
        self.chat_view.append(f'<p><b style="color:{color};">{label}:</b><br>{escaped}</p>')
        self.chat_view.moveCursor(QTextCursor.MoveOperation.End)

    def _append_markdown_block(self, label: str, markdown_text: str, color: str) -> None:
        cursor = self.chat_view.textCursor()
        cursor.movePosition(QTextCursor.MoveOperation.End)
        # Force a fresh, format-reset paragraph before the label - moving to
        # End only puts the cursor after the previous message's last
        # character (inheriting its block format, e.g. a Pygments code
        # block's background color), not on a clean new line.
        if not self.chat_view.document().isEmpty():
            _reset_cursor_format(cursor)
        cursor.insertHtml(f'<b style="color:{color};">{label}:</b>')

        pos = 0
        for match in _CODE_FENCE_RE.finditer(markdown_text):
            before = markdown_text[pos : match.start()]
            if before.strip():
                _reset_cursor_format(cursor)
                fragment = QTextDocumentFragment.fromMarkdown(before, _MARKDOWN_FEATURES)
                cursor.insertFragment(fragment)
            lang, code = match.group(1), match.group(2)
            for line in _highlight_code_lines(code, lang):
                cursor.insertBlock(_CODE_BLOCK_FORMAT)
                start = cursor.position()
                cursor.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(cursor)
                line_cursor.setPosition(start)
                line_cursor.setPosition(cursor.position(), QTextCursor.MoveMode.KeepAnchor)
                line_cursor.mergeCharFormat(_CODE_CHAR_FORMAT)
            pos = match.end()
        remainder = markdown_text[pos:]
        if remainder.strip():
            _reset_cursor_format(cursor)
            fragment = QTextDocumentFragment.fromMarkdown(remainder, _MARKDOWN_FEATURES)
            cursor.insertFragment(fragment)

        # Reset the block format one more time so a code block's background/
        # monospace font doesn't bleed into whatever message gets appended next.
        _reset_cursor_format(cursor)
        self.chat_view.setTextCursor(cursor)
        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)
    window = ChatWindow()
    window.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    main()