//go:build integration

package aws

import (
	"context"
	"errors"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"

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

type userManagedKeyPairFixture struct {
	publicKey string
	err       error
	calls     int
	input     *ec2.DescribeKeyPairsInput
}

func (f *userManagedKeyPairFixture) DescribeKeyPairs(
	_ context.Context,
	input *ec2.DescribeKeyPairsInput,
	_ ...func(*ec2.Options),
) (*ec2.DescribeKeyPairsOutput, error) {
	f.calls++
	f.input = input
	if f.err != nil {
		return nil, f.err
	}
	return &ec2.DescribeKeyPairsOutput{KeyPairs: []types.KeyPairInfo{{
		KeyName:   awssdk.String("caller-key"),
		PublicKey: awssdk.String(f.publicKey),
	}}}, nil
}

func TestIntegration_UserManagedEC2KeyValidationMatchesRSAAndED25519(t *testing.T) {
	for _, algorithm := range []string{"ed25519", "rsa"} {
		t.Run(algorithm, func(t *testing.T) {
			privatePath := generateUserManagedKey(t, algorithm)
			publicPath := privatePath + ".pub"
			publicBytes, err := os.ReadFile(publicPath)
			if err != nil {
				t.Fatal(err)
			}
			fields := strings.Fields(string(publicBytes))
			if len(fields) < 2 {
				t.Fatalf("public key = %q", publicBytes)
			}
			fixture := &userManagedKeyPairFixture{
				publicKey: fields[0] + " " + strings.TrimRight(fields[1], "=") + " test-comment",
			}
			result, err := validateUserManagedEC2Key(
				context.Background(),
				fixture,
				"123456789012",
				"us-west-2",
				"caller-key",
				privatePath,
				"",
			)
			if err != nil {
				t.Fatal(err)
			}
			if result.EqualitySkipped || result.Algorithm != "ssh-"+algorithm ||
				result.Account != "123456789012" || result.Region != "us-west-2" ||
				!strings.HasPrefix(result.PublicKeySHA256, "SHA256:") {
				t.Fatalf("validation result = %#v", result)
			}
			if fixture.calls != 1 || fixture.input == nil || len(fixture.input.KeyNames) != 1 ||
				fixture.input.KeyNames[0] != "caller-key" || !awssdk.ToBool(fixture.input.IncludePublicKey) {
				t.Fatalf("DescribeKeyPairs input = %#v, calls=%d", fixture.input, fixture.calls)
			}
		})
	}
}

func TestIntegration_UserManagedEC2KeyValidationRejectsMismatchWithoutKeyMaterial(t *testing.T) {
	for _, algorithm := range []string{"ed25519", "rsa"} {
		t.Run(algorithm, func(t *testing.T) {
			privatePath := generateUserManagedKey(t, algorithm)
			otherPath := generateUserManagedKey(t, algorithm)
			otherPublic, err := os.ReadFile(otherPath + ".pub")
			if err != nil {
				t.Fatal(err)
			}
			fixture := &userManagedKeyPairFixture{publicKey: strings.TrimSpace(string(otherPublic))}
			_, err = validateUserManagedEC2Key(
				context.Background(),
				fixture,
				"123456789012",
				"us-west-2",
				"caller-key",
				privatePath,
				"",
			)
			if err == nil || !strings.Contains(err.Error(), "does not match") {
				t.Fatalf("mismatch error = %v", err)
			}
			if fixture.calls != 1 || strings.Contains(err.Error(), strings.TrimSpace(string(otherPublic))) {
				t.Fatalf("mismatch validation leaked material or made unexpected calls: calls=%d err=%v", fixture.calls, err)
			}
			if diagnostic := DiagnoseUserManagedEC2Key(err); diagnostic.Remediation == "" {
				t.Fatalf("mismatch diagnostic = %#v", diagnostic)
			}
		})
	}
}

func TestIntegration_UserManagedEC2KeyValidationNegativeProviderCasesMakeNoMutation(t *testing.T) {
	privatePath := generateUserManagedKey(t, "ed25519")
	otherPath := generateUserManagedKey(t, "ed25519")
	otherPublic, err := os.ReadFile(otherPath + ".pub")
	if err != nil {
		t.Fatal(err)
	}
	for _, testCase := range []struct {
		name        string
		err         error
		publicKey   string
		failure     string
		remediation string
	}{
		{
			name:        "not found",
			err:         &smithy.GenericAPIError{Code: "InvalidKeyPair.NotFound"},
			failure:     UserManagedEC2KeyFailureNotFound,
			remediation: "create or select",
		},
		{
			name:        "permission denied",
			err:         &smithy.GenericAPIError{Code: "UnauthorizedOperation"},
			failure:     UserManagedEC2KeyFailurePermission,
			remediation: "grant ec2:DescribeKeyPairs",
		},
		{
			name:        "malformed remote public key",
			publicKey:   "ssh-dss AAAA test-comment",
			failure:     UserManagedEC2KeyFailureInvalid,
			remediation: "valid RSA or ED25519 public key",
		},
		{
			name:        "mismatch",
			publicKey:   strings.TrimSpace(string(otherPublic)),
			failure:     UserManagedEC2KeyFailureInvalid,
			remediation: "matching existing EC2 key pair",
		},
	} {
		t.Run(testCase.name, func(t *testing.T) {
			fixture := &userManagedKeyPairFixture{err: testCase.err, publicKey: testCase.publicKey}
			_, err := validateUserManagedEC2Key(
				context.Background(), fixture, "123456789012", "us-west-2", "caller-key", privatePath, "",
			)
			if err == nil {
				t.Fatal("negative validation unexpectedly succeeded")
			}
			if fixture.calls != 1 {
				t.Fatalf("provider calls=%d, want exactly one read-only DescribeKeyPairs call", fixture.calls)
			}
			diagnostic := DiagnoseUserManagedEC2Key(err)
			if diagnostic.FailureClass != testCase.failure || !strings.Contains(diagnostic.Remediation, testCase.remediation) {
				t.Fatalf("diagnostic = %#v, err=%v", diagnostic, err)
			}
			for _, secret := range []string{strings.TrimSpace(string(otherPublic)), "ssh-dss AAAA"} {
				if strings.Contains(err.Error(), secret) {
					t.Fatalf("error exposed public key material %q: %v", secret, err)
				}
			}
			for _, want := range []string{"account=123456789012", "region=us-west-2", "key_pair=\"caller-key\"", "path=\"" + privatePath + "\""} {
				if !strings.Contains(err.Error(), want) {
					t.Fatalf("error %q missing context %q", err, want)
				}
			}
		})
	}
}

func TestIntegration_UserManagedEC2KeyValidationDiagnosesMissingAndDeniedKeys(t *testing.T) {
	privatePath := generateUserManagedKey(t, "ed25519")
	for _, testCase := range []struct {
		name        string
		code        string
		failure     string
		field       string
		remediation string
	}{
		{
			name:        "not found",
			code:        "InvalidKeyPair.NotFound",
			failure:     UserManagedEC2KeyFailureNotFound,
			field:       "host.ec2.key_pair",
			remediation: "existing EC2 key pair",
		},
		{
			name:        "permission denied",
			code:        "UnauthorizedOperation",
			failure:     UserManagedEC2KeyFailurePermission,
			field:       "target.aws.iam.role_arn",
			remediation: "ec2:DescribeKeyPairs",
		},
	} {
		t.Run(testCase.name, func(t *testing.T) {
			_, err := validateUserManagedEC2Key(
				context.Background(),
				&userManagedKeyPairFixture{err: &smithy.GenericAPIError{Code: testCase.code}},
				"123456789012",
				"us-west-2",
				"caller-key",
				privatePath,
				"",
			)
			if err == nil {
				t.Fatal("validation unexpectedly succeeded")
			}
			diagnostic := DiagnoseUserManagedEC2Key(err)
			if diagnostic.FailureClass != testCase.failure || diagnostic.Field != testCase.field ||
				!strings.Contains(diagnostic.Remediation, testCase.remediation) {
				t.Fatalf("diagnostic = %#v, err=%v", diagnostic, err)
			}
		})
	}
}

func TestIntegration_UserManagedEC2KeyValidationRequiresAbsoluteRegularKey(t *testing.T) {
	privatePath := generateUserManagedKey(t, "ed25519")
	for _, path := range []string{
		"relative-key",
		filepath.Dir(privatePath),
		filepath.Join(t.TempDir(), "missing-key"),
	} {
		_, err := validateUserManagedEC2Key(
			context.Background(),
			&userManagedKeyPairFixture{},
			"123456789012",
			"us-west-2",
			"caller-key",
			path,
			"",
		)
		if err == nil || DiagnoseUserManagedEC2Key(err).Field != "host.ec2.ssh_private_key" {
			t.Fatalf("path %q validation error = %v", path, err)
		}
	}
}

func TestIntegration_UserManagedEC2KeyValidationDiagnosesLocalKeyFailures(t *testing.T) {
	for _, testCase := range []struct {
		name        string
		path        func(*testing.T) string
		failure     string
		remediation string
	}{
		{
			name: "missing path",
			path: func(t *testing.T) string {
				return filepath.Join(t.TempDir(), "missing-key")
			},
			failure:     UserManagedEC2KeyFailureInvalid,
			remediation: "existing absolute readable regular RSA or ED25519 private-key file",
		},
		{
			name: "unreadable path",
			path: func(t *testing.T) string {
				if os.Geteuid() == 0 {
					t.Skip("root can read mode-000 files")
				}
				path := filepath.Join(t.TempDir(), "unreadable-key")
				if err := os.WriteFile(path, []byte("not an SSH private key\n"), 0o600); err != nil {
					t.Fatal(err)
				}
				if err := os.Chmod(path, 0o000); err != nil {
					t.Fatal(err)
				}
				t.Cleanup(func() { _ = os.Chmod(path, 0o600) })
				return path
			},
			failure:     UserManagedEC2KeyFailureInvalid,
			remediation: "existing absolute readable regular RSA or ED25519 private-key file",
		},
		{
			name: "malformed key",
			path: func(t *testing.T) string {
				path := filepath.Join(t.TempDir(), "malformed-key")
				if err := os.WriteFile(path, []byte("not an SSH private key\n"), 0o600); err != nil {
					t.Fatal(err)
				}
				return path
			},
			failure:     UserManagedEC2KeyFailureInvalid,
			remediation: "valid RSA or ED25519 private-key file",
		},
		{
			name: "encrypted key",
			path: func(t *testing.T) string {
				path := filepath.Join(t.TempDir(), "encrypted-key")
				if output, err := exec.Command("ssh-keygen", "-q", "-t", "ed25519", "-N", "test-passphrase", "-f", path).CombinedOutput(); err != nil {
					t.Fatalf("generate encrypted key: %v\n%s", err, output)
				}
				return path
			},
			failure:     UserManagedEC2KeyFailureInvalid,
			remediation: "unencrypted RSA or ED25519 private-key file",
		},
	} {
		t.Run(testCase.name, func(t *testing.T) {
			path := testCase.path(t)
			_, err := validateUserManagedEC2Key(
				context.Background(), &userManagedKeyPairFixture{}, "123456789012", "us-west-2", "caller-key", path, "",
			)
			if err == nil {
				t.Fatal("local key validation unexpectedly succeeded")
			}
			diagnostic := DiagnoseUserManagedEC2Key(err)
			if diagnostic.FailureClass != testCase.failure || !strings.Contains(diagnostic.Remediation, testCase.remediation) {
				t.Fatalf("diagnostic = %#v, err=%v", diagnostic, err)
			}
			for _, want := range []string{"account=123456789012", "region=us-west-2", "key_pair=\"caller-key\"", "path=\"" + path + "\""} {
				if !strings.Contains(err.Error(), want) {
					t.Fatalf("error %q missing context %q", err, want)
				}
			}
		})
	}
}

func TestIntegration_UserManagedEC2KeyValidationMissingSSHKeygenIsActionable(t *testing.T) {
	privatePath := generateUserManagedKey(t, "ed25519")
	t.Setenv("PATH", t.TempDir())
	fixture := &userManagedKeyPairFixture{}
	_, err := validateUserManagedEC2Key(
		context.Background(), fixture, "123456789012", "us-west-2", "caller-key", privatePath, "",
	)
	if err == nil || !errors.Is(err, errUserManagedSSHKeygen) || fixture.calls != 0 {
		t.Fatalf("missing ssh-keygen validation = calls=%d err=%v", fixture.calls, err)
	}
	diagnostic := DiagnoseUserManagedEC2Key(err)
	if diagnostic.FailureClass != UserManagedEC2KeyFailureUnavailable ||
		!strings.Contains(diagnostic.Remediation, "install OpenSSH ssh-keygen") {
		t.Fatalf("missing ssh-keygen diagnostic = %#v", diagnostic)
	}
}

func TestIntegration_UserManagedEC2KeyValidationSkipsAWSForPersistedInstance(t *testing.T) {
	privatePath := generateUserManagedKey(t, "ed25519")
	result, err := ValidateUserManagedEC2Key(
		context.Background(),
		nil,
		"123456789012",
		"us-west-2",
		"caller-key",
		privatePath,
		"i-existing",
	)
	if err != nil || result == nil || !result.EqualitySkipped || result.Account != "123456789012" {
		t.Fatalf("persisted-instance validation = %#v, err=%v", result, err)
	}
}

func generateUserManagedKey(t *testing.T, algorithm string) string {
	t.Helper()
	privatePath := filepath.Join(t.TempDir(), "id_"+algorithm)
	args := []string{"-q", "-t", algorithm, "-N", "", "-f", privatePath}
	if algorithm == "rsa" {
		args = append(args[:len(args)-2], "-b", "2048")
		args = append(args, "-f", privatePath)
	}
	if output, err := exec.Command("ssh-keygen", args...).CombinedOutput(); err != nil {
		t.Fatalf("generate %s key: %v\n%s", algorithm, err, output)
	}
	return privatePath
}