ai-chat-client/src/main/java/com/aiclient/controller/MainController.java

package com.aiclient.controller;

import javafx.fxml.FXML;
import javafx.scene.control.*;
import javafx.scene.layout.VBox;
import javafx.scene.layout.HBox;
import javafx.scene.text.Text;
import javafx.scene.text.TextAlignment;
import javafx.scene.text.TextFlow;
import com.aiclient.service.ChatService;
import com.aiclient.service.MarkdownService;
import com.aiclient.service.ThemeService;
import com.aiclient.model.Chat;
import com.aiclient.model.Message;
import javafx.scene.input.KeyCode;
import javafx.scene.input.KeyCodeCombination;
import javafx.scene.input.KeyCombination;
import com.aiclient.service.ShortcutService;
import javafx.scene.Node;
import javafx.scene.Scene;
import javafx.scene.input.Clipboard;
import javafx.scene.input.ClipboardContent;
import javafx.geometry.Pos;
import javafx.scene.text.Font;
import javafx.scene.control.ScrollPane;

public class MainController {

    @FXML
    public ScrollPane chatScrollPane;

    @FXML
    private ListView<Chat> chatListView;

    @FXML
    private ComboBox<String> modelSelector;

    @FXML
    private VBox chatContainer;

    @FXML
    private TextArea inputArea;

    private ChatService chatService;
    private MarkdownService markdownService;
    private ThemeService themeService;
    private Chat currentChat;
    private ShortcutService shortcutService;
    private Scene scene;

    @FXML
    public void initialize() {
        chatService = new ChatService();
        markdownService = new MarkdownService();
        themeService = new ThemeService();
        shortcutService = new ShortcutService();

        // 初始化模型选择器
        modelSelector.getItems().addAll("deepseek-coder:latest");
        modelSelector.setValue("deepseek-coder:latest");

        // 初始化聊天列表
        chatListView.setCellFactory(param -> new ChatListCell());
        chatListView.getSelectionModel().selectedItemProperty().addListener(
                (observable, oldValue, newValue) -> loadChat(newValue)
        );

        // 创建初始会话
        createNewChat();

        // 配置快捷键
        setupShortcuts();

        chatContainer.heightProperty().addListener((observable, oldValue, newValue) -> {
            if (chatScrollPane.getVmax() > 0) {
                chatScrollPane.setVvalue(1.0);
            }
        });
    }

    private void setupShortcuts() {
        // 发送消息 (Ctrl/Cmd + Enter)
        KeyCombination sendCombination = new KeyCodeCombination(KeyCode.ENTER,
                KeyCombination.SHORTCUT_DOWN);
        inputArea.setOnKeyPressed(event -> {
            if (sendCombination.match(event)) {
                event.consume();
                sendMessage();
            }
        });

        // 新建会话 (Ctrl/Cmd + N)
        shortcutService.addShortcut(
                new KeyCodeCombination(KeyCode.N, KeyCombination.SHORTCUT_DOWN),
                this::createNewChat
        );

        // 切换主题 (Ctrl/Cmd + T)
        shortcutService.addShortcut(
                new KeyCodeCombination(KeyCode.T, KeyCombination.SHORTCUT_DOWN),
                this::toggleTheme
        );

        // 清空当前会话 (Ctrl/Cmd + L)
        shortcutService.addShortcut(
                new KeyCodeCombination(KeyCode.L, KeyCombination.SHORTCUT_DOWN),
                this::clearCurrentChat
        );

        // 复制选中文本 (Ctrl/Cmd + C)
        shortcutService.addShortcut(
                new KeyCodeCombination(KeyCode.C, KeyCombination.SHORTCUT_DOWN),
                this::copySelectedText
        );
    }

    @FXML
    private void createNewChat() {
        Chat newChat = chatService.createNewChat();
        chatListView.getItems().add(newChat);
        chatListView.getSelectionModel().select(newChat);
        currentChat = newChat;
    }

    @FXML
    private void sendMessage() {
        String messageText = inputArea.getText();
        if (messageText.trim().isEmpty()) return;

        // 禁用输入框和发送按钮
        inputArea.setDisable(true);

        Message userMessage = new Message("user", messageText);
        currentChat.addMessage(userMessage);
        addMessageToChat(userMessage);

        // 清空输入框
        inputArea.clear();

        // 创建助手消息
        Message assistantMessage = new Message("assistant", "");
        currentChat.addMessage(assistantMessage);
        HBox messageBox = addMessageToChat(assistantMessage);
        TextFlow textFlow = (TextFlow) messageBox.getChildren().get(0);

        StringBuilder contentBuilder = new StringBuilder();

        // 发送到服务器并获取流式响应
        chatService.sendMessageStream(
                messageText,
                modelSelector.getValue(),
                chunk -> {
                    // 在JavaFX线程中更新UI
                    contentBuilder.append(chunk);
                    javafx.application.Platform.runLater(() -> {
                        assistantMessage.setContent(contentBuilder.toString());
                        updateMessage(textFlow, chunk);
                    });
                },
                () -> {
                    // 完成后在JavaFX线程中更新UI
                    javafx.application.Platform.runLater(() -> {
                        inputArea.setDisable(false);
                        inputArea.requestFocus();
                    });
                },
                currentChat
        );
    }

    private HBox addMessageToChat(Message message) {

        TextFlow textFlow = new TextFlow();
        textFlow.prefWidthProperty().bind(chatContainer.widthProperty());

        textFlow.setLineSpacing(15);

        Text text = new Text(message.getContent());
        textFlow.getChildren().add(text);

        // 设置消息样式
        HBox messageBox = new HBox(textFlow);
        messageBox.prefWidthProperty().bind(chatContainer.widthProperty());
        messageBox.getStyleClass().add(message.getRole() + "-message");
        messageBox.setPadding(new javafx.geometry.Insets(10));

        if ("user".equals(message.getRole())) {
            textFlow.setTextAlignment(TextAlignment.RIGHT);
            messageBox.setAlignment(Pos.CENTER_RIGHT);
        } else {
            messageBox.setAlignment(Pos.CENTER_LEFT);
            textFlow.setTextAlignment(TextAlignment.LEFT);
        }

        chatContainer.getChildren().add(messageBox);
        return messageBox;
    }

    private void updateMessage(TextFlow textFlow, String content) {

        if (content.contains("<think>")) {
            content = content.split("</think>")[1];
            if (content.equals("")){
                return;
            }
        }

        if (content.contains("</think>")) {
            content = content.split("</think>")[1];
        }


        Text text = new Text(content);
        Text spacer = new Text("\r\n");
        textFlow.getChildren().add(text);
        textFlow.getChildren().add(spacer);

    }

    private void loadChat(Chat chat) {
        if (chat == null) return;
        currentChat = chat;
        chatContainer.getChildren().clear();
        chat.getMessages().forEach(this::addMessageToChat);
    }

    @FXML
    private void toggleTheme() {
        if (themeService == null) return;
        themeService.toggleTheme();
        markdownService.setDarkTheme(themeService.isDarkTheme());
        reloadMessages();
    }

    private void reloadMessages() {
        VBox oldContainer = chatContainer;
        chatContainer = new VBox(10);
        oldContainer.getChildren().clear();
        loadChat(currentChat);
    }

    @FXML
    private void clearCurrentChat() {
        if (currentChat != null) {
            currentChat.getMessages().clear();
            chatContainer.getChildren().clear();
        }
    }

    private void copySelectedText() {
        if (scene == null) return;

        // 获取当前焦点节点
        Node focusedNode = scene.getFocusOwner();
        String selectedText = null;

        // 根据不同的节点类型获取选中的文本
        if (focusedNode instanceof TextArea) {
            TextArea textArea = (TextArea) focusedNode;
            selectedText = textArea.getSelectedText();
        } else if (focusedNode instanceof ScrollPane) {
            // 如果是消息容器，获取其中的文本内容
            ScrollPane scrollPane = (ScrollPane) focusedNode;
            Node content = scrollPane.getContent();
            if (content instanceof TextFlow) {
                TextFlow textFlow = (TextFlow) content;
                // 获取文本内容
                for (Node node : textFlow.getChildren()) {
                    if (node instanceof Text) {
                        Text text = (Text) node;
                        selectedText = text.getText();
                        break;
                    }
                }
            }
        }

        // 如果有文本，复制到剪贴板
        if (selectedText != null && !selectedText.isEmpty()) {
            final Clipboard clipboard = Clipboard.getSystemClipboard();
            final ClipboardContent content = new ClipboardContent();
            content.putString(selectedText);
            clipboard.setContent(content);
        }
    }

    public void setScene(Scene scene) {
        if (scene == null) return;
        this.scene = scene;
        themeService.setScene(scene);
        shortcutService.setScene(scene);
    }
}