package ssh

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"net"
	"os"
	"syscall"
	"time"

	"github.com/mimecast/dtail/internal/io/dlog"

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

// GeneratePrivateRSAKey is used by the server to generate its key.
func GeneratePrivateRSAKey(size int) (*rsa.PrivateKey, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, size)
	if err != nil {
		return nil, fmt.Errorf("failed to generate RSA key: %w", err)
	}
	if err = privateKey.Validate(); err != nil {
		return nil, fmt.Errorf("failed to validate generated RSA key: %w", err)
	}
	return privateKey, nil
}

// EncodePrivateKeyToPEM is a helper function for converting a key to PEM format.
func EncodePrivateKeyToPEM(privateKey *rsa.PrivateKey) []byte {
	derFormat := x509.MarshalPKCS1PrivateKey(privateKey)

	block := pem.Block{
		Type:    "RSA PRIVATE KEY",
		Headers: nil,
		Bytes:   derFormat,
	}
	return pem.EncodeToMemory(&block)
}

// Agent used for SSH auth.
func Agent() (gossh.AuthMethod, error) {
	return AgentWithKeyIndex(-1)
}

// AgentSignersWithKeyIndex returns SSH agent signers.
// If keyIndex is -1, all keys are used. Otherwise, only the specified key is used.
func AgentSignersWithKeyIndex(keyIndex int) ([]gossh.Signer, error) {
	// Use context-aware dialing for SSH agent connection (local Unix socket).
	// 2-second timeout is reasonable for local socket connections.
	dialer := &net.Dialer{
		Timeout: 2 * time.Second,
	}

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	sshAgent, err := dialer.DialContext(ctx, "unix", os.Getenv("SSH_AUTH_SOCK"))
	if err != nil {
		return nil, fmt.Errorf("failed to connect to SSH agent: %w", err)
	}
	agentClient := agent.NewClient(sshAgent)
	keys, err := agentClient.List()
	if err != nil {
		return nil, fmt.Errorf("failed to list SSH agent keys: %w", err)
	}
	for i, key := range keys {
		dlog.Common.Debug("Public key", i, key)
	}

	signers, err := agentClient.Signers()
	if err != nil {
		return nil, fmt.Errorf("failed to load SSH agent signers: %w", err)
	}

	// If no specific key index requested, use all keys (backwards compatible default)
	if keyIndex < 0 {
		return signers, nil
	}

	// Use only the specified key index (0-based)
	if keyIndex >= len(signers) {
		return nil, fmt.Errorf("key index %d out of range (agent has %d signers)", keyIndex, len(signers))
	}

	dlog.Common.Debug("Using SSH agent key at index", keyIndex)
	return []gossh.Signer{signers[keyIndex]}, nil
}

// AgentWithKeyIndex used for SSH auth with a specific key index from the agent.
// If keyIndex is -1, all keys are used. Otherwise, only the specified key is used.
func AgentWithKeyIndex(keyIndex int) (gossh.AuthMethod, error) {
	signers, err := AgentSignersWithKeyIndex(keyIndex)
	if err != nil {
		return nil, err
	}
	return gossh.PublicKeys(signers...), nil
}

// EnterKeyPhrase is required to read phrase protected private keys.
func EnterKeyPhrase(keyFile string) []byte {
	fmt.Printf("Enter phrase for key %s: ", keyFile)
	phrase, err := term.ReadPassword(int(syscall.Stdin))
	if err != nil {
		panic(err)
	}
	fmt.Printf("%s\n", string(phrase))
	return phrase
}

// PrivateKeySigner returns an SSH signer from the provided private key file.
func PrivateKeySigner(keyFile string) (gossh.Signer, error) {
	buffer, err := os.ReadFile(keyFile)
	if err != nil {
		return nil, err
	}
	key, err := gossh.ParsePrivateKey(buffer)
	if err != nil {
		var passphraseMissingErr *gossh.PassphraseMissingError
		if !errors.As(err, &passphraseMissingErr) {
			return nil, err
		}

		passphrase := os.Getenv("DTAIL_KEY_PASSPHRASE")
		if passphrase == "" {
			return nil, err
		}

		key, err = gossh.ParsePrivateKeyWithPassphrase(buffer, []byte(passphrase))
		if err != nil {
			return nil, err
		}
	}
	return key, nil
}

// KeyFile returns the key as a SSH auth method.
func KeyFile(keyFile string) (gossh.AuthMethod, error) {
	key, err := PrivateKeySigner(keyFile)
	if err != nil {
		return nil, err
	}

	// Key phrase support disabled as password will be printed to stdout!
	/*
		if err == nil {
			return gossh.PublicKeys(key), nil
		}

		keyPhrase := EnterKeyPhrase(keyFile)
		key, err = gossh.ParsePrivateKeyWithPassphrase(buffer, keyPhrase)
		if err != nil {
			return nil, err
		}
	*/

	return gossh.PublicKeys(key), nil
}

// PrivateKey returns the private key as a SSH auth method.
func PrivateKey(keyFile string) (gossh.AuthMethod, error) {
	signer, err := KeyFile(keyFile)
	if err != nil {
		dlog.Common.Debug(keyFile, err)
		return nil, err
	}
	return gossh.AuthMethod(signer), nil
}