package com.axis.innovators.box.ui;

import com.vladsch.flexmark.util.ast.Node;
import org.tzd.lm.LM;
import com.vladsch.flexmark.html.HtmlRenderer;
import com.vladsch.flexmark.parser.Parser;
import com.vladsch.flexmark.util.data.MutableDataSet;

import javax.swing.*;
import javax.swing.border.AbstractBorder;
import javax.swing.text.html.HTMLDocument;
import javax.swing.text.html.HTMLEditorKit;
import javax.swing.text.html.StyleSheet;
import java.awt.*;
import java.awt.datatransfer.Clipboard;
import java.awt.datatransfer.StringSelection;
import java.awt.event.*;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 *
 * @author tzdwindows 7
 */
public class AIChatDialog extends JFrame {
    private static final Font UI_FONT = createSystemFont();
    private static final Color TEXT_COLOR = new Color(0xF0F0F0);
    private static final int COLLAPSED_HEIGHT = 100;
    private static final int EXPANDED_HEIGHT = 400;
    private final JTextPane outputPane = new JTextPane();
    private final JTextField inputField = new JTextField();
    private final Timer resizeTimer;
    private final StringBuilder currentResponse = new StringBuilder();
    private final HtmlRenderer htmlRenderer;
    private final Parser markdownParser;
    private long modelHandle;
    private long ctxHandle;
    private JScrollPane scrollPane;
    private final AtomicBoolean isRendering = new AtomicBoolean(false);
    private final Timer borderAnimationTimer;
    private float hue = 0f;
    private boolean isUserResized = false;
    private boolean isProcessing = false;
    private final AnimatedBorder inputAnimatedBorder;
    private final AnimatedBorder outputAnimatedBorder;
    private boolean borderActive = false;
    // 用于窗口拖动的变量
    private Point initialClick;

    static {
        try {
            LM.loadLibrary(LM.CUDA);
        } catch (Exception ex) {
            JOptionPane.showMessageDialog(null, "无法加载AI推理库",
                    "无法加载AI推理库", JOptionPane.ERROR_MESSAGE);
        }
    }

    private static Font createSystemFont() {
        String[] fontNames = {
                "Microsoft YaHei",
                "PingFang SC",
                "Noto Sans CJK SC",
                "SimHei",
                "sans-serif"
        };

        for (String name : fontNames) {
            Font font = new Font(name, Font.PLAIN, 14);
            if (font.getFamily().equals(name)) {
                return font.deriveFont(Font.PLAIN, 14);
            }
        }
        return new Font("sans-serif", Font.PLAIN, 14);
    }

    public AIChatDialog() {
        super("AI助手");
        setAlwaysOnTop(true);

        initResources();

        inputAnimatedBorder = new AnimatedBorder(20, true);
        outputAnimatedBorder = new AnimatedBorder(20, false);

        resizeTimer = new Timer(15, e -> updateWindowSize());
        borderAnimationTimer = new Timer(30, e -> updateBorderAnimation());
        initUI();
        initListeners();

        MutableDataSet options = new MutableDataSet();
        markdownParser = Parser.builder(options).build();
        htmlRenderer = HtmlRenderer.builder(options).build();
    }

    private void updateBorderAnimation() {
        hue = (hue + 0.015f) % 1f;
        inputField.repaint();
        outputPane.repaint();
    }

    private void initResources() {
        modelHandle = LM.llamaLoadModelFromFile(LM.DEEP_SEEK);
        ctxHandle = LM.createContext(modelHandle);
    }

    private void initListeners() {
        // 窗口关闭时释放资源
        addWindowListener(new WindowAdapter() {
            @Override
            public void windowClosed(WindowEvent e) {
                releaseResources();
            }
        });

        addComponentListener(new ComponentAdapter() {
            @Override
            public void componentMoved(ComponentEvent e) {
                repaintImmediately();
            }

            @Override
            public void componentResized(ComponentEvent e) {
                repaintImmediately();
            }
        });

        // ESC键关闭窗口
        getRootPane().registerKeyboardAction(e -> dispose(),
                KeyStroke.getKeyStroke(KeyEvent.VK_ESCAPE, 0),
                JComponent.WHEN_IN_FOCUSED_WINDOW
        );

        // 输入框回车事件
        inputField.addActionListener(e -> processQuery());

        // 窗口拖动功能
        addMouseListener(new MouseAdapter() {
            @Override
            public void mousePressed(MouseEvent e) {
                initialClick = e.getPoint(); // 记录鼠标点击的初始位置
            }
        });

        addComponentListener(new ComponentAdapter() {
            @Override
            public void componentResized(ComponentEvent e) {
                if (!resizeTimer.isRunning()) {
                    isUserResized = true;
                }
            }
        });

        addMouseMotionListener(new MouseMotionAdapter() {
            @Override
            public void mouseDragged(MouseEvent e) {
                // 计算窗口新位置
                int deltaX = e.getX() - initialClick.x;
                int deltaY = e.getY() - initialClick.y;
                setLocation(getX() + deltaX, getY() + deltaY);
            }
        });
    }

    private void repaintImmediately() {
        SwingUtilities.invokeLater(() -> {
            getContentPane().repaint();
            outputPane.repaint();
            inputField.repaint();
        });
    }

    private void initUI() {
        // 窗口基础设置
        //getRootPane().setOpaque(false);
        //setUndecorated(true);
        //setBackground(new Color(0, 0, 0, 0));
        //setShape(new RoundRectangle2D.Double(0, 0, getWidth(), getHeight(), 30, 30));

        // 窗口尺寸和位置配置
        setSize(600, COLLAPSED_HEIGHT);
        setMinimumSize(new Dimension(400, COLLAPSED_HEIGHT));
        setLocationRelativeTo(null);

        // 窗口阴影效果
        JPanel shadowPanel = new JPanel(new BorderLayout()) {
            @Override
            protected void paintComponent(Graphics g) {
                super.paintComponent(g);
                Graphics2D g2 = (Graphics2D) g.create();
                g2.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);

                for (int i = 0; i < 8; i++) {
                    g2.setColor(new Color(0, 0, 0, 30 - i * 3));
                    g2.fillRoundRect(i, i, getWidth() - i * 2, getHeight() - i * 2, 30, 30);
                }
                g2.dispose();
            }
        };

        shadowPanel.setOpaque(false);

        // 主容器
        JPanel mainPanel = new JPanel(new BorderLayout()) {
            @Override
            protected void paintComponent(Graphics g) {
                Graphics2D g2 = (Graphics2D) g.create();
                g2.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);

                g2.setColor(new Color(70, 73, 75, 255));
                g2.fillRoundRect(0, 0, getWidth(), getHeight(), 20, 20);
                g2.dispose();
            }

            @Override
            public Dimension getPreferredSize() {
                return new Dimension(getWidth(), getHeight());
            }
        };
        mainPanel.setOpaque(false);
        mainPanel.setBorder(BorderFactory.createEmptyBorder(20, 25, 20, 25));

        // 输入面板（浏览器搜索框风格）
        JPanel inputPanel = new JPanel(new BorderLayout());
        inputPanel.setOpaque(false);
        inputPanel.setBorder(BorderFactory.createEmptyBorder(0, 0, 8, 0));

        // 搜索图标
        //JLabel searchIcon = new JLabel(LoadIcon.loadIcon("search_icon.png", 32));
        //inputPanel.add(searchIcon, BorderLayout.WEST);

        // 输入框样式
        inputField.setFont(new Font("Microsoft YaHei", Font.PLAIN, 14));
        inputField.setForeground(Color.WHITE);

        inputField.setBorder(BorderFactory.createCompoundBorder(
                inputAnimatedBorder,
                null
        ));

        // 获取字体度量
        FontMetrics metrics = inputField.getFontMetrics(inputField.getFont());
        // 计算最小显示高度 = 字体高度 + 上下边距
        int minHeight = metrics.getHeight() + (5 + 5); // 5+5是EmptyBorder的上下值
        // 设置动态高度
        inputField.setPreferredSize(new Dimension(
                inputField.getWidth(),
                Math.max(minHeight, 36)
        ));

        inputField.putClientProperty("Text.antialias", true);
        inputField.putClientProperty("Text.renderer", "subpixel");

        inputField.putClientProperty("JTextField.placeholderText", "向我提出问题...");
        inputPanel.add(inputField, BorderLayout.CENTER);

        // 在initUI方法开头添加
        UIManager.put("TextPane.font", new Font("JetBrains Mono", Font.PLAIN, 14));

        // 修改输出区域配置
        outputPane.putClientProperty("JTextPane.w3cLengthUnits", true);
        outputPane.putClientProperty("JTextPane.honorDisplayProperties", true);

        // 代码高亮样式增强
        StyleSheet styleSheet = new StyleSheet();

        HTMLEditorKit editorKit = new HTMLEditorKit() {
            @Override
            public void install(JEditorPane c) {
                super.install(c);
            }
        };
        editorKit.setStyleSheet(styleSheet);

        outputPane.setFont(new Font("Microsoft YaHei", Font.PLAIN, 16));

        outputPane.setBorder(BorderFactory.createCompoundBorder(
                outputAnimatedBorder,
                BorderFactory.createEmptyBorder(15, 15, 15, 15)
        ));

        outputPane.setOpaque(false);

        JPopupMenu popupMenu = new JPopupMenu();
        JMenuItem copyItem = new JMenuItem("复制");
        copyItem.addActionListener(e -> {
            String text = outputPane.getSelectedText();
            if (text != null && !text.isEmpty()) {
                copyToClipboard(text);
            }
        });
        popupMenu.add(copyItem);
        outputPane.setComponentPopupMenu(popupMenu);

        outputPane.setEditorKit(editorKit);
        outputPane.setContentType("text/html; charset=UTF-8");

        HTMLDocument doc = (HTMLDocument) editorKit.createDefaultDocument();
        doc.setPreservesUnknownTags(false);

        doc.setPreservesUnknownTags(false);
        outputPane.setDocument(doc);

        scrollPane = new JScrollPane(outputPane);
        scrollPane.setVisible(false);
        scrollPane.setBorder(null);
        scrollPane.getViewport().setOpaque(false);
        scrollPane.setVerticalScrollBarPolicy(ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED);

        outputPane.setMinimumSize(new Dimension(580, 200));
        scrollPane.setMinimumSize(new Dimension(580, 200));

        mainPanel.addComponentListener(new ComponentAdapter() {
            @Override
            public void componentResized(ComponentEvent e) {
                if (getHeight() > COLLAPSED_HEIGHT + 50) {
                    scrollPane.setVisible(true);
                } else {
                    scrollPane.setVisible(false);
                }
            }
        });

        // 布局组装
        mainPanel.add(inputPanel, BorderLayout.NORTH);
        mainPanel.add(scrollPane, BorderLayout.CENTER);
        shadowPanel.add(mainPanel);
        setContentPane(shadowPanel);

    }

    class AnimatedBorder extends AbstractBorder {
        private final int radius;
        private final boolean isInput;
        private boolean active = false;

        public AnimatedBorder(int radius, boolean isInput) {
            this.radius = radius;
            this.isInput = isInput;
        }

        public void setActive(boolean active) {
            this.active = active;
        }

        @Override
        public void paintBorder(Component c, Graphics g, int x, int y, int width, int height) {
            Graphics2D g2 = (Graphics2D) g.create();
            try {
                g2.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
                g2.setRenderingHint(RenderingHints.KEY_STROKE_CONTROL, RenderingHints.VALUE_STROKE_PURE);

                if (active) {
                    Color color1 = Color.getHSBColor(hue, 0.8f, 1f);
                    Color color2 = Color.getHSBColor(hue + 0.3f, 0.8f, 1f);
                    GradientPaint gp = new GradientPaint(0, 0, color1, width, height, color2);
                    g2.setPaint(gp);
                } else {
                    float fixedHue = 0.6f; // 固定的色调，例如蓝色
                    Color color1 = Color.getHSBColor(fixedHue, 0.8f, 1f);
                    Color color2 = Color.getHSBColor(fixedHue + 0.3f, 0.8f, 1f);
                    GradientPaint gp = new GradientPaint(0, 0, color1, width, height, color2);
                    g2.setPaint(gp);
                }

                g2.setStroke(new BasicStroke(isInput ? 3.5f : 3f)); // 加粗线条
                g2.drawRoundRect(x, y, width - 1, height - 1, radius, radius);
            } finally {
                g2.dispose();
            }
        }

        @Override
        public Insets getBorderInsets(Component c) {
            return new Insets(radius/2, radius/2, radius/2, radius/2);
        }

        @Override
        public Insets getBorderInsets(Component c, Insets insets) {
            insets.left = insets.right = radius/2;
            insets.top = insets.bottom = radius/2;
            return insets;
        }
    }

    @Override
    public void paint(Graphics g) {
        Graphics2D g2d = (Graphics2D) g;
        g2d.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
        g2d.setRenderingHint(RenderingHints.KEY_TEXT_ANTIALIASING, RenderingHints.VALUE_TEXT_ANTIALIAS_LCD_HRGB);
        g2d.setRenderingHint(RenderingHints.KEY_FRACTIONALMETRICS, RenderingHints.VALUE_FRACTIONALMETRICS_ON);

        getContentPane().paint(g2d);

        Image buffer = createImage(getWidth(), getHeight());
        Graphics bufferGraphics = buffer.getGraphics();
        super.paint(bufferGraphics);
        g2d.drawImage(buffer, 0, 0, null);
    }

    private void copyToClipboard(String text) {
        StringSelection selection = new StringSelection(text);
        Clipboard clipboard = Toolkit.getDefaultToolkit().getSystemClipboard();
        clipboard.setContents(selection, null);
    }

    private void processQuery() {
        String query = inputField.getText().trim();
        if (query.isEmpty()) return;

        // 修正边框激活状态
        isProcessing = true;
        borderActive = true;
        inputAnimatedBorder.setActive(true);
        outputAnimatedBorder.setActive(true);
        borderAnimationTimer.start();

        // 展开对话框
        if (!resizeTimer.isRunning()) {
            resizeTimer.start();
        }

        // 清空输入
        inputField.setText("");
        appendMessage(query, true);

        // 异步推理
        new SwingWorker<Void, String>() {
            @Override
            protected Void doInBackground() {
                appendMessage("", false);

                String systemPrompt = """
                        你是一个基于DeepSeek-R1的智能助手
                        """;

                LM.inference(modelHandle,
                        ctxHandle,
                        0.7f,
                        query,
                        systemPrompt,
                        this::publishChunk);
                return null;
            }

            private void publishChunk(String chunk) {
                publish(chunk);
            }

            @Override
            protected void process(java.util.List<String> chunks) {
                chunks.forEach(chunk -> {
                    currentResponse.append(chunk);
                    renderMarkdown(currentResponse.toString());
                });
            }

            @Override
            protected void done() {
                inputAnimatedBorder.setActive(false);
                outputAnimatedBorder.setActive(false);
                super.done();
            }
        }.execute();
        if (!shouldExpand) {
            shouldExpand = true;
            resizeTimer.start();
        }
    }

    private void renderMarkdown(String markdown) {
        if (!isRendering.compareAndSet(false, true)) return;

        new SwingWorker<Void, Void>() {
            private String processedHtml;

            @Override
            protected Void doInBackground() {
                try {
                    Node document = markdownParser.parse(markdown);
                    processedHtml = htmlRenderer.render(document);

                    processedHtml = String.format("""
                    <!DOCTYPE html>
                    <html>
                    <head>
                        <script>
                        MathJax = {
                          tex: {
                            inlineMath: [['$', '$'], ['\\(', '\\)']],
                            displayMath: [['$$', '$$'], ['\\[', '\\]']],
                            processEscapes: true
                          }
                        };
                        </script>
                        <script src='https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js'></script>
                        <style>
                            @font-face { font-family: 'LocalFallback'; src: local('%s'); }
                            body { font-family: 'LocalFallback', sans-serif; }
                        </style>
                    </head>
                    <body>%s</body>
                    </html>
                    """, UI_FONT.getFamily(), processedHtml);
                } catch (Exception e) {
                    processedHtml = "<div class='error'>渲染错误：" + e.getMessage() + "</div>";
                }
                return null;
            }

            @Override
            protected void done() {
                try {
                    if (processedHtml == null) return;

                    HTMLDocument doc = (HTMLDocument) outputPane.getDocument();
                    doc.setInnerHTML(doc.getDefaultRootElement(), processedHtml);

                } catch (Exception ex) {
                    ex.printStackTrace();
                } finally {
                    isRendering.set(false);
                }
            }
        }.execute();
    }

    private void appendMessage(String message, boolean isUser) {
        String prefix;
        if (isUser) {
            prefix = "<div style='color:#1a73e8; margin:16px 0;'>You:</div>";
        } else {
            prefix = "<div style='color:#34a853; margin:16px 0;'>AI:</div>";
            message = message + "<think>";
        }

        currentResponse.setLength(0);
        currentResponse.append(prefix).append(message);
        renderMarkdown(currentResponse.toString());
    }

    private boolean shouldExpand = false;
    private void updateWindowSize() {
        if (!SwingUtilities.isEventDispatchThread()) {
            SwingUtilities.invokeLater(this::updateWindowSize);
            return;
        }

        Dimension currentSize = getSize();
        int targetHeight = shouldExpand ? EXPANDED_HEIGHT : COLLAPSED_HEIGHT;

        // 改进的插值算法（使用弹簧物理模型）
        int delta = targetHeight - currentSize.height;
        int acceleration = delta / 15;  // 控制加速度的系数
        int newHeight = currentSize.height + acceleration;

        // 严格边界约束
        if ((delta > 0 && newHeight > targetHeight) ||
                (delta < 0 && newHeight < targetHeight)) {
            newHeight = targetHeight;
        }

        // 当接近目标时直接锁定
        if (Math.abs(delta) <= 2) {
            newHeight = targetHeight;
        }

        // 更新窗口尺寸（保持水平居中）
        Point currentLocation = getLocation();
        int newWidth = Math.max(getWidth(), 400);
        setLocation(
                currentLocation.x - (newWidth - currentSize.width)/2,  // 保持水平居中
                currentLocation.y
        );
        setSize(new Dimension(newWidth, newHeight));

        // 终局状态处理
        if (newHeight == targetHeight) {
            resizeTimer.stop();

            // 强制设置精确尺寸
            setSize(new Dimension(getWidth(), targetHeight));

            // 单次布局更新
            SwingUtilities.invokeLater(() -> {
                scrollPane.setVisible(shouldExpand);
                if (shouldExpand) {
                    scrollPane.getVerticalScrollBar().setValue(0);
                }
            });
        }

        // 优化重绘策略
        if (Math.abs(currentSize.height - newHeight) > 1) {
            getContentPane().repaint();
        }
    }

    private void releaseResources() {
        if (ctxHandle != 0) {
            LM.llamaFreeContext(ctxHandle);
            ctxHandle = 0;
        }
        if (modelHandle != 0) {
            LM.llamaFreeModel(modelHandle);
            modelHandle = 0;
        }
    }

    public void toggleVisibility() {
        if (!isVisible()) {
            setVisible(true);
        }
        updateWindowSize();

        SwingUtilities.invokeLater(() -> {
            setVisible(true);
            setLocationRelativeTo(null);

            // 添加自动聚焦逻辑
            if (inputField.isDisplayable() && inputField.isEnabled()) {
                inputField.requestFocusInWindow();
                // 触发动画边框
                inputAnimatedBorder.setActive(true);
                outputAnimatedBorder.setActive(false);
            }
        });
    }
}