agr-assistent/src/agr_assistent/ui/chat_window.py

"""Окно чата: история диалога и поле ввода."""

from __future__ import annotations

from dataclasses import dataclass

from PySide6.QtCore import Qt, QTimer, Signal
from PySide6.QtGui import (
    QCloseEvent,
    QColor,
    QFont,
    QKeyEvent,
    QTextBlockFormat,
    QTextCharFormat,
    QTextCursor,
)
from PySide6.QtWidgets import (
    QHBoxLayout,
    QLabel,
    QPlainTextEdit,
    QPushButton,
    QTextBrowser,
    QVBoxLayout,
    QWidget,
)

from agr_assistent import APP_NAME
from agr_assistent.core.assistant import Assistant, AssistantState
from agr_assistent.ui.icons import state_icon

_ROLE_STYLES = {
    "user": ("Вы", "#1e88e5"),
    "assistant": ("Ассистент", "#43a047"),
    "error": ("Ошибка", "#e53935"),
}

_VOICE_BUTTON_TEXTS = {
    AssistantState.LISTENING: "Готово",
    AssistantState.RECOGNIZING: "Распознаю…",
}

# Во время стриминга перерисовываем историю не чаще, чем раз в N мс
_RENDER_INTERVAL_MS = 50


@dataclass
class _Entry:
    role: str
    text: str


class _MessageInput(QPlainTextEdit):
    """Enter отправляет сообщение, Shift+Enter переносит строку."""

    submitted = Signal()

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


class ChatWindow(QWidget):
    def __init__(self, assistant: Assistant) -> None:
        super().__init__()
        self._assistant = assistant
        self._entries: list[_Entry] = []

        self.setWindowTitle(APP_NAME)
        self.setWindowIcon(state_icon(assistant.state))
        self.resize(560, 680)

        self._provider_label = QLabel()
        new_chat_button = QPushButton("Новый диалог")
        new_chat_button.clicked.connect(assistant.clear_history)

        self._transcript = QTextBrowser()
        self._transcript.setOpenExternalLinks(True)

        self._input = _MessageInput()
        self._input.setPlaceholderText(
            "Напишите сообщение… (Enter — отправить, Shift+Enter — новая строка)"
        )
        self._input.setMaximumHeight(90)
        self._input.submitted.connect(self._submit)

        self._send_button = QPushButton()
        self._send_button.setMinimumWidth(100)
        self._send_button.clicked.connect(self._on_send_clicked)

        self._voice_button = QPushButton()
        self._voice_button.setMinimumWidth(100)
        self._voice_button.clicked.connect(assistant.toggle_listening)
        self._voice_button.setVisible(assistant.voice_available)

        header = QHBoxLayout()
        header.addWidget(self._provider_label, 1)
        header.addWidget(new_chat_button)

        buttons = QVBoxLayout()
        buttons.addWidget(self._send_button)
        buttons.addWidget(self._voice_button)
        buttons.addStretch()

        input_row = QHBoxLayout()
        input_row.addWidget(self._input, 1)
        input_row.addLayout(buttons)

        layout = QVBoxLayout(self)
        layout.addLayout(header)
        layout.addWidget(self._transcript, 1)
        layout.addLayout(input_row)

        self._render_timer = QTimer(self)
        self._render_timer.setSingleShot(True)
        self._render_timer.setInterval(_RENDER_INTERVAL_MS)
        self._render_timer.timeout.connect(self._render)

        assistant.state_changed.connect(self._on_state_changed)
        assistant.provider_changed.connect(self._update_provider_label)
        assistant.user_message_added.connect(lambda text: self._append("user", text))
        assistant.reply_started.connect(lambda: self._append("assistant", ""))
        assistant.reply_chunk.connect(self._on_reply_chunk)
        assistant.reply_finished.connect(self._on_reply_finished)
        assistant.error_occurred.connect(lambda message: self._append("error", message))
        assistant.history_cleared.connect(self._on_history_cleared)

        self._update_provider_label()
        self._on_state_changed(assistant.state)

    def show_and_raise(self) -> None:
        self.showNormal()
        self.raise_()
        self.activateWindow()
        self._input.setFocus()

    def toggle_visibility(self) -> None:
        if self.isVisible() and not self.isMinimized():
            self.hide()
        else:
            self.show_and_raise()

    def closeEvent(self, event: QCloseEvent) -> None:
        # Закрытие окна сворачивает приложение в трей; выход — через меню трея
        event.ignore()
        self.hide()

    def _submit(self) -> None:
        text = self._input.toPlainText().strip()
        if not text or self._assistant.is_generating:
            return
        self._input.clear()
        self._assistant.send(text)

    def _on_send_clicked(self) -> None:
        if self._assistant.is_busy:
            self._assistant.cancel()
        else:
            self._submit()

    def _on_state_changed(self, state: object) -> None:
        current = self._assistant.state
        self._send_button.setText("Стоп" if self._assistant.is_busy else "Отправить")
        self._voice_button.setText(_VOICE_BUTTON_TEXTS.get(current, "Говорить"))
        self._voice_button.setEnabled(current is not AssistantState.RECOGNIZING)
        self.setWindowIcon(state_icon(current))

    def _update_provider_label(self) -> None:
        self._provider_label.setText(
            f"Модель: <b>{self._assistant.provider_name}</b> · {self._assistant.model_name}"
        )

    def _on_reply_chunk(self, piece: str) -> None:
        if self._entries and self._entries[-1].role == "assistant":
            self._entries[-1].text += piece
            if not self._render_timer.isActive():
                self._render_timer.start()

    def _on_reply_finished(self, text: str) -> None:
        last = self._entries[-1] if self._entries else None
        if last and last.role == "assistant" and not last.text.strip():
            self._entries.pop()
        self._render_now()

    def _on_history_cleared(self) -> None:
        self._entries.clear()
        self._render_now()

    def _append(self, role: str, text: str) -> None:
        self._entries.append(_Entry(role, text))
        self._render_now()

    def _render_now(self) -> None:
        self._render_timer.stop()
        self._render()

    def _render(self) -> None:
        scrollbar = self._transcript.verticalScrollBar()
        previous_position = scrollbar.value()
        stick_to_bottom = previous_position >= scrollbar.maximum() - 4

        document = self._transcript.document()
        document.clear()
        cursor = QTextCursor(document)
        for index, entry in enumerate(self._entries):
            title, color = _ROLE_STYLES[entry.role]

            title_block = QTextBlockFormat()
            title_block.setTopMargin(14 if index else 0)
            if index:
                cursor.insertBlock(title_block, QTextCharFormat())
            else:
                cursor.setBlockFormat(title_block)

            title_format = QTextCharFormat()
            title_format.setFontWeight(QFont.Weight.Bold)
            title_format.setForeground(QColor(color))
            cursor.insertText(title, title_format)

            cursor.insertBlock(QTextBlockFormat(), QTextCharFormat())
            if entry.role == "assistant":
                cursor.insertMarkdown(entry.text or "…")
            else:
                cursor.insertText(entry.text)

        # Диапазон прокрутки обновляется после раскладки документа
        QTimer.singleShot(
            0,
            lambda: scrollbar.setValue(
                scrollbar.maximum() if stick_to_bottom else previous_position
            ),
        )