package com.sshmanager.service;

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.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 && password != null) {
            session.setPassword(password);
        }

        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 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 disconnect() {
            if (channel != null && channel.isConnected()) {
                channel.disconnect();
            }
            if (session != null && session.isConnected()) {
                session.disconnect();
            }
        }

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