package aws

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	smithy "github.com/aws/smithy-go"
)

const (
	UserManagedEC2KeyFailureInvalid     = "aws-key-pair-invalid"
	UserManagedEC2KeyFailureNotFound    = "aws-key-pair-not-found"
	UserManagedEC2KeyFailurePermission  = "aws-key-pair-permission"
	UserManagedEC2KeyFailureUnavailable = "aws-key-pair-unavailable"
)

var (
	errUserManagedKeyName         = errors.New("EC2 key-pair name is invalid")
	errUserManagedPrivateKey      = errors.New("EC2 private-key input is invalid")
	errUserManagedSSHKeygen       = errors.New("ssh-keygen is unavailable")
	errUserManagedEncryptedKey    = errors.New("EC2 private key is encrypted or unusable")
	errUserManagedMalformedKey    = errors.New("EC2 private key is malformed or unsupported")
	errUserManagedRemotePublicKey = errors.New("EC2 public key is malformed or unsupported")
	errUserManagedKeyMismatch     = errors.New("EC2 private key does not match key pair")
	errUserManagedKeyPairShape    = errors.New("EC2 key pair result is ambiguous")
	errUserManagedKeyNotFound     = errors.New("EC2 key pair was not found")
	errUserManagedKeyPermission   = errors.New("EC2 key-pair permission is missing")
)

// UserManagedEC2KeyValidation records the non-secret result of the EC2 key
// prerequisite check.
type UserManagedEC2KeyValidation struct {
	KeyPair         string `json:"key_pair"`
	Algorithm       string `json:"algorithm"`
	PublicKeySHA256 string `json:"public_key_sha256"`
	Account         string `json:"account,omitempty"`
	Region          string `json:"region"`
	EqualitySkipped bool   `json:"equality_skipped,omitempty"`
}

// Summary returns the same safe, non-key-material result used by setup,
// dry-run, and doctor output.
func (v *UserManagedEC2KeyValidation) Summary() string {
	if v == nil {
		return ""
	}
	equality := "verified"
	if v.EqualitySkipped {
		equality = "skipped-persisted-instance"
	}
	return fmt.Sprintf(
		"key_pair=%s algorithm=%s public_key_sha256=%s account=%s region=%s equality=%s",
		v.KeyPair,
		v.Algorithm,
		v.PublicKeySHA256,
		v.Account,
		v.Region,
		equality,
	)
}

// UserManagedEC2KeyDiagnostic describes a safe, actionable failure without
// exposing private or public key material.
type UserManagedEC2KeyDiagnostic struct {
	FailureClass string
	Field        string
	Remediation  string
}

// DiagnoseUserManagedEC2Key classifies key prerequisite failures for all
// command surfaces that report setup, dry-run, or doctor diagnostics.
func DiagnoseUserManagedEC2Key(err error) UserManagedEC2KeyDiagnostic {
	switch {
	case errors.Is(err, errUserManagedKeyNotFound):
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailureNotFound,
			Field:        "host.ec2.key_pair",
			Remediation:  "create or select an existing EC2 key pair in the target account and region, then rerun cp9 setup",
		}
	case errors.Is(err, errUserManagedKeyPermission):
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailurePermission,
			Field:        "target.aws.iam.role_arn",
			Remediation:  "grant ec2:DescribeKeyPairs to the configured role in the target account and region, then rerun cp9 setup",
		}
	case errors.Is(err, errUserManagedKeyName):
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailureInvalid,
			Field:        "host.ec2.key_pair",
			Remediation:  "set host.ec2.key_pair to the existing EC2 key-pair name in the reported account and region",
		}
	case errors.Is(err, errUserManagedSSHKeygen):
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailureUnavailable,
			Field:        "host.ec2.ssh_private_key",
			Remediation:  "install OpenSSH ssh-keygen and ensure it is on PATH, then rerun cp9 setup",
		}
	case errors.Is(err, errUserManagedEncryptedKey):
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailureInvalid,
			Field:        "host.ec2.ssh_private_key",
			Remediation:  "provide an unencrypted RSA or ED25519 private-key file that ssh-keygen can read; do not import or delete a key pair",
		}
	case errors.Is(err, errUserManagedMalformedKey):
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailureInvalid,
			Field:        "host.ec2.ssh_private_key",
			Remediation:  "provide a valid RSA or ED25519 private-key file in a format supported by OpenSSH ssh-keygen",
		}
	case errors.Is(err, errUserManagedRemotePublicKey):
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailureInvalid,
			Field:        "host.ec2.key_pair",
			Remediation:  "select an existing EC2 key pair with a valid RSA or ED25519 public key matching the configured private key",
		}
	case errors.Is(err, errUserManagedKeyMismatch):
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailureInvalid,
			Field:        "host.ec2",
			Remediation:  "use the matching existing EC2 key pair and private key; do not import, rotate, or delete either key",
		}
	case errors.Is(err, errUserManagedKeyPairShape):
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailureInvalid,
			Field:        "host.ec2.key_pair",
			Remediation:  "select one existing EC2 key pair name that resolves to exactly one RSA or ED25519 key",
		}
	case errors.Is(err, errUserManagedPrivateKey):
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailureInvalid,
			Field:        "host.ec2.ssh_private_key",
			Remediation:  "provide an existing absolute readable regular RSA or ED25519 private-key file at host.ec2.ssh_private_key",
		}
	case IsContractError(err):
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailureInvalid,
			Field:        "host.ec2",
			Remediation:  "verify the existing EC2 key pair and local private key match, then rerun cp9 setup",
		}
	default:
		return UserManagedEC2KeyDiagnostic{
			FailureClass: UserManagedEC2KeyFailureUnavailable,
			Field:        "target.aws.iam.role_arn",
			Remediation:  "verify AWS access and EC2 DescribeKeyPairs availability in the target account and region, then rerun cp9 setup",
		}
	}
}

type userManagedEC2KeyPairAPI interface {
	DescribeKeyPairs(
		ctx context.Context,
		input *ec2.DescribeKeyPairsInput,
		optFns ...func(*ec2.Options),
	) (*ec2.DescribeKeyPairsOutput, error)
}

var newUserManagedEC2KeyPairClient = func(cfg *awssdk.Config) userManagedEC2KeyPairAPI {
	return ec2.NewFromConfig(*cfg)
}

// ValidateUserManagedEC2Key proves that the configured local private key is
// the same RSA or ED25519 public key already present in EC2. It never creates,
// tags, imports, rotates, or deletes a key pair. A persisted instance owns the
// resume/destroy path, so its current key object is not a resume blocker.
func ValidateUserManagedEC2Key(
	ctx context.Context,
	cfg *awssdk.Config,
	account,
	region,
	keyPair,
	privateKeyPath,
	persistedInstanceID string,
) (*UserManagedEC2KeyValidation, error) {
	if strings.TrimSpace(region) == "" && cfg != nil {
		region = cfg.Region
	}
	var client userManagedEC2KeyPairAPI
	if cfg != nil {
		client = newUserManagedEC2KeyPairClient(cfg)
	}
	return validateUserManagedEC2Key(
		ctx,
		client,
		account,
		region,
		keyPair,
		privateKeyPath,
		persistedInstanceID,
	)
}

func validateUserManagedEC2Key(
	ctx context.Context,
	client userManagedEC2KeyPairAPI,
	account,
	region,
	keyPair,
	privateKeyPath,
	persistedInstanceID string,
) (*UserManagedEC2KeyValidation, error) {
	keyPair = strings.TrimSpace(keyPair)
	privateKeyPath = strings.TrimSpace(privateKeyPath)
	keyContext := userManagedEC2KeyContext(account, region, keyPair, privateKeyPath)
	if keyPair == "" {
		return nil, newContractError(fmt.Errorf(
			"%w: %s; set host.ec2.key_pair to an existing EC2 key-pair name before the first remote host mutation",
			errUserManagedKeyName,
			keyContext,
		))
	}
	if privateKeyPath == "" || !filepath.IsAbs(privateKeyPath) {
		return nil, newContractError(fmt.Errorf(
			"%w: %s; set host.ec2.ssh_private_key to an absolute path before the first remote host mutation",
			errUserManagedPrivateKey,
			keyContext,
		))
	}
	info, err := os.Stat(privateKeyPath)
	if err != nil {
		action := "provide an existing readable regular private-key file"
		if !os.IsNotExist(err) {
			action = "grant read access to the configured regular private-key file"
		}
		return nil, newContractError(fmt.Errorf(
			"%w: %s; %s",
			errUserManagedPrivateKey,
			keyContext,
			action,
		))
	}
	if !info.Mode().IsRegular() {
		return nil, newContractError(fmt.Errorf(
			"%w: %s; replace it with an existing readable regular private-key file",
			errUserManagedPrivateKey,
			keyContext,
		))
	}
	keyFile, err := os.Open(privateKeyPath)
	if err != nil {
		return nil, newContractError(fmt.Errorf(
			"%w: %s; grant read access to the configured regular private-key file",
			errUserManagedPrivateKey,
			keyContext,
		))
	}
	_ = keyFile.Close()

	local, err := deriveUserManagedPublicKey(ctx, privateKeyPath)
	if err != nil {
		return nil, newContractError(fmt.Errorf("%w: %s; %v", err, keyContext, userManagedKeyRemediation(err)))
	}
	validation := &UserManagedEC2KeyValidation{
		KeyPair:         keyPair,
		Algorithm:       local.algorithm,
		PublicKeySHA256: sshPublicKeyFingerprint(local.blob),
		Account:         strings.TrimSpace(account),
		Region:          strings.TrimSpace(region),
	}
	if strings.TrimSpace(persistedInstanceID) != "" {
		validation.EqualitySkipped = true
		return validation, nil
	}
	if client == nil {
		return nil, fmt.Errorf(
			"%s; AWS config is required before the first remote host mutation",
			keyContext,
		)
	}
	out, err := client.DescribeKeyPairs(ctx, &ec2.DescribeKeyPairsInput{
		IncludePublicKey: awssdk.Bool(true),
		KeyNames:         []string{keyPair},
	})
	if err != nil {
		return nil, describeUserManagedKeyPairError(account, region, keyPair, privateKeyPath, err)
	}
	if out == nil || len(out.KeyPairs) == 0 {
		return nil, fmt.Errorf(
			"%w: %s; create or select this existing EC2 key pair in the reported account and region",
			errUserManagedKeyNotFound,
			keyContext,
		)
	}
	if len(out.KeyPairs) != 1 {
		return nil, newContractError(fmt.Errorf(
			"%w: %s; select one existing EC2 key pair name that resolves to exactly one key",
			errUserManagedKeyPairShape,
			keyContext,
		))
	}
	remote, err := parseOpenSSHKey(awssdk.ToString(out.KeyPairs[0].PublicKey))
	if err != nil {
		return nil, newContractError(fmt.Errorf(
			"%w: %s; EC2 returned an unusable public key (%w); select an existing RSA or ED25519 key pair with a valid public key",
			errUserManagedRemotePublicKey,
			keyContext,
			err,
		))
	}
	if local.algorithm != remote.algorithm || !bytes.Equal(local.blob, remote.blob) {
		return nil, newContractError(fmt.Errorf(
			"%w: %s; use the matching existing EC2 key pair and private key, without importing, rotating, or deleting either key",
			errUserManagedKeyMismatch,
			keyContext,
		))
	}
	return validation, nil
}

func describeUserManagedKeyPairError(account, region, keyPair, privateKeyPath string, err error) error {
	keyContext := userManagedEC2KeyContext(account, region, keyPair, privateKeyPath)
	code := userManagedEC2KeyAPIErrorCode(err)
	switch {
	case userManagedEC2KeyNotFoundCode(code):
		return fmt.Errorf(
			"%w: %s; create or select this existing EC2 key pair in the reported account and region",
			errUserManagedKeyNotFound,
			keyContext,
		)
	case userManagedEC2KeyPermissionCode(code):
		return fmt.Errorf(
			"%w: %s; grant ec2:DescribeKeyPairs to the configured role in the reported account and region, then rerun cp9 setup",
			errUserManagedKeyPermission,
			keyContext,
		)
	default:
		return fmt.Errorf(
			"describe configured EC2 key pair failed: %s; verify AWS access and ec2:DescribeKeyPairs in the reported account and region (api_code=%s): %w",
			keyContext,
			code,
			err,
		)
	}
}

func userManagedEC2KeyContext(account, region, keyPair, privateKeyPath string) string {
	if strings.TrimSpace(account) == "" {
		account = "<unresolved>"
	}
	if strings.TrimSpace(region) == "" {
		region = "<unset>"
	}
	if strings.TrimSpace(keyPair) == "" {
		keyPair = "<unset>"
	}
	if strings.TrimSpace(privateKeyPath) == "" {
		privateKeyPath = "<unset>"
	}
	return fmt.Sprintf("account=%s region=%s key_pair=%q path=%q", account, region, keyPair, privateKeyPath)
}

func userManagedKeyRemediation(err error) string {
	switch {
	case errors.Is(err, errUserManagedSSHKeygen):
		return "install OpenSSH ssh-keygen and ensure it is on PATH"
	case errors.Is(err, errUserManagedEncryptedKey):
		return "provide an unencrypted RSA or ED25519 private-key file that ssh-keygen can read"
	default:
		return "provide a valid RSA or ED25519 private-key file in a format supported by OpenSSH ssh-keygen"
	}
}

func userManagedEC2KeyAPIErrorCode(err error) string {
	var apiErr smithy.APIError
	if errors.As(err, &apiErr) {
		return apiErr.ErrorCode()
	}
	return ""
}

func userManagedEC2KeyNotFoundCode(code string) bool {
	return code == "InvalidKeyPair.NotFound" || code == "InvalidKeyPair.NotFoundException"
}

func userManagedEC2KeyPermissionCode(code string) bool {
	switch code {
	case "UnauthorizedOperation", "AccessDenied", "AccessDeniedException", "AuthFailure":
		return true
	default:
		return false
	}
}

type parsedOpenSSHKey struct {
	algorithm string
	blob      []byte
}

func deriveUserManagedPublicKey(ctx context.Context, privateKeyPath string) (*parsedOpenSSHKey, error) {
	command := exec.CommandContext(ctx, "ssh-keygen", "-y", "-f", privateKeyPath)
	command.Stdin = strings.NewReader("")
	output, err := command.CombinedOutput()
	if err != nil {
		if errors.Is(err, exec.ErrNotFound) {
			return nil, errUserManagedSSHKeygen
		}
		message := strings.ToLower(string(output))
		if strings.Contains(message, "passphrase") || strings.Contains(message, "encrypted") {
			return nil, errUserManagedEncryptedKey
		}
		return nil, errUserManagedMalformedKey
	}
	parsed, err := parseOpenSSHKey(string(output))
	if err != nil {
		return nil, fmt.Errorf("%w: %w", errUserManagedMalformedKey, err)
	}
	return parsed, nil
}

func parseOpenSSHKey(value string) (*parsedOpenSSHKey, error) {
	fields := strings.Fields(value)
	if len(fields) < 2 {
		return nil, fmt.Errorf("public key is not in OpenSSH format")
	}
	algorithm := fields[0]
	if algorithm != "ssh-rsa" && algorithm != "ssh-ed25519" {
		return nil, fmt.Errorf("unsupported public key algorithm %q", algorithm)
	}
	blob, err := decodeUserManagedPublicKeyBlob(fields[1])
	if err != nil {
		return nil, fmt.Errorf("decode public key blob: %w", err)
	}
	return &parsedOpenSSHKey{algorithm: algorithm, blob: blob}, nil
}

func decodeUserManagedPublicKeyBlob(value string) ([]byte, error) {
	for _, encoding := range []*base64.Encoding{base64.StdEncoding, base64.RawStdEncoding} {
		blob, err := encoding.DecodeString(value)
		if err == nil {
			return blob, nil
		}
	}
	return nil, fmt.Errorf("invalid base64 public key blob")
}

func sshPublicKeyFingerprint(blob []byte) string {
	digest := sha256.Sum256(blob)
	return "SHA256:" + base64.RawStdEncoding.EncodeToString(digest[:])
}