package com.njzscloud.common.sshtunnel;

import cn.hutool.core.util.StrUtil;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.njzscloud.common.core.tuple.Tuple2;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

@Slf4j
public class SSHTunnel {

    private final List<Tuple2<JSch, Session>> tunnels = new LinkedList<>();
    private final ScheduledExecutorService reconnectExecutor = Executors.newSingleThreadScheduledExecutor();
    private final AtomicBoolean running = new AtomicBoolean(false);

    public SSHTunnel(SSHTunnelProperties sshTunnelProperties) {
        if (sshTunnelProperties.isEnable()) {
            log.info("SSH 隧道已启用");
            running.set(true);
            List<SSHTunnelProperties.SSHProperties> tunnelsProperties = sshTunnelProperties.getTunnels();
            tunnelsProperties.stream()
                    .filter(SSHTunnelProperties.SSHProperties::isEnable)
                    .forEach(it -> {
                        Tuple2<JSch, Session> tunnel = createTunnel(it);
                        tunnels.add(tunnel);
                    });
            reconnectTask();
        }
    }

    private Tuple2<JSch, Session> createTunnel(SSHTunnelProperties.SSHProperties sshProperties) {
        try {
            String host = sshProperties.getHost();
            int port = sshProperties.getPort();
            String user = sshProperties.getUser();
            String pwd = sshProperties.getPwd();
            String credentials = sshProperties.getCredentials();
            String passphrase = sshProperties.getPassphrase();
            log.info("正在连接服务器：{}:{}", host, port);
            JSch jsch = new JSch();
            Session session;
            if (StrUtil.isNotBlank(credentials)) {
                // 使用私钥认证
                File keyFile = new File(credentials);
                if (!keyFile.exists()) {
                    throw new RuntimeException("私钥文件不存在: " + credentials);
                }

                if (passphrase != null) {
                    jsch.addIdentity(credentials, passphrase);
                } else {
                    jsch.addIdentity(credentials);
                }
                session = jsch.getSession(user, host, port);
            } else if (StrUtil.isNotBlank(pwd)) {
                // 创建SSH会话
                session = jsch.getSession(user, host, port);
                // 使用密码认证
                session.setPassword(pwd);
            } else {
                throw new RuntimeException("未提供有效认证方式（私钥或密码）");
            }


            session.setConfig("StrictHostKeyChecking", "no"); // 禁用主机密钥检查

            // 设置保持活动状态
            session.setServerAliveInterval(60000); // 每分钟发送一次保持活动的数据包
            session.setServerAliveCountMax(3); // 允许3次失败

            session.connect();

            List<SSHTunnelProperties.ProxyProperties> proxyProperties = sshProperties.getProxy();
            for (SSHTunnelProperties.ProxyProperties proxyProperty : proxyProperties) {
                int localPort = proxyProperty.getLocalPort();
                String targetHost = proxyProperty.getTargetHost();
                int targetPort = proxyProperty.getTargetPort();
                // 设置端口转发 (本地端口 -> 目标地址)
                session.setPortForwardingL(localPort, targetHost, targetPort);
                log.info("已创建隧道,本地端口：{} --> 服务器：{}:{} --> 目标地址：{}:{}", localPort, host, port, targetHost, targetPort);
            }
            return Tuple2.create(jsch, session);
        } catch (JSchException e) {
            log.error("SSH 隧道连接失败", e);
            throw new RuntimeException("SSH 隧道连接失败", e);
        }
    }

    private void reconnectTask() {
        reconnectExecutor.scheduleAtFixedRate(() -> {
            if (!running.get()) return;
            try {
                for (Tuple2<JSch, Session> tunnel : tunnels) {
                    Session session = tunnel.get_1();
                    if (!session.isConnected()) {
                        log.warn("SSH 隧道连接已断开，尝试重新连接...");
                        session.connect();
                    }
                }
            } catch (Exception e) {
                log.error("重连任务执行失败", e);
            }
        }, 30, 30, TimeUnit.SECONDS); // 每30秒检查一次连接状态
    }

    public void close() {
        running.set(false);
        reconnectExecutor.shutdownNow();
        for (Tuple2<JSch, Session> tunnel : tunnels) {
            Session session = tunnel.get_1();
            if (session != null && session.isConnected()) {
                String host = session.getHost();
                int port = session.getPort();
                session.disconnect();
                log.info("SSH 隧道已关闭,服务器：{}:{}", host, port);
            }
        }
        tunnels.clear();
    }
}
