package sshkey

import (
	"context"
	"fmt"
	"net"
	"os"
	"path/filepath"

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

// Agent is an SSH agent that holds one key and serves it on a unix
// socket. The socket sits in a directory of its own that only its
// owner may enter, and the key stays in memory: nothing is written to
// disk.
type Agent struct {
	socket   string
	listener net.Listener
}

// Serve starts an agent holding this key under the given comment.
// Stop takes it down again.
func (k *Key) Serve(ctx context.Context, comment string) (*Agent, error) {
	keyring := agent.NewKeyring()

	err := keyring.Add(agent.AddedKey{
		PrivateKey: k.private,
		Comment:    comment,
	})
	if err != nil {
		return nil, fmt.Errorf("giving the key to the agent: %w", err)
	}

	// A temporary directory is made enterable by its owner alone,
	// which is the protection the socket inside it has.
	directory, err := os.MkdirTemp("", "keyfunc-agent-")
	if err != nil {
		return nil, fmt.Errorf("making the agent directory: %w", err)
	}

	socket := filepath.Join(directory, "socket")

	var listen net.ListenConfig

	listener, err := listen.Listen(ctx, "unix", socket)
	if err != nil {
		_ = os.RemoveAll(directory)

		return nil, fmt.Errorf("listening on the agent socket: %w", err)
	}

	served := &Agent{socket: socket, listener: listener}

	go served.accept(keyring)

	return served, nil
}

// Socket is the path to point ssh at.
func (a *Agent) Socket() string {
	return a.socket
}

// Stop takes the agent down and removes the socket and the directory
// it is in.
func (a *Agent) Stop() {
	_ = a.listener.Close()
	_ = os.RemoveAll(filepath.Dir(a.socket))
}

// accept answers connections until Stop closes the listener.
func (a *Agent) accept(keyring agent.Agent) {
	for {
		connection, err := a.listener.Accept()
		if err != nil {
			return
		}

		go func() {
			defer func() { _ = connection.Close() }()

			_ = agent.ServeAgent(keyring, connection)
		}()
	}
}