qt_gui.py

from __future__ import annotations

import html
import json
import os
from pathlib import Path
import sys
from threading import Event

WEBENGINE_DATA_DIR = Path(__file__).resolve().parent / ".qtwebengine"
WEBENGINE_DATA_DIR.mkdir(exist_ok=True)
os.environ.setdefault("XDG_CACHE_HOME", str(WEBENGINE_DATA_DIR / "cache"))
os.environ.setdefault("XDG_CONFIG_HOME", str(WEBENGINE_DATA_DIR / "config"))

from PySide6.QtCore import QObject, QThread, Qt, Signal, Slot
from PySide6.QtGui import QCloseEvent, QKeyEvent
from PySide6.QtWebEngineWidgets import QWebEngineView
from PySide6.QtWidgets import (
    QApplication,
    QFrame,
    QHBoxLayout,
    QLabel,
    QMainWindow,
    QPushButton,
    QTextEdit,
    QVBoxLayout,
    QWidget,
)

from chat_memory import ChatMemory
from config import DEEPSEEK_MODEL, MAX_USER_MESSAGE_CHARS
from deepseek_client import (
    DeepSeekAPIError,
    DeepSeekClient,
    DeepSeekConfigurationError,
    DeepSeekNetworkError,
    DeepSeekResponseError,
    RequestCancelled,
)
from rich_text import render_message_html


WELCOME_TEXT = "\u4f60\u597d\uff01\u6211\u662f\u4f60\u7684\u684c\u9762 AI \u52a9\u624b\u3002\u8f93\u5165\u95ee\u9898\u540e\u6309 Enter \u5c31\u53ef\u4ee5\u5f00\u59cb\u804a\u5929\u3002"
CLEAR_TEXT = "\u6e05\u7a7a\u5bf9\u8bdd"
SEND_TEXT = "\u53d1\u9001"
STOP_TEXT = "\u505c\u6b62"
STOPPING_TEXT = "\u505c\u6b62\u4e2d..."
PLACEHOLDER_TEXT = "\u8f93\u5165\u6d88\u606f\uff0cEnter \u53d1\u9001\uff0cShift+Enter \u6362\u884c"
THINKING_TEXT = "AI \u6b63\u5728\u601d\u8003..."
REASONING_TITLE = "\u601d\u8003\u8fc7\u7a0b"
STOPPED_TEXT = "\u5df2\u505c\u6b62\u751f\u6210\u3002"
ERROR_PREFIX = "\u51fa\u9519\u4e86\uff1a"
INPUT_TOO_LONG_TEXT = "\u8f93\u5165\u5185\u5bb9\u8fc7\u957f\uff0c\u5f53\u524d\u6700\u591a\u652f\u6301 {limit} \u4e2a\u5b57\u7b26\u3002"
CHAT_PAGE_LOAD_ERROR_TEXT = "\u804a\u5929\u9875\u9762\u52a0\u8f7d\u5931\u8d25\uff0c\u8bf7\u91cd\u65b0\u542f\u52a8\u5e94\u7528\u540e\u518d\u8bd5\u3002"
MAX_PENDING_SCRIPTS = 100


class ApiWorker(QObject):
    chunk = Signal(str, str)
    finished = Signal(str, object)

    def __init__(self, client: DeepSeekClient, messages: list[dict[str, str]]) -> None:
        super().__init__()
        self.client = client
        self.messages = messages
        self.cancel_event = Event()

    def cancel(self) -> None:
        self.cancel_event.set()

    @Slot()
    def run(self) -> None:
        full_reply = []
        try:
            for kind, delta in self.client.stream_chat_events(self.messages, self.cancel_event.is_set):
                if kind == "content":
                    full_reply.append(delta)
                self.chunk.emit(kind, delta)
            self.finished.emit("ok", "".join(full_reply))
        except RequestCancelled as exc:
            self.finished.emit("cancelled", str(exc))
        except Exception as exc:
            self.finished.emit("error", exc)


class ChatInput(QTextEdit):
    send_requested = Signal()

    def keyPressEvent(self, event: QKeyEvent) -> None:
        if event.key() in (Qt.Key_Return, Qt.Key_Enter) and not event.modifiers() & Qt.ShiftModifier:
            self.send_requested.emit()
            return
        super().keyPressEvent(event)


class ChatWindow(QMainWindow):
    def __init__(self) -> None:
        super().__init__()
        self.setWindowTitle("PyDeepChat")
        self.resize(980, 720)
        self.setMinimumSize(760, 560)

        self.client = DeepSeekClient()
        self.configuration_error: DeepSeekConfigurationError | None = None
        try:
            self.client.validate_configuration()
        except DeepSeekConfigurationError as exc:
            self.configuration_error = exc
        self.memory = ChatMemory()
        self.is_waiting = False
        self.thread: QThread | None = None
        self.worker: ApiWorker | None = None
        self.pending_scripts: list[str] = []
        self.page_ready = False
        self.message_counter = 0
        self.active_assistant_id: str | None = None
        self.active_assistant_text = ""
        self.active_reasoning_text = ""

        self._build_layout()
        self._apply_styles()
        self._init_chat_document()

    def _build_layout(self) -> None:
        root = QWidget()
        self.setCentralWidget(root)

        layout = QVBoxLayout(root)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)

        header = QFrame()
        header.setObjectName("header")
        header_layout = QHBoxLayout(header)
        header_layout.setContentsMargins(24, 14, 24, 14)

        title_box = QVBoxLayout()
        title = QLabel("PyDeepChat")
        title.setObjectName("title")
        subtitle = QLabel(f"模型：{DEEPSEEK_MODEL}")
        subtitle.setObjectName("subtitle")
        title_box.addWidget(title)
        title_box.addWidget(subtitle)
        header_layout.addLayout(title_box)
        header_layout.addStretch()

        self.clear_button = QPushButton(CLEAR_TEXT)
        self.clear_button.clicked.connect(self._clear_chat)
        header_layout.addWidget(self.clear_button)
        layout.addWidget(header)

        self.chat_view = QWebEngineView()
        self.chat_view.setObjectName("chatView")
        self.chat_view.loadFinished.connect(self._on_chat_page_ready)
        layout.addWidget(self.chat_view, stretch=1)

        input_panel = QFrame()
        input_panel.setObjectName("inputPanel")
        input_layout = QHBoxLayout(input_panel)
        input_layout.setContentsMargins(18, 14, 18, 14)
        input_layout.setSpacing(12)

        self.input_box = ChatInput()
        self.input_box.setObjectName("inputBox")
        self.input_box.setPlaceholderText(PLACEHOLDER_TEXT)
        self.input_box.setFixedHeight(92)
        self.input_box.send_requested.connect(self._send_message)
        input_layout.addWidget(self.input_box, stretch=1)

        self.send_button = QPushButton(SEND_TEXT)
        self.send_button.setObjectName("sendButton")
        self.send_button.setFixedSize(96, 46)
        self.send_button.clicked.connect(self._on_send_button_clicked)
        input_layout.addWidget(self.send_button)

        layout.addWidget(input_panel)

    def _apply_styles(self) -> None:
        self.setStyleSheet(
            """
            QMainWindow { background: #eef2f7; }
            #header { background: #172033; }
            #title { color: #ffffff; font-size: 18px; font-weight: 700; }
            #subtitle { color: #cbd5e1; font-size: 12px; }
            QPushButton {
                border: 0;
                border-radius: 7px;
                padding: 8px 14px;
                background: #e5e7eb;
                color: #111827;
                font-weight: 600;
            }
            QPushButton:hover { background: #d1d5db; }
            QPushButton:disabled { background: #cbd5e1; color: #64748b; }
            #sendButton { background: #2563eb; color: #ffffff; }
            #sendButton:hover { background: #1d4ed8; }
            #inputPanel { background: #ffffff; border-top: 1px solid #e5e7eb; }
            #inputBox {
                border: 1px solid #cbd5e1;
                border-radius: 8px;
                background: #f8fafc;
                color: #111827;
                padding: 10px;
                font-size: 14px;
            }
            """
        )

    def _init_chat_document(self) -> None:
        self.page_ready = False
        self.pending_scripts.clear()
        self.chat_view.setHtml(_document_html())

    @Slot(bool)
    def _on_chat_page_ready(self, ok: bool) -> None:
        self.page_ready = ok
        if not ok:
            self.pending_scripts.clear()
            self.chat_view.setHtml(_fallback_error_document(CHAT_PAGE_LOAD_ERROR_TEXT))
            return

        self._add_message("assistant", WELCOME_TEXT)
        if self.configuration_error is not None:
            self._add_message("assistant", format_user_error(self.configuration_error))
            self.send_button.setDisabled(True)
        for script in self.pending_scripts:
            self.chat_view.page().runJavaScript(script)
        self.pending_scripts.clear()

    def _on_send_button_clicked(self) -> None:
        if self.is_waiting:
            self._cancel_generation()
            return
        self._send_message()

    def _send_message(self) -> None:
        if self.is_waiting:
            return
        if self.configuration_error is not None:
            self._add_message("assistant", format_user_error(self.configuration_error))
            return

        content = _clean_user_input(self.input_box.toPlainText()).strip()
        if not content:
            return
        if len(content) > MAX_USER_MESSAGE_CHARS:
            self._add_message("assistant", INPUT_TOO_LONG_TEXT.format(limit=MAX_USER_MESSAGE_CHARS))
            return

        self.input_box.clear()
        self._add_message("user", content)
        self.memory.add_user_message(content)
        self._set_waiting(True)

        self.active_assistant_text = ""
        self.active_assistant_id = self._add_message("assistant", THINKING_TEXT, temporary=True)

        self.thread = QThread()
        self.worker = ApiWorker(self.client, self.memory.get_messages())
        self.worker.moveToThread(self.thread)
        self.thread.started.connect(self.worker.run)
        self.worker.chunk.connect(self._append_assistant_chunk)
        self.worker.finished.connect(self._handle_response)
        self.worker.finished.connect(self.thread.quit)
        self.worker.finished.connect(self.worker.deleteLater)
        self.thread.finished.connect(self.thread.deleteLater)
        self.thread.finished.connect(self._clear_worker_refs)
        self.thread.start()

    @Slot(str, str)
    def _append_assistant_chunk(self, kind: str, delta: str) -> None:
        if self.active_assistant_id is None:
            return
        if kind == "reasoning":
            self.active_reasoning_text += delta
            if self.active_assistant_text:
                live_html = _render_plain_text(self.active_assistant_text)
            else:
                live_html = _render_reasoning_preview(self.active_reasoning_text)
        elif kind == "content":
            self.active_assistant_text += delta
            live_html = _render_plain_text(self.active_assistant_text or THINKING_TEXT)
        else:
            live_html = _render_plain_text(self.active_assistant_text or THINKING_TEXT)
        self._update_message(self.active_assistant_id, live_html, typeset=False)

    @Slot(str, str)
    def _handle_response(self, status: str, content: object) -> None:
        self._set_waiting(False)

        if self.active_assistant_id is None:
            self.active_assistant_id = self._add_message("assistant", "")

        if status == "ok":
            content = str(content)
            self.memory.add_assistant_message(content)
            final_html = render_message_html(content)
            if self.active_reasoning_text.strip():
                final_html = _render_collapsed_reasoning(self.active_reasoning_text) + final_html
            self._update_message(self.active_assistant_id, final_html, typeset=True)
        elif status == "cancelled":
            stopped = self.active_assistant_text or STOPPED_TEXT
            if self.active_assistant_text:
                stopped = f"{stopped}\n\n{STOPPED_TEXT}"
            self._update_message(self.active_assistant_id, render_message_html(stopped), typeset=True)
        else:
            self._update_message(self.active_assistant_id, render_message_html(format_user_error(content)), typeset=True)

        self.active_assistant_id = None
        self.active_assistant_text = ""
        self.active_reasoning_text = ""

    def _cancel_generation(self, update_button: bool = True) -> None:
        if self.worker is not None:
            self.worker.cancel()
        if self.thread is not None:
            self.thread.requestInterruption()
        if update_button and self.is_waiting:
            self.send_button.setDisabled(True)
            self.send_button.setText(STOPPING_TEXT)

    @Slot()
    def _clear_worker_refs(self) -> None:
        self.thread = None
        self.worker = None

    def _set_waiting(self, waiting: bool) -> None:
        self.is_waiting = waiting
        self.send_button.setDisabled(False)
        self.send_button.setText(STOP_TEXT if waiting else SEND_TEXT)
        self.clear_button.setDisabled(waiting)
        self.input_box.setDisabled(waiting)
        if not waiting:
            self.input_box.setFocus()

    def _add_message(self, role: str, content: str, temporary: bool = False) -> str:
        self.message_counter += 1
        message_id = f"msg-{self.message_counter}"
        is_user = role == "user"
        align = "right" if is_user else "left"
        bubble_class = "user-bubble" if is_user else "assistant-bubble"
        text_color = "#ffffff" if is_user else "#111827"
        rendered = _render_plain_user_text(content) if is_user else render_message_html(content, text_color)

        self._run_script(
            "window.chatApp.appendMessage("
            f"{json.dumps(message_id)}, "
            f"{json.dumps(align)}, "
            f"{json.dumps(bubble_class)}, "
            f"{json.dumps(rendered)}, "
            f"{json.dumps(temporary)}"
            ");"
        )
        return message_id

    def _update_message(self, message_id: str, rendered_html: str, typeset: bool) -> None:
        self._run_script(
            "window.chatApp.updateMessage("
            f"{json.dumps(message_id)}, "
            f"{json.dumps(rendered_html)}, "
            f"{json.dumps(typeset)}"
            ");"
        )

    def _clear_chat(self) -> None:
        if self.is_waiting:
            return
        self.memory.reset()
        self.message_counter = 0
        self.active_assistant_id = None
        self.active_assistant_text = ""
        self.active_reasoning_text = ""
        self._run_script("window.chatApp.clearMessages();")
        self._add_message("assistant", WELCOME_TEXT)
        self.input_box.setFocus()

    def _run_script(self, script: str) -> None:
        if self.page_ready:
            self.chat_view.page().runJavaScript(script)
        else:
            if len(self.pending_scripts) >= MAX_PENDING_SCRIPTS:
                self.pending_scripts.pop(0)
            self.pending_scripts.append(script)

    def closeEvent(self, event: QCloseEvent) -> None:
        self._cancel_generation(update_button=False)
        if self.thread is not None and self.thread.isRunning():
            self.thread.quit()
            if not self.thread.wait(3000):
                event.ignore()
                return
        self.client.close()
        event.accept()


def _render_plain_user_text(content: str) -> str:
    return _render_plain_text(content)


def _clean_user_input(content: str) -> str:
    return "".join(
        char
        for char in content
        if char in "\n\t" or (ord(char) >= 32 and ord(char) != 127)
    )


def _render_reasoning_preview(content: str) -> str:
    escaped = html.escape(content).replace("\n", "<br>")
    return (
        f"<div class='reasoning-preview'>"
        f"<div class='reasoning-title'>{REASONING_TITLE}</div>"
        f"<div>{escaped}</div>"
        f"</div>"
    )


def _render_collapsed_reasoning(content: str) -> str:
    escaped = html.escape(content).replace("\n", "<br>")
    return (
        "<details class='reasoning-details'>"
        f"<summary>{REASONING_TITLE}</summary>"
        f"<div class='reasoning-body'>{escaped}</div>"
        "</details>"
    )


def format_user_error(error: str | Exception) -> str:
    if isinstance(error, DeepSeekConfigurationError):
        return (
            f"{ERROR_PREFIX}\n\n"
            "DeepSeek API Key 未正确配置。请在当前 PowerShell 中设置有效的 `DEEPSEEK_API_KEY`，"
            "例如：`$env:DEEPSEEK_API_KEY=\"sk-...\"`。"
        )
    if isinstance(error, DeepSeekNetworkError):
        return f"{ERROR_PREFIX}\n\n网络连接失败或请求超时，请检查网络后重试。"
    if isinstance(error, DeepSeekAPIError):
        return f"{ERROR_PREFIX}\n\nDeepSeek API 返回错误，请检查 API Key、账户额度、模型名称或请求参数。"
    if isinstance(error, DeepSeekResponseError):
        return f"{ERROR_PREFIX}\n\nDeepSeek API 返回了暂不支持的响应格式，请稍后重试。"
    return f"{ERROR_PREFIX}\n\n{error}"


def _render_plain_text(content: str) -> str:
    return "<p>" + html.escape(content).replace("\n", "<br>") + "</p>"


def _document_html() -> str:
    return """
    <!doctype html>
    <html>
    <head>
        <meta charset="utf-8">
        <script>
            window.MathJax = {
                tex: {
                    inlineMath: [['\\\\(', '\\\\)']],
                    displayMath: [['\\\\[', '\\\\]'], ['$$', '$$']]
                },
                svg: { fontCache: 'global' },
                startup: {
                    pageReady: () => {
                        return MathJax.startup.defaultPageReady().then(() => {
                            window.mathJaxReady = true;
                            if (window.chatApp) {
                                window.chatApp.typesetPending();
                            }
                        });
                    }
                }
            };
        </script>
        <script async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-svg.js"></script>
        <style>
            body {
                margin: 0;
                padding: 22px 24px;
                background: #f8fafc;
                color: #111827;
                font-family: "Microsoft YaHei UI", "Segoe UI", sans-serif;
                font-size: 14px;
                line-height: 1.62;
            }
            .message-row { margin-bottom: 14px; }
            .assistant-bubble {
                max-width: 720px;
                background: #ffffff;
                border: 1px solid #e5e7eb;
                border-radius: 8px;
                padding: 15px 18px;
                overflow-wrap: anywhere;
                word-break: break-word;
            }
            .user-bubble {
                max-width: 620px;
                background: #2563eb;
                color: #ffffff;
                border-radius: 8px;
                padding: 12px 15px;
                overflow-wrap: anywhere;
                word-break: break-word;
            }
            p { margin: 5px 0 9px 0; }
            h1, h2, h3, h4 { margin: 13px 0 7px 0; font-weight: 700; }
            h1 { font-size: 20px; }
            h2 { font-size: 18px; }
            h3 { font-size: 16px; }
            h4 { font-size: 15px; }
            ul, ol { margin: 6px 0 9px 24px; padding: 0; }
            li { margin: 4px 0; }
            hr { border: 0; border-top: 1px solid #e5e7eb; margin: 13px 0; }
            blockquote {
                border-left: 4px solid #cbd5e1;
                color: #475569;
                margin: 8px 0;
                padding: 4px 0 4px 12px;
            }
            table {
                border-collapse: collapse;
                margin: 10px 0;
                width: 100%;
            }
            th, td {
                border: 1px solid #e5e7eb;
                padding: 7px 9px;
                text-align: left;
            }
            th { background: #f1f5f9; }
            code {
                background: #eef2f7;
                border-radius: 4px;
                padding: 1px 4px;
                font-family: Consolas, "Cascadia Mono", monospace;
            }
            pre {
                background: #0f172a;
                color: #e2e8f0;
                border-radius: 8px;
                padding: 12px;
                overflow-x: auto;
            }
            pre code {
                background: transparent;
                color: inherit;
                padding: 0;
            }
            .codehilite {
                background: #0f172a;
                color: #e2e8f0;
                border-radius: 8px;
                margin: 10px 0;
                overflow-x: auto;
            }
            .codehilite pre {
                margin: 0;
            }
            .codehilite .hll { background-color: #334155; }
            .codehilite .c, .codehilite .ch, .codehilite .cm, .codehilite .c1, .codehilite .cs { color: #94a3b8; font-style: italic; }
            .codehilite .k, .codehilite .kc, .codehilite .kd, .codehilite .kn, .codehilite .kp, .codehilite .kr, .codehilite .kt { color: #c084fc; }
            .codehilite .s, .codehilite .sa, .codehilite .sb, .codehilite .sc, .codehilite .dl, .codehilite .sd, .codehilite .s2, .codehilite .se, .codehilite .sh, .codehilite .si, .codehilite .sx, .codehilite .sr, .codehilite .s1, .codehilite .ss { color: #86efac; }
            .codehilite .m, .codehilite .mb, .codehilite .mf, .codehilite .mh, .codehilite .mi, .codehilite .mo { color: #fbbf24; }
            .codehilite .nf, .codehilite .fm { color: #67e8f9; }
            .codehilite .nc, .codehilite .nn { color: #93c5fd; }
            .codehilite .o, .codehilite .ow { color: #f472b6; }
            .codehilite .bp, .codehilite .nb { color: #fca5a5; }
            a { color: #2563eb; text-decoration: none; }
            a:hover { text-decoration: underline; }
            .formula-block {
                background: #f8fafc;
                border: 1px solid #e5e7eb;
                border-radius: 7px;
                padding: 12px;
                margin: 10px 0;
                text-align: center;
                overflow-x: auto;
            }
            .math-inline { white-space: nowrap; }
            .reasoning-preview {
                color: #64748b;
                background: #f8fafc;
                border-left: 4px solid #cbd5e1;
                padding: 8px 10px;
                border-radius: 6px;
            }
            .reasoning-title {
                color: #334155;
                font-weight: 700;
                margin-bottom: 6px;
            }
            .reasoning-details {
                color: #64748b;
                background: #f8fafc;
                border: 1px solid #e5e7eb;
                border-radius: 7px;
                margin-top: 12px;
                padding: 8px 10px;
            }
            .reasoning-details summary {
                color: #334155;
                cursor: pointer;
                font-weight: 700;
            }
            .reasoning-body {
                border-top: 1px solid #e5e7eb;
                margin-top: 8px;
                padding-top: 8px;
            }
            mjx-container { font-size: 108%; }
        </style>
    </head>
    <body>
        <main id="chat-root"></main>
        <script>
            window.mathJaxReady = false;
            window.pendingMathNodes = [];
            window.chatApp = {
                appendMessage(id, align, bubbleClass, html, temporary) {
                    const row = document.createElement('table');
                    row.className = 'message-row';
                    row.width = '100%';
                    row.cellSpacing = '0';
                    row.cellPadding = '0';
                    if (temporary) row.dataset.temporary = '1';
                    row.innerHTML = `<tr><td align="${align}"><div id="${id}" class="${bubbleClass}">${html}</div></td></tr>`;
                    document.getElementById('chat-root').appendChild(row);
                    this.typeset(row);
                },
                updateMessage(id, html, typeset) {
                    const node = document.getElementById(id);
                    if (!node) return;
                    node.innerHTML = html;
                    if (typeset) this.typeset(node);
                    this.scrollBottom();
                },
                clearMessages() {
                    document.getElementById('chat-root').innerHTML = '';
                },
                typeset(node) {
                    if (window.mathJaxReady && window.MathJax && window.MathJax.typesetPromise) {
                        window.MathJax.typesetPromise([node]).finally(() => this.scrollBottom());
                    } else {
                        window.pendingMathNodes.push(node);
                        this.scrollBottom();
                    }
                },
                typesetPending() {
                    const nodes = window.pendingMathNodes.splice(0);
                    if (nodes.length && window.MathJax && window.MathJax.typesetPromise) {
                        window.MathJax.typesetPromise(nodes).finally(() => this.scrollBottom());
                    }
                },
                scrollBottom() {
                    window.scrollTo(0, document.body.scrollHeight);
                }
            };
        </script>
    </body>
    </html>
    """


def _fallback_error_document(message: str) -> str:
    return f"""
    <!doctype html>
    <html>
    <head>
        <meta charset="utf-8">
        <style>
            body {{
                margin: 0;
                padding: 24px;
                background: #f8fafc;
                color: #111827;
                font-family: "Microsoft YaHei UI", "Segoe UI", sans-serif;
                font-size: 14px;
            }}
            .error {{
                background: #ffffff;
                border: 1px solid #fecaca;
                border-radius: 8px;
                color: #991b1b;
                padding: 14px 16px;
            }}
        </style>
    </head>
    <body>
        <div class="error">{html.escape(message)}</div>
    </body>
    </html>
    """


def run_app() -> None:
    app = QApplication(sys.argv)
    window = ChatWindow()
    window.show()
    sys.exit(app.exec())