package com.axis.innovators.box.register;

import com.axis.innovators.box.AxisInnovatorsBox;
import com.axis.innovators.box.gui.FridaWindow;
import com.axis.innovators.box.gui.LocalWindow;
import com.axis.innovators.box.gui.MainWindow;
import com.axis.innovators.box.plugins.PluginDescriptor;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.tzd.lm.LM;

import javax.swing.*;
import java.awt.*;
import java.awt.event.ActionEvent;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;

/**
 * 负责注册Registration项（限定时机，在执行.run之前可设置）
 * @author tzdwindows 7
 */
public class RegistrationTool {
    private static final Logger logger = LogManager.getLogger(RegistrationTool.class);
    private final List<MainWindow.ToolCategory> toolCategories = new ArrayList<>();
    private final AxisInnovatorsBox main;
    private final List<UUID> uuidList = new ArrayList<>();
    private final List<String> registeredNameList = new ArrayList<>();
    private final List<PluginDescriptor> pluginDescriptors = new ArrayList<>();

    public RegistrationTool(AxisInnovatorsBox main) {
        this.main = main;
        int id = 0;
        MainWindow.ToolCategory debugCategory = new MainWindow.ToolCategory("调试工具",
                "debug/debug.png",
                "用于调试指定Windows工具的一个分类");

        debugCategory.addTool(new MainWindow.ToolItem("Frida注入工具", "debug/frida/frida_main.png",
                "使用frida注入目标进程的脚本程序 " +
                        "\n作者：tzdwindows 7", ++id, new AbstractAction() {
            @Override
            public void actionPerformed(ActionEvent e) {
                Window owner = SwingUtilities.windowForComponent((Component) e.getSource());
                FridaWindow fridaWindow = new FridaWindow(owner);
                main.popupWindow(fridaWindow);
            }
        }));

        MainWindow.ToolCategory aICategory = new MainWindow.ToolCategory("AI工具",
                "ai/ai.png",
                "人工智能/大语言模型");

        aICategory.addTool(new MainWindow.ToolItem("本地AI执行工具", "ai/local/local_main.png",
                "在本机对开源大语言模型进行推理" +
                        "\n作者：tzdwindows 7", ++id, new AbstractAction() {
            @Override
            public void actionPerformed(ActionEvent e) {

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

                Window owner = SwingUtilities.windowForComponent((Component) e.getSource());
                LocalWindow dialog = new LocalWindow(owner);
                main.popupWindow(dialog);
            }
        }));

        addToolCategory(debugCategory, "system:debugTools");
        addToolCategory(aICategory,"system:fridaTools");
    }

    /**
     * 注册ToolCategory
     * @param toolCategory ToolCategory
     */
    public boolean addToolCategory(MainWindow.ToolCategory toolCategory,
                                String registeredName) {
        if (!main.isWindow()) {
            if (registeredNameList.contains(registeredName)) {
                throw new RegistrationError(registeredName + " duplicate registered names");
            }
            uuidList.add(toolCategory.getId());
            registeredNameList.add(registeredName);
            toolCategories.add(toolCategory);
            return true;
        } else {
            logger.warn("Wrong time to add tools");
            return false;
        }
    }

    /**
     * 注册ToolCategory
     * @param toolCategory ToolCategory
     */
    public void addToolCategory(MainWindow.ToolCategory toolCategory,
                                PluginDescriptor pluginDescriptor,
                                String registeredName) {
        if (addToolCategory(toolCategory,
                pluginDescriptor.getRegistrationName() + ":" +
                        registeredName)) {
            pluginDescriptors.add(pluginDescriptor);
        }
    }

    /**
     * 根据UUID获取ToolCategory
     * @param id UUID
     * @return ToolCategory对象
     */
    public MainWindow.ToolCategory getToolCategory(UUID id) {
        if (!main.isWindow()) {
            for (MainWindow.ToolCategory toolCategory : toolCategories) {
                if (toolCategory.getId().equals(id)) {
                    return toolCategory;
                }
            }
        } else {
            logger.warn("Wrong time to operation tools");
        }
        return null;
    }

    /**
     * 根据注册名获取UUID
     * @param registeredName 注册名
     * @return UUID
     */
    public UUID getUUID(String registeredName){
        if (!main.isWindow()) {
            for (int i = 0; i < registeredNameList.size(); i++) {
                if (registeredNameList.get(i).equals(registeredName)) {
                    return uuidList.get(i);
                }
            }
        } else {
            logger.warn("Wrong time to operation tools");
        }
        return null;
    }

    public List<MainWindow.ToolCategory> getToolCategories() {
        return toolCategories;
    }
}
