package com.sshmanager.service;

import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.ChannelShell;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.Session;
import com.sshmanager.entity.Connection;
import org.springframework.stereotype.Service;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.PipedInputStream;
import java.io.PipedOutputStream;

@Service
public class SshService {

    public SshSession createShellSession(Connection conn, String password, String privateKey, String passphrase)
        throws Exception {
        JSch jsch = new JSch();

        if (conn.getAuthType() == Connection.AuthType.PRIVATE_KEY && privateKey != null && !privateKey.isEmpty()) {
            byte[] keyBytes = privateKey.getBytes(StandardCharsets.UTF_8);
            byte[] passphraseBytes = (passphrase != null && !passphrase.isEmpty())
                ? passphrase.getBytes(StandardCharsets.UTF_8) : null;
            jsch.addIdentity("key", keyBytes, null, passphraseBytes);
        }

        Session session = jsch.getSession(conn.getUsername(), conn.getHost(), conn.getPort());
        session.setConfig("StrictHostKeyChecking", "no");
        // Use only DH-based kex to avoid "Algorithm ECDH not available" on Java 8 / minimal JRE
        session.setConfig("kex", "diffie-hellman-group-exchange-sha256,diffie-hellman-group14-sha256,diffie-hellman-group14-sha1,diffie-hellman-group-exchange-sha1");

        if (conn.getAuthType() == Connection.AuthType.PASSWORD) {
            if (password == null || password.isEmpty()) {
                throw new IllegalArgumentException("Password is required for password authentication");
            }
            session.setConfig("PreferredAuthentications", "password");
            session.setPassword(password);
        } else {
            session.setConfig("PreferredAuthentications", "publickey");
        }

        session.connect(10000);

        ChannelShell channel = (ChannelShell) session.openChannel("shell");
        channel.setPtyType("xterm");
        channel.connect(5000);

        PipedInputStream pipedIn = new PipedInputStream();
        PipedOutputStream pipeToChannel = new PipedOutputStream(pipedIn);
        OutputStream channelIn = channel.getOutputStream();
        InputStream channelOut = channel.getInputStream();

        new Thread(() -> {
            try {
                byte[] buf = new byte[1024];
                int n;
                while ((n = pipedIn.read(buf)) > 0) {
                    channelIn.write(buf, 0, n);
                    channelIn.flush();
                }
            } catch (Exception e) {
                // Channel closed
            }
        }).start();

        return new SshSession(session, channel, channelOut, pipeToChannel);
    }

    // 执行单次命令并返回输出
    public String executeCommand(Connection conn, String password, String privateKey, String passphrase, String command) throws Exception {
        JSch jsch = new JSch();

        if (conn.getAuthType() == Connection.AuthType.PRIVATE_KEY && privateKey != null && !privateKey.isEmpty()) {
            byte[] keyBytes = privateKey.getBytes(StandardCharsets.UTF_8);
            byte[] passphraseBytes = (passphrase != null && !passphrase.isEmpty())
                ? passphrase.getBytes(StandardCharsets.UTF_8) : null;
            jsch.addIdentity("key", keyBytes, null, passphraseBytes);
        }

        Session session = jsch.getSession(conn.getUsername(), conn.getHost(), conn.getPort());
        session.setConfig("StrictHostKeyChecking", "no");
        session.setConfig("kex", "diffie-hellman-group-exchange-sha256,diffie-hellman-group14-sha256,diffie-hellman-group14-sha1,diffie-hellman-group-exchange-sha1");
        session.setConfig("PreferredAuthentications", conn.getAuthType() == Connection.AuthType.PASSWORD ? "password" : "publickey");

        if (conn.getAuthType() == Connection.AuthType.PASSWORD) {
            session.setPassword(password);
        }

        session.connect(8000);

        ChannelExec channel = (ChannelExec) session.openChannel("exec");
        channel.setCommand(command);
        channel.setErrStream(System.err);

        InputStream in = channel.getInputStream();
        channel.connect(3000);

        BufferedReader reader = new BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8));
        StringBuilder result = new StringBuilder();
        String line;
        while ((line = reader.readLine()) != null) {
            result.append(line).append("\n");
        }

        channel.disconnect();
        session.disconnect();

        return result.toString().trim();
    }

    public static class SshSession {
        private final Session session;
        private final ChannelShell channel;
        private final InputStream outputStream;
        private final OutputStream inputStream;

        public SshSession(Session session, ChannelShell channel, InputStream outputStream, OutputStream inputStream) {
            this.session = session;
            this.channel = channel;
            this.outputStream = outputStream;
            this.inputStream = inputStream;
        }

        public InputStream getOutputStream() {
            return outputStream;
        }

        public OutputStream getInputStream() {
            return inputStream;
        }

        public void resize(int cols, int rows) {
            if (cols <= 0 || rows <= 0) {
                return;
            }
            if (channel == null) {
                return;
            }
            try {
                channel.setPtySize(cols, rows, 0, 0);
            } catch (Exception ignored) {
            }
        }

        public void disconnect() {
            if (channel != null && channel.isConnected()) {
                channel.disconnect();
            }
            if (session != null && session.isConnected()) {
                session.disconnect();
            }
        }

        public boolean isConnected() {
            return channel != null && channel.isConnected();
        }
    }
}