package api.common.util;


import api.common.pojo.dto.GpuDTO;
import api.common.pojo.dto.ProcessDTO;
import lombok.extern.slf4j.Slf4j;
import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.channel.ChannelExec;
import org.apache.sshd.client.channel.ClientChannelEvent;
import org.apache.sshd.client.future.AuthFuture;
import org.apache.sshd.client.session.ClientSession;
import org.dom4j.Document;
import org.dom4j.DocumentException;
import org.dom4j.DocumentHelper;
import org.dom4j.Element;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.*;

/**
 * https://github.com/apache/mina-sshd
 * https://github.com/apache/mina-sshd/blob/master/docs/client-setup.md
 * <!-- ssh -->
 * <dependency>
 * <groupId>org.apache.sshd</groupId>
 * <artifactId>sshd-core</artifactId>
 * <version>${sshd-core.version}</version>
 * </dependency>
 */
@Slf4j
public class SshUtil {

    private static final int port = 22;


    /**
     *
     */
    public static SshClient getClient() {
        return SshClient.setUpDefaultClient();
    }

    /**
     * @param sshClient ssh 客户端对象
     * @param hostname  主机名或 IP
     * @param username  用户名
     * @param password  密码
     * @return 执行结果
     * @throws IOException 异常
     */
    public static ClientSession getSession(SshClient sshClient, String hostname, String username, String password) throws IOException {
        sshClient.start();
        ClientSession clientSession = sshClient.connect(username, hostname, port).verify(3600000).getSession();
        clientSession.addPasswordIdentity(password);
        AuthFuture verify = clientSession.auth().verify(3600000);
        if (verify.isFailure()) {
            throw new RuntimeException("------- ssh 用户名密码验证失败！");
        }
        return clientSession;
    }

    /**
     * @param clientSession 主机名或 IP
     * @throws IOException 异常
     */
    public static void stop(SshClient sshClient, ClientSession clientSession) throws IOException {
        clientSession.close();
        sshClient.stop();
    }


    /**
     * @param session 会话
     * @param command 命令
     * @return 执行结果
     * @throws IOException 异常
     */
    public static String execute(ClientSession session, String command) throws IOException {
        log.info("SshUtil--execute 通过远程链接执行命令：" + command);
        String result;
        ChannelExec execChannel = session.createExecChannel(command);
        // 创建输出流
        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(2048);
        execChannel.setOut(byteArrayOutputStream);
        execChannel.open();
        execChannel.waitFor(Collections.singleton(ClientChannelEvent.CLOSED), 0);
        execChannel.close();
        // 结果写入
        result = byteArrayOutputStream.toString();
        byteArrayOutputStream.close();
        return result;
    }


    /**
     * 获取主机名
     */
    public static String hostname(ClientSession session) throws IOException {
        return SshUtil.execute(session, "hostname").replace("\n", "");
    }


    /**
     * 获取 cpu 使用率
     */
    public static double cpuUsage(ClientSession session) throws IOException, InterruptedException {

        Map<?, ?> map1 = cpuInfo(session);
        Thread.sleep(5 * 1000);
        Map<?, ?> map2 = cpuInfo(session);

        long user1 = Long.parseLong(map1.get("user").toString());
        long nice1 = Long.parseLong(map1.get("nice").toString());
        long system1 = Long.parseLong(map1.get("system").toString());
        long idle1 = Long.parseLong(map1.get("idle").toString());

        long user2 = Long.parseLong(map2.get("user").toString());
        long nice2 = Long.parseLong(map2.get("nice").toString());
        long system2 = Long.parseLong(map2.get("system").toString());
        long idle2 = Long.parseLong(map2.get("idle").toString());

        long total1 = user1 + system1 + nice1;
        long total2 = user2 + system2 + nice2;
        double total = total2 - total1;

        long totalIdle1 = user1 + nice1 + system1 + idle1;
        long totalIdle2 = user2 + nice2 + system2 + idle2;
        double totalIdle = totalIdle2 - totalIdle1;

        return total / totalIdle;

    }

    /**
     * 获取 cpu 总量
     */
    public static long cpuTotal(ClientSession session) throws IOException, InterruptedException {

        Map<?, ?> map1 = cpuInfo(session);
        Thread.sleep(5 * 1000);
        Map<?, ?> map2 = cpuInfo(session);

        long user1 = Long.parseLong(map1.get("user").toString());
        long nice1 = Long.parseLong(map1.get("nice").toString());
        long system1 = Long.parseLong(map1.get("system").toString());
        long idle1 = Long.parseLong(map1.get("idle").toString());

        long user2 = Long.parseLong(map2.get("user").toString());
        long nice2 = Long.parseLong(map2.get("nice").toString());
        long system2 = Long.parseLong(map2.get("system").toString());
        long idle2 = Long.parseLong(map2.get("idle").toString());

        long total1 = user1 + system1 + nice1;
        long total2 = user2 + system2 + nice2;
        double total = total2 - total1;

        long totalIdle1 = user1 + nice1 + system1 + idle1;
        long totalIdle2 = user2 + nice2 + system2 + idle2;

        return totalIdle2 - totalIdle1;

    }


    /**
     * 功能：CPU使用信息
     */
    private static Map<?, ?> cpuInfo(ClientSession session) throws IOException {

        String execute = SshUtil.execute(session, "cat /proc/stat");
        String[] split = execute.split("\n");
        Map<String, Object> map = new HashMap<>();
        for (String line : split) {
            if (line.startsWith("cpu")) {
                StringTokenizer tokenizer = new StringTokenizer(line);
                List<String> temp = new ArrayList<>();
                while (tokenizer.hasMoreElements()) {
                    String value = tokenizer.nextToken();
                    temp.add(value);
                }
                map.put("user", temp.get(1));
                map.put("nice", temp.get(2));
                map.put("system", temp.get(3));
                map.put("idle", temp.get(4));
                map.put("iowait", temp.get(5));
                map.put("irq", temp.get(6));
                map.put("softirq", temp.get(7));
                map.put("stealstolen", temp.get(8));
                break;
            }
        }

        return map;
    }


    /**
     * 显存使用率
     *
     * @param session 会话
     * @return 显存使用率
     */
    public static double gpuUsage(ClientSession session) throws DocumentException, IOException {
        List<GpuDTO> gpuDTOList = SshUtil.gpuInfo(session);
        int used = gpuDTOList.stream().mapToInt(gpu -> Integer.parseInt(gpu.getUsedMemory().split(" ")[0])).sum();
        int total = gpuDTOList.stream().mapToInt(gpu -> Integer.parseInt(gpu.getTotalMemory().split(" ")[0])).sum();
        return used * 1.0 / total;
    }

    /**
     * 显存总量 kb
     *
     * @param session 会话
     * @return 显存总量，单位 kb
     */
    public static long gpuTotal(ClientSession session) throws DocumentException, IOException {
        List<GpuDTO> gpuDTOList = SshUtil.gpuInfo(session);
        return gpuDTOList.stream().mapToLong(gpu -> Integer.parseInt(gpu.getTotalMemory().split(" ")[0])).sum() * 1024;
    }


    /**
     * 功能：CPU使用信息
     */
    public static List<GpuDTO> gpuInfo(ClientSession session) throws IOException, DocumentException {
        String execute = SshUtil.execute(session, "nvidia-smi -q -x");
        String REG = "<!DOCTYPE.*.dtd\">";
        execute = execute.replaceAll(REG, "");
        Document document = DocumentHelper.parseText(execute);
        List<Element> gpu = document.getRootElement().elements("gpu");
        List<GpuDTO> gpuDTOList = new ArrayList<>();
        gpu.forEach(element -> {
            GpuDTO gpuDTOInfo = new GpuDTO();
            String uuid = element.element("uuid").getText();
            Element fbMemoryUsage = element.element("fb_memory_usage");
            String total = fbMemoryUsage.element("total").getText();
            String used = fbMemoryUsage.element("used").getText();
            String free = fbMemoryUsage.element("free").getText();
            gpuDTOInfo.setTotalMemory(total);
            gpuDTOInfo.setUsedMemory(used);
            gpuDTOInfo.setFreeMemory(free);
            gpuDTOInfo.setName(uuid);
            Element processes = element.element("processes");
            List<Element> infos = processes.elements("process_info");
            List<ProcessDTO> processDTOInfos = new ArrayList<>();
            infos.forEach(info -> {
                ProcessDTO processDTO = new ProcessDTO();
                String pid = info.element("pid").getText();
                String name = info.element("process_name").getText();
                String usedMemory = info.element("used_memory").getText();
                processDTO.setPid(pid);
                processDTO.setName(name);
                processDTO.setUsedMemory(usedMemory);
                processDTOInfos.add(processDTO);
            });
            gpuDTOInfo.setProcessDTOS(processDTOInfos);
            int intTotal = Integer.parseInt(total.split(" ")[0]);
            int intUsed = Integer.parseInt(used.split(" ")[0]);
            gpuDTOInfo.setUsageRate((int) ((float) intUsed / intTotal * 100));
            gpuDTOList.add(gpuDTOInfo);
        });
        return gpuDTOList;
    }


    /**
     * 内存使用率
     *
     * @param session 会话
     * @return 剩余可用内存
     */
    public static double memoryUsage(ClientSession session) throws IOException {
        String execute = SshUtil.execute(session, "cat /proc/meminfo");
        String[] split = execute.split("\n");
        Map<String, Object> map = new HashMap<>();
        for (String line : split) {
            if (line == null)
                break;
            int beginIndex = 0;
            int endIndex = line.indexOf(":");
            if (endIndex != -1) {
                String key = line.substring(beginIndex, endIndex);
                beginIndex = endIndex + 1;
                endIndex = line.length();
                String memory = line.substring(beginIndex, endIndex);
                String value = memory.replace("kB", "").trim();
                map.put(key, value);
            }
        }

        long memTotal = Long.parseLong(map.get("MemTotal").toString());
        long memFree = Long.parseLong(map.get("MemFree").toString());
        long memused = memTotal - memFree;
        long buffers = Long.parseLong(map.get("Buffers").toString());
        long cached = Long.parseLong(map.get("Cached").toString());

        return (double) (memused - buffers - cached) / memTotal;
    }

    /**
     * 剩余可用内存(kb)
     *
     * @param session 会话
     * @return 剩余可用内存(kb)
     */
    public static long memoryAvailable(ClientSession session) throws IOException {
        String execute = SshUtil.execute(session, "cat /proc/meminfo");
        String[] split = execute.split("\n");
        Map<String, Object> map = new HashMap<>();
        for (String line : split) {
            if (line == null)
                break;
            int beginIndex = 0;
            int endIndex = line.indexOf(":");
            if (endIndex != -1) {
                String key = line.substring(beginIndex, endIndex);
                beginIndex = endIndex + 1;
                endIndex = line.length();
                String memory = line.substring(beginIndex, endIndex);
                String value = memory.replace("kB", "").trim();
                map.put(key, value);
            }
        }

        long memTotal = Long.parseLong(map.get("MemTotal").toString());
        long memFree = Long.parseLong(map.get("MemFree").toString());
        long memused = memTotal - memFree;
        long buffers = Long.parseLong(map.get("Buffers").toString());
        long cached = Long.parseLong(map.get("Cached").toString());

        return memTotal - (memused - buffers - cached);
    }


    /**
     * 总内存(kb)
     *
     * @param session 会话
     * @return 总内存(kb)
     */
    public static long memoryTotal(ClientSession session) throws IOException {
        String execute = SshUtil.execute(session, "cat /proc/meminfo");
        String[] split = execute.split("\n");
        Map<String, Object> map = new HashMap<>();
        for (String line : split) {
            if (line == null)
                break;
            int beginIndex = 0;
            int endIndex = line.indexOf(":");
            if (endIndex != -1) {
                String key = line.substring(beginIndex, endIndex);
                beginIndex = endIndex + 1;
                endIndex = line.length();
                String memory = line.substring(beginIndex, endIndex);
                String value = memory.replace("kB", "").trim();
                map.put(key, value);
            }
        }

        long memTotal = Long.parseLong(map.get("MemTotal").toString());
        long memFree = Long.parseLong(map.get("MemFree").toString());
        long memused = memTotal - memFree;
        long buffers = Long.parseLong(map.get("Buffers").toString());
        long cached = Long.parseLong(map.get("Cached").toString());

        return memTotal;
    }


    /**
     * 磁盘使用率
     *
     * @param session 会话
     * @return 磁盘使用率
     */
    public static double diskUsage(ClientSession session) throws IOException {
        String execute = execute(session, "df --total");
        String[] split = execute.split("\n");
        String[] split1 = split[split.length - 1].split("\\s+");
        long totalSize = Long.parseLong(split1[1]);
        long totalUsed = Long.parseLong(split1[2]);
        long totalAvailable = Long.parseLong(split1[3]);
        String totalUsedPercent = split1[4];
        return totalUsed * 1.0 / totalSize;
    }

    /**
     * 剩余可用磁盘(kb)
     *
     * @param session 会话
     * @return 剩余可用磁盘(kb)
     */
    public static long diskAvailable(ClientSession session) throws IOException {
        String execute = execute(session, "df --total");
        String[] split = execute.split("\n");
        String[] split1 = split[split.length - 1].split("\\s+");
        long total = Long.parseLong(split1[1]);
        long used = Long.parseLong(split1[2]);
        long available = Long.parseLong(split1[3]);
        String usedPercent = split1[4];
        return available;
    }

    /**
     * 总磁盘(kb)
     *
     * @param session 会话
     * @return 总磁盘(kb)
     */
    public static long diskTotal(ClientSession session) throws IOException {
        String execute = execute(session, "df --total");
        String[] split = execute.split("\n");
        String[] split1 = split[split.length - 1].split("\\s+");
        long total = Long.parseLong(split1[1]);
        long used = Long.parseLong(split1[2]);
        long available = Long.parseLong(split1[3]);
        String usedPercent = split1[4];
        return total;
    }

}
