package com.axis.innovators.box.gui;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.tzd.lm.LM;
import javax.swing.*;
import javax.swing.border.EmptyBorder;
import java.awt.*;
import java.awt.event.ActionEvent;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
 * 本地AI执行工具
 * @author tzdwindows 7
 */
public class LocalWindow extends JDialog {
    private static final Logger logger = LogManager.getLogger(LocalWindow.class);
    private final Color bgColor = new Color(45, 45, 48);
    private final Color textColor = new Color(240, 240, 240);
    private final Font mainFont = new Font("微软雅黑", Font.PLAIN, 14);
    private final Color bgColor1 = new Color(235, 241, 250);
    private final Color bgColor2 = new Color(255, 255, 255);

    private JTextArea chatArea;
    private JTextField inputField;
    private JProgressBar progressBar;
    private JComboBox<String> contextBox;
    private final LinkedList<Long> contextHandles = new LinkedList<>();
    private long currentModelHandle = -1;

    private JSpinner tempSpinner;

    public LocalWindow(Window owner) {
        super(owner, "本地AI执行工具", ModalityType.APPLICATION_MODAL);
        initializeUI();
        loadModelAsync();
    }

    private void initializeUI() {
        JPanel mainPanel = getjPanel();

        // 聊天区域
        chatArea = new JTextArea();
        chatArea.setEditable(false);
        chatArea.setFont(mainFont);
        chatArea.setForeground(textColor);
        chatArea.setOpaque(false);
        chatArea.setLineWrap(true);
        chatArea.setWrapStyleWord(true);

        chatArea.setBackground(Color.lightGray);
        chatArea.setForeground(Color.black);

        JScrollPane scrollPane = new JScrollPane(chatArea);
        scrollPane.setOpaque(false);
        scrollPane.getViewport().setOpaque(false);
        scrollPane.setBorder(BorderFactory.createEmptyBorder());

        JToolBar toolBar = createToolBar();
        JPanel inputPanel = createInputPanel();

        progressBar = new JProgressBar();
        progressBar.setIndeterminate(true);
        progressBar.setVisible(false);

        mainPanel.add(toolBar, BorderLayout.NORTH);
        mainPanel.add(scrollPane, BorderLayout.CENTER);
        mainPanel.add(inputPanel, BorderLayout.SOUTH);

        mainPanel.add(scrollPane, BorderLayout.CENTER);
        this.add(mainPanel);
        this.setSize(800, 600);
        this.setLocationRelativeTo(getOwner());
    }

    private JPanel getjPanel() {
        JPanel mainPanel = new JPanel(new BorderLayout()) {
            @Override
            protected void paintComponent(Graphics g) {
                super.paintComponent(g);
                Graphics2D g2d = (Graphics2D) g;
                g2d.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
                GradientPaint gp = new GradientPaint(0, 0, bgColor1, getWidth(), getHeight(), bgColor2);
                g2d.setPaint(gp);
                g2d.fillRect(0, 0, getWidth(), getHeight());
            }
        };
        mainPanel.setBorder(new EmptyBorder(10, 10, 10, 10));
        return mainPanel;
    }

    private JToolBar createToolBar() {
        JToolBar toolBar = new JToolBar();
        toolBar.setFloatable(false);
        toolBar.setBackground(bgColor);

        // 新对话按钮
        JButton newBtn = new JButton("新对话");
        styleButton(newBtn);
        newBtn.addActionListener(e -> createNewContext());

        // 保存记录按钮
        JButton saveBtn = new JButton("保存记录");
        styleButton(saveBtn);
        saveBtn.addActionListener(e -> saveConversation());

        // CUDA切换
        JCheckBox cudaCheck = new JCheckBox("启用CUDA", LM.CUDA);
        cudaCheck.setFont(mainFont);
        cudaCheck.setBackground(Color.lightGray);
        cudaCheck.setForeground(Color.black);
        cudaCheck.addActionListener(e -> LM.CUDA = cudaCheck.isSelected());

        // 上下文选择
        contextBox = new JComboBox<>();
        contextBox.setBackground(bgColor);
        contextBox.setForeground(textColor);

        toolBar.add(newBtn);
        toolBar.add(saveBtn);
        toolBar.addSeparator();
        toolBar.add(cudaCheck);
        toolBar.addSeparator();

        toolBar.add(contextBox);

        return toolBar;
    }

    private JPanel createInputPanel() {
        JPanel panel = new JPanel(new BorderLayout());
        panel.setBackground(bgColor);

        inputField = new JTextField();
        inputField.setFont(mainFont);
        inputField.setForeground(Color.black);
        inputField.setBackground(Color.white);

        JButton sendButton = new JButton("发送");
        styleButton(sendButton);
        sendButton.addActionListener(this::handleSendMessage);

        tempSpinner = new JSpinner(new SpinnerNumberModel(0.8, 0.0, 1.0, 0.1));
        tempSpinner.setFont(mainFont);
        tempSpinner.setBackground(bgColor);
        tempSpinner.setForeground(textColor);
        tempSpinner.setPreferredSize(new Dimension(80, tempSpinner.getPreferredSize().height));

        JPanel rightPanel = new JPanel(new BorderLayout());
        rightPanel.add(tempSpinner, BorderLayout.WEST);
        rightPanel.add(sendButton, BorderLayout.EAST);

        panel.add(inputField, BorderLayout.CENTER);
        panel.add(rightPanel, BorderLayout.EAST);
        return panel;
    }

    private void styleButton(JButton button) {
        button.setFont(mainFont);
        button.setBackground(Color.lightGray);
        button.setForeground(Color.black);
        button.setFocusPainted(false);
        button.setBorder(BorderFactory.createEmptyBorder(5, 15, 5, 15));
    }

    private void loadModelAsync() {
        progressBar.setVisible(true);
        new SwingWorker<Long, Void>() {
            @Override
            protected Long doInBackground() {
                return LM.llamaLoadModelFromFile(LM.DEEP_SEEK);
            }

            @Override
            protected void done() {
                try {
                    currentModelHandle = get();
                    createNewContext();
                } catch (Exception ex) {
                    logger.error("Model loading failed , loading location {}", LM.DEEP_SEEK ,ex);
                    JOptionPane.showMessageDialog(LocalWindow.this, "模型加载失败");
                }
                progressBar.setVisible(false);
            }
        }.execute();
    }

    private void createNewContext() {
        if (currentModelHandle != -1) {
            long newCtx = LM.createContext(currentModelHandle);
            contextHandles.add(newCtx);
            contextBox.addItem("上下文 " + contextHandles.size());
            contextBox.setSelectedIndex(contextHandles.size()-1);
            contextBox.setBackground(Color.lightGray);
            contextBox.setForeground(Color.black);
            chatArea.setText("");
        }
    }

    private void handleSendMessage(ActionEvent e) {
        String prompt = inputField.getText();
        if (!prompt.isEmpty()) {
            appendMessage(prompt);
            inputField.setText("");
            new SwingWorker<String, String>() {
                @Override
                protected String doInBackground() {
                    float temperature = ((Number) tempSpinner.getValue()).floatValue();
                    return LM.inference(
                            currentModelHandle,
                            contextHandles.getLast(),
                            temperature,
                            prompt,
                            "你是一个乐于助人的AI助手，请用友好自然的语气回答用户问题。",
                            this::publish
                    );
                }

                @Override
                protected void process(java.util.List<String> chunks) {
                    chunks.forEach(msg -> appendAIMessage(msg));
                }

                @Override
                protected void done() {
                    try {
                        String result = get();
                        logger.info(result);
                    } catch (Exception ex) {
                        logger.error("SwingWorker execution error: ", ex);
                    }
                }
            }.execute();
        }
    }

    static boolean isTemporary = true;
    private void appendAIMessage(String message) {
        SwingUtilities.invokeLater(() -> {
            if (isTemporary) {
                chatArea.append("【AI】\n");
                logger.info("【AI】");
                isTemporary = false;
            }
            // 处理Markdown语法
            String message2 = message.replaceAll("\\*\\*(.*?)\\*\\*", "<b>$1</b>")
                    .replaceAll("##(.*?)\\n", "<h2>$1</h2>")
                    .replaceAll("</think>", "</think></fold>");
            chatArea.append(message2);
            chatArea.setCaretPosition(chatArea.getDocument().getLength());
        });
    }

    private void appendMessage(String message) {
        SwingUtilities.invokeLater(() -> {
            logger.info("\n【用户】\n{}", message);
            chatArea.append("\n【用户】\n" + message + "\n\n");
            chatArea.setCaretPosition(chatArea.getDocument().getLength());
        });
    }

    private void saveConversation() {
        //String content = chatArea.getText();
        // TODO: 调用保存实现
    }

    @Override
    public void dispose() {
        if (contextHandles.isEmpty()){
            logger.warn("No context to free");
        } else {
            contextHandles.forEach(LM::llamaFreeContext);
        }
        if (currentModelHandle != -1) {
            LM.llamaFreeModel(currentModelHandle);
        } else {
            logger.warn("Wrong model handle");
        }
        super.dispose();
    }
}