package ssh

import (
	"fmt"
	"net"

	"golang.org/x/crypto/ssh"
)

// RunCommandViaSSH runs a command on a remote host via SSH, returns the output and error if any.
func RunCommandViaSSH(host, user, privateKey, cmd string) (string, error) {
	authMethod, err := publicKeyAuthFunc(privateKey)
	if err != nil {
		return "", fmt.Errorf("failed to create auth method: %w", err)
	}
	config := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{authMethod},
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}
	addr := fmt.Sprintf("%s:22", host)
	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		return "", err
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()

	output, err := session.CombinedOutput(cmd)
	if err != nil {
		return "", err
	}
	return string(output), nil
}

func publicKeyAuthFunc(privateKey string) (ssh.AuthMethod, error) {
	signer, err := ssh.ParsePrivateKey([]byte(privateKey))
	if err != nil {
		return nil, fmt.Errorf("failed to parse private key: %w", err)
	}
	return ssh.PublicKeys(signer), nil
}