package auth

import (
	"crypto"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/rsa"
	"encoding/pem"
	"errors"
	"fmt"
	"strings"

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

const (
	SSHKeyAlgoEd25519 = "ed25519"
	SSHKeyAlgoRSA     = "rsa"

	DefaultSSHUsername   = "clustercanvas"
	DefaultRSABits       = 4096
	DefaultEd25519Rounds = 100
)

// GeneratedSSHKey is the result of generating or importing an SSH key pair.
type GeneratedSSHKey struct {
	PrivateKeyPEM string
	PublicKey     string
	Algorithm     string
	RSABits       int
	KDFRounds     int
}

// GenerateSSHKey creates an OpenSSH private key and authorized_keys public line.
//
// Private keys are written without an OpenSSH passphrase. At-rest protection is
// provided by ClusterCanvas AES-GCM (node-keys.enc). KDFRounds is retained as
// generation metadata (ssh-keygen -a) for future passphrase-protected exports;
// golang.org/x/crypto/ssh hardcodes bcrypt rounds when encrypting.
func GenerateSSHKey(algorithm string, rsaBits int, kdfRounds int) (GeneratedSSHKey, error) {
	algorithm = strings.ToLower(strings.TrimSpace(algorithm))
	switch algorithm {
	case SSHKeyAlgoEd25519:
		if kdfRounds <= 0 {
			kdfRounds = DefaultEd25519Rounds
		}
		_, privateKey, err := ed25519.GenerateKey(rand.Reader)
		if err != nil {
			return GeneratedSSHKey{}, fmt.Errorf("generate ed25519 key: %w", err)
		}
		return marshalGeneratedKey(privateKey, SSHKeyAlgoEd25519, 0, kdfRounds)
	case SSHKeyAlgoRSA:
		if rsaBits == 0 {
			rsaBits = DefaultRSABits
		}
		if !ValidRSABits(rsaBits) {
			return GeneratedSSHKey{}, fmt.Errorf("rsa bits must be one of 2048, 3072, or 4096")
		}
		privateKey, err := rsa.GenerateKey(rand.Reader, rsaBits)
		if err != nil {
			return GeneratedSSHKey{}, fmt.Errorf("generate rsa key: %w", err)
		}
		return marshalGeneratedKey(privateKey, SSHKeyAlgoRSA, rsaBits, 0)
	default:
		return GeneratedSSHKey{}, fmt.Errorf("algorithm must be %q or %q", SSHKeyAlgoEd25519, SSHKeyAlgoRSA)
	}
}

// ParseSSHPrivateKey accepts an OpenSSH/PEM private key and derives the public key.
func ParseSSHPrivateKey(privateKeyPEM string) (GeneratedSSHKey, error) {
	trimmed := strings.TrimSpace(privateKeyPEM)
	if trimmed == "" {
		return GeneratedSSHKey{}, errors.New("private_key is required")
	}

	rawKey, err := ssh.ParseRawPrivateKey([]byte(trimmed))
	if err != nil {
		return GeneratedSSHKey{}, fmt.Errorf("invalid private key: %w", err)
	}

	signer, ok := rawKey.(crypto.Signer)
	if !ok {
		return GeneratedSSHKey{}, errors.New("private key does not support signing")
	}

	sshPublicKey, err := ssh.NewPublicKey(signer.Public())
	if err != nil {
		return GeneratedSSHKey{}, fmt.Errorf("derive public key: %w", err)
	}

	algorithm := SSHKeyAlgoEd25519
	rsaBits := 0
	switch key := rawKey.(type) {
	case *ed25519.PrivateKey, ed25519.PrivateKey:
		algorithm = SSHKeyAlgoEd25519
	case *rsa.PrivateKey:
		algorithm = SSHKeyAlgoRSA
		rsaBits = key.N.BitLen()
	default:
		// Keep OpenSSH type string for unusual keys.
		algorithm = sshPublicKey.Type()
	}

	return GeneratedSSHKey{
		PrivateKeyPEM: trimmed + "\n",
		PublicKey:     strings.TrimSpace(string(ssh.MarshalAuthorizedKey(sshPublicKey))),
		Algorithm:     algorithm,
		RSABits:       rsaBits,
	}, nil
}

// ValidRSABits reports whether bits is an allowed RSA size.
func ValidRSABits(bits int) bool {
	switch bits {
	case 2048, 3072, 4096:
		return true
	default:
		return false
	}
}

func marshalGeneratedKey(privateKey crypto.PrivateKey, algorithm string, rsaBits int, kdfRounds int) (GeneratedSSHKey, error) {
	block, err := ssh.MarshalPrivateKey(privateKey, "")
	if err != nil {
		return GeneratedSSHKey{}, fmt.Errorf("marshal private key: %w", err)
	}

	signer, ok := privateKey.(crypto.Signer)
	if !ok {
		return GeneratedSSHKey{}, errors.New("private key does not support signing")
	}
	sshPublicKey, err := ssh.NewPublicKey(signer.Public())
	if err != nil {
		return GeneratedSSHKey{}, fmt.Errorf("derive public key: %w", err)
	}

	privatePEM := string(pem.EncodeToMemory(block))
	return GeneratedSSHKey{
		PrivateKeyPEM: privatePEM,
		PublicKey:     strings.TrimSpace(string(ssh.MarshalAuthorizedKey(sshPublicKey))),
		Algorithm:     algorithm,
		RSABits:       rsaBits,
		KDFRounds:     kdfRounds,
	}, nil
}