package com.axis.innovators.box.gui;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.tzd.frida.windows.Frida;

import javax.swing.*;
import javax.swing.border.EmptyBorder;
import javax.swing.event.DocumentEvent;
import javax.swing.event.DocumentListener;
import javax.swing.table.DefaultTableModel;
import java.awt.*;
import java.awt.event.ActionEvent;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;

/**
 * Frida注入工具窗口
 * @author tzdwindows 7
 */
public class FridaWindow extends JDialog {
    private static final Logger logger = LogManager.getLogger(FridaWindow.class);
    private JTextArea scriptArea;
    private JTextArea logArea;
    private JTextField pidField;
    private boolean isRepetition = false;

    public FridaWindow(Window owner) {
        super(owner, "Frida 注入工具", ModalityType.APPLICATION_MODAL);
        initializeUI();
    }

    private void initializeUI() {
        setSize(800, 600);
        setLocationRelativeTo(getOwner());
        setDefaultCloseOperation(DISPOSE_ON_CLOSE);

        JPanel mainPanel = new JPanel() {
            @Override
            protected void paintComponent(Graphics g) {
                super.paintComponent(g);
                Graphics2D g2d = (Graphics2D) g;
                g2d.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
                Color color1 = new Color(235, 241, 250);
                Color color2 = new Color(255, 255, 255);
                GradientPaint gp = new GradientPaint(0, 0, color1, getWidth(), getHeight(), color2);
                g2d.setPaint(gp);
                g2d.fillRect(0, 0, getWidth(), getHeight());
            }
        };
        mainPanel.setLayout(new BorderLayout(10, 10));
        mainPanel.setBorder(new EmptyBorder(15, 15, 15, 15));

        // 输入面板
        JPanel inputPanel = createInputPanel();
        mainPanel.add(inputPanel, BorderLayout.NORTH);

        // 脚本编辑区
        JPanel scriptPanel = createScriptPanel();
        mainPanel.add(scriptPanel, BorderLayout.CENTER);

        // 日志输出
        JPanel logPanel = createLogPanel();
        mainPanel.add(logPanel, BorderLayout.SOUTH);

        setContentPane(mainPanel);
    }

    private JPanel createInputPanel() {
        JPanel panel = new JPanel(new GridLayout(1, 3, 10, 10));
        panel.setOpaque(false);

        pidField = new JTextField("请输入进程PID");
        pidField.setFont(new Font("微软雅黑", Font.PLAIN, 14));
        pidField.setBorder(BorderFactory.createCompoundBorder(
                BorderFactory.createLineBorder(new Color(57, 56, 56)),
                BorderFactory.createEmptyBorder(5, 10, 5, 10)
        ));

        //JButton attachButton = new JButton("附加进程");
        //styleButton(attachButton, new Color(70, 130, 180));
        //attachButton.setEnabled(false);

        JButton browseButton = new JButton("选择进程");
        styleButton(browseButton, new Color(60, 179, 113));
        browseButton.addActionListener(this::openProcessSelectionWindow);
        panel.add(pidField);
        //panel.add(attachButton);
        panel.add(browseButton);

        return panel;
    }

    private JPanel createScriptPanel() {
        JPanel panel = new JPanel(new BorderLayout());
        panel.setOpaque(false);
        panel.setBorder(BorderFactory.createTitledBorder("脚本编辑器"));

        scriptArea = new JTextArea();
        scriptArea.setFont(new Font("Consolas", Font.PLAIN, 14));
        JScrollPane scrollPane = new JScrollPane(scriptArea);
        scrollPane.setBorder(BorderFactory.createLineBorder(new Color(0, 0, 0)));

        JPanel buttonPanel = new JPanel(new FlowLayout(FlowLayout.RIGHT, 10, 5));
        buttonPanel.setOpaque(false);

        JButton injectButton = new JButton("注入脚本");
        styleButton(injectButton, new Color(220, 20, 60));
        injectButton.addActionListener(this::handleInject);

        buttonPanel.add(injectButton);

        panel.add(scrollPane, BorderLayout.CENTER);
        panel.add(buttonPanel, BorderLayout.SOUTH);

        return panel;
    }

    private JPanel createLogPanel() {
        JPanel panel = new JPanel(new BorderLayout());
        panel.setPreferredSize(new Dimension(0, 150));
        panel.setBorder(BorderFactory.createTitledBorder("日志输出"));
        panel.setOpaque(false);

        logArea = new JTextArea();
        logArea.setEditable(false);
        logArea.setFont(new Font("微软雅黑", Font.PLAIN, 12));
        JScrollPane scrollPane = new JScrollPane(logArea);
        scrollPane.setBorder(BorderFactory.createLineBorder(new Color(0, 0, 0)));

        panel.add(scrollPane);
        return panel;
    }

    private void styleButton(JButton button, Color bgColor) {
        button.setFont(new Font("微软雅黑", Font.BOLD, 14));
        button.setFocusPainted(false);
        button.setBackground(bgColor);
        button.setForeground(Color.black);
        button.setBorder(BorderFactory.createEmptyBorder(8, 20, 8, 20));
    }

    private void handleInject(ActionEvent e) {
        try {
            long pid = Long.parseLong(pidField.getText());
            String script = scriptArea.getText();

            Frida frida = new Frida(script, pid);
            if (!isRepetition) {
                try {
                    frida.run(() -> {
                    }).execute(frida1 -> frida1.addCallbackMessage(message ->
                            SwingUtilities.invokeLater(() -> {
                                logger.info("[LOG] {}", message);
                                logArea.append("[LOG] " + message + "\n");
                            })
                    )).start();
                } catch (Exception ex) {
                    logger.error("Injection failure", ex);
                    throw new RuntimeException(ex);
                }
                isRepetition = true;
            } else {
                try {
                    frida.run(() -> {}).execute(frida12 -> {}).start();
                } catch (Exception ex) {
                    logger.error("Injection failure", ex);
                }
            }

        } catch (NumberFormatException ex) {
            logger.error("无效的进程ID", ex);
            JOptionPane.showMessageDialog(this, "无效的进程ID", "错误", JOptionPane.ERROR_MESSAGE);
        }
    }

    private void openProcessSelectionWindow(ActionEvent e) {
        ProcessSelectionWindow selectionWindow = new ProcessSelectionWindow(this);
        selectionWindow.setVisible(true);
    }

    private static class ProcessSelectionWindow extends JDialog {
        private List<ProcessInfo> processList;
        private JTable table;

        public ProcessSelectionWindow(Window owner) {
            super(owner, "选择进程", ModalityType.APPLICATION_MODAL);
            setSize(600, 400);
            setLocationRelativeTo(owner);
            initializeUI();
        }

        private void initializeUI() {
            JPanel panel = new JPanel();
            panel.setLayout(new BorderLayout());

            // 获取所有进程信息
            processList = getProcesses();

            // 创建搜索框
            JPanel searchPanel = new JPanel();
            searchPanel.setLayout(new BorderLayout());
            JTextField searchField = new JTextField();
            searchField.getDocument().addDocumentListener(new DocumentListener() {
                @Override
                public void insertUpdate(DocumentEvent e) {
                    filterProcesses(searchField.getText());
                }

                @Override
                public void removeUpdate(DocumentEvent e) {
                    filterProcesses(searchField.getText());
                }

                @Override
                public void changedUpdate(DocumentEvent e) {
                    filterProcesses(searchField.getText());
                }
            });
            searchPanel.add(new JLabel("搜索进程: "), BorderLayout.WEST);
            searchPanel.add(searchField, BorderLayout.CENTER);
            panel.add(searchPanel, BorderLayout.NORTH);

            String[] columns = {"进程名称", "进程ID"};
            Object[][] data = getTableData(processList);
            table = new JTable(data, columns);
            JScrollPane scrollPane = new JScrollPane(table);
            panel.add(scrollPane, BorderLayout.CENTER);

            JButton selectButton = new JButton("选择");
            selectButton.addActionListener(e -> {
                int selectedRow = table.getSelectedRow();
                if (selectedRow >= 0) {
                    long selectedPid = (Long) table.getValueAt(selectedRow, 1); // 注意这里获取的是PID列
                    ((FridaWindow) getOwner()).pidField.setText(String.valueOf(selectedPid));
                    dispose();
                }
            });
            panel.add(selectButton, BorderLayout.SOUTH);

            setContentPane(panel);
        }

        private void filterProcesses(String query) {
            List<ProcessInfo> filteredList = new ArrayList<>();
            for (ProcessInfo process : processList) {
                if (process.name().toLowerCase().contains(query.toLowerCase())) {
                    filteredList.add(process);
                }
            }

            Object[][] filteredData = getTableData(filteredList);
            table.setModel(new DefaultTableModel(filteredData, new String[]{"进程名称", "进程ID"}));
        }

        private Object[][] getTableData(List<ProcessInfo> processes) {
            Object[][] data = new Object[processes.size()][2];
            for (int i = 0; i < processes.size(); i++) {
                ProcessInfo process = processes.get(i);
                data[i] = new Object[]{process.name(), process.pid()};
            }
            return data;
        }

        /**
         * 执行cmd命令获取所有进程信息
         * @return 进程信息列表
         */
        private List<ProcessInfo> getProcesses() {
            List<ProcessInfo> processList = new ArrayList<>();
            try {
                Process process = Runtime.getRuntime().exec("tasklist /fo csv /nh");
                BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
                String line;

                while ((line = reader.readLine()) != null) {
                    String[] columns = line.split(",");
                    if (columns.length > 1) {
                        long pid = Long.parseLong(columns[1].replaceAll("\"", "").trim());
                        String processName = columns[0].replaceAll("\"", "").trim();
                        processList.add(new ProcessInfo(pid, processName));
                    }
                }
                reader.close();
            } catch (IOException e) {
                logger.error("Failed to get process information ", e);
            }
            return processList;
        }
    }

    private record ProcessInfo(long pid, String name) {}
}
