package ssh

import (
	"fmt"
	"io"
	"net"
	"os"

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

type agentDialer func(network, address string) (net.Conn, error)

func authMethods(keyFile string) ([]cryptossh.AuthMethod, func(), error) {
	return authMethodsWithAgent(os.Getenv("SSH_AUTH_SOCK"), net.Dial, keyFile)
}

func authMethodsWithAgent(agentSocket string, dial agentDialer, keyFile string) ([]cryptossh.AuthMethod, func(), error) {
	var methods []cryptossh.AuthMethod
	var closers []io.Closer
	if agentSocket != "" {
		methods = append(methods, cryptossh.PublicKeysCallback(func() ([]cryptossh.Signer, error) {
			conn, err := dial("unix", agentSocket)
			if err != nil {
				return nil, err
			}
			closers = append(closers, conn)
			return agent.NewClient(conn).Signers()
		}))
	}
	if keyFile != "" {
		signer, err := signerFromKeyFile(keyFile)
		if err != nil {
			return nil, nil, err
		}
		methods = append(methods, cryptossh.PublicKeys(signer))
	}
	if len(methods) == 0 {
		return nil, nil, fmt.Errorf("no SSH auth methods configured; set SSH_AUTH_SOCK or ssh_key_file")
	}
	return methods, func() { closeAll(closers) }, nil
}

func signerFromKeyFile(path string) (cryptossh.Signer, error) {
	key, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read ssh_key_file %q: %w", path, err)
	}
	signer, err := cryptossh.ParsePrivateKey(key)
	if err != nil {
		return nil, fmt.Errorf("parse ssh_key_file %q: %w", path, err)
	}
	return signer, nil
}

func closeAll(closers []io.Closer) {
	for _, closer := range closers {
		_ = closer.Close()
	}
}