package com.axis.innovators.box.gui;

import com.axis.innovators.box.register.LanguageManager;
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.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.LinkedList;
import java.io.IOException;

import org.jsoup.Jsoup;
import org.jsoup.nodes.Document;
import org.jsoup.nodes.Element;
import org.jsoup.select.Elements;

/**
 * 本地AI执行工具
 * @author tzdwindows 7
 */
public class LocalWindow extends WindowsJDialog {
    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;
    private boolean searchEnabled = false;

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

    @Override
    public void initUI() {
        super.initUI();
        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(language.getText("localWindow.newBtn"));
        styleButton(newBtn);
        newBtn.addActionListener(e -> createNewContext());

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

        // CUDA切换
        JCheckBox cudaCheck = new JCheckBox(language.getText("localWindow.cudaCheck"), 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);

        // 搜索功能按钮
        JButton searchBtn = new JButton(language.getText("localWindow.searchBtn"));
        styleButton(searchBtn);
        searchBtn.addActionListener(e -> {
            searchEnabled = !searchEnabled;
            searchBtn.setText(searchEnabled ? language.getText("localWindow.searchBtn.1")
                    : language.getText("localWindow.searchBtn"));
            if (searchEnabled) {
                appendMessage(language.getText("localWindow.searchBtn.appendMessage"));
            } else {
                appendMessage(language.getText("localWindow.searchBtn.appendMessage.1"));
            }
        });

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

        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(language.getText("localWindow.sendButton"));
        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("");
            if (searchEnabled) {
                // 如果搜索功能开启，先进行搜索
                new SwingWorker<String, Void>() {
                    @Override
                    protected String doInBackground() {
                        return searchWeb(prompt);
                    }

                    @Override
                    protected void done() {
                        try {
                            String searchResult = get();
                            logger.info("Search result: {}", searchResult);
                            if (searchResult.isEmpty()) {
                                while (true){
                                    searchResult = searchWeb(prompt);
                                    if (!searchResult.isEmpty()) {
                                        break;
                                    }
                                }
                                seekAI(prompt, searchResult);
                                return;
                            }
                            seekAI(prompt, searchResult);
                        } catch (Exception ex) {
                            logger.error("Search failed: ", ex);
                        }
                    }
                }.execute();
            } else {
                analyzeWithAI(prompt, null);
            }
        }
    }

    private static String searchWeb(String query) {
        int maxRetries = 2;
        int retry = 0;

        while (true) {
            try {
                // 设置请求 URL
                String url = "https://cn.bing.com/search?q=" + URLEncoder.encode(query, StandardCharsets.UTF_8);

                // 设置请求头和 Cookies
                Document doc = Jsoup.connect(url)
                        .userAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/107.0.0.0 Safari/537.36")
                        .header("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8")
                        .header("Accept-Language", "zh-CN,zh;q=0.9")
                        .header("Accept-Encoding", "gzip, deflate, br")
                        .header("Connection", "keep-alive")
                        .cookie("MUID", "你的MUID值")
                        .cookie("SRCHD", "你的SRCHD值")
                        .timeout(10000)
                        .get();

                StringBuilder result = new StringBuilder();
                Elements items = doc.select("li.b_algo");

                for (Element item : items) {
                    String title = item.select("h2").text();
                    String urlLink = item.select("a").attr("href");
                    String snippet = item.select("div.b_caption p").text();
                    String urlContent = "未爬取";

                    result.append("标题: ").append(title)
                            .append("\n链接: ").append(urlLink)
                            .append("\n摘要: ").append(snippet)
                            .append("\n爬取到的内容: ").append(urlContent)
                            .append("\n\n");
                }

                return result.toString();
            } catch (IOException e) {
                retry++;
                if (retry > maxRetries) {
                    System.err.println("搜索失败（重试 " + maxRetries + " 次）: " + e.getMessage());
                    return "搜索失败: " + e.getMessage();
                }
            }
        }
    }

    private void analyzeWithAI(String prompt, String searchResult) {
        new SwingWorker<String, String>() {
            @Override
            protected String doInBackground() {
                float temperature = ((Number) tempSpinner.getValue()).floatValue();
                String fullPrompt = searchResult != null ?  ("\n\n搜索结果:\n" + searchResult + "。") + prompt : prompt;
                return LM.inference(
                        currentModelHandle,
                        contextHandles.getLast(),
                        temperature,
                        fullPrompt,
                        language.getText("localWindow.prompt.system"),
                        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();
    }

    private void seekAI(String prompt, String searchResult) {
        new SwingWorker<String, String>() {
            @Override
            protected String doInBackground() {
                float temperature = ((Number) tempSpinner.getValue()).floatValue();
                String fullPrompt = searchResult != null ?  ("\n\n搜索结果:\n" + searchResult + "。") + prompt : prompt;
                return LM.inference(
                        currentModelHandle,
                        contextHandles.getLast(),
                        temperature,
                        fullPrompt,
                        language.getText("localWindow.prompt.system.2"),
                        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"+ language.getText("localWindow.prompt.user")+"\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();
    }
}