//go:build integration

package doctor

import (
	"context"
	"strings"
	"testing"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"

	cpaws "github.com/tidbcloud/controlplane9/internal/aws"
	"github.com/tidbcloud/controlplane9/internal/config"
	"github.com/tidbcloud/controlplane9/internal/host"
	"github.com/tidbcloud/controlplane9/internal/host/remoteimage"
)

func TestIntegration_CheckSessionIAMSurfacesSafeUserManagedKeyResult(t *testing.T) {
	session := &config.SessionConfig{
		APIVersion: config.APIVersionV1Alpha2,
		Name:       "doctor-key-result",
		Host: config.SessionHostConfig{
			Mode: config.HostModeRemoteEC2K3d,
			EC2: config.HostEC2Config{
				InstanceType:  "m7i.xlarge",
				KeyPair:       "caller-key",
				SSHPrivateKey: "/absolute/private-key",
			},
		},
		Target: config.SessionTargetConfig{
			Region: "us-west-2",
			AWS: config.SessionAWSConfig{IAM: config.SessionIAMConfig{
				SourceProfile: "source",
				RoleARN:       "arn:aws:iam::123456789012:role/cp9-developer",
			}},
		},
	}

	oldLoad := doctorLoadAWSConfig
	oldIAM := doctorValidateIAM
	oldCleanup := doctorValidateIAMCleanup
	oldKey := doctorValidateUserManagedEC2Key
	oldImage := doctorResolveRemoteImagePlan
	t.Cleanup(func() {
		doctorLoadAWSConfig = oldLoad
		doctorValidateIAM = oldIAM
		doctorValidateIAMCleanup = oldCleanup
		doctorValidateUserManagedEC2Key = oldKey
		doctorResolveRemoteImagePlan = oldImage
	})

	doctorLoadAWSConfig = func(context.Context, string, string) (awssdk.Config, error) {
		return awssdk.Config{Region: session.Target.Region}, nil
	}
	doctorValidateIAM = func(context.Context, cpaws.STSAPI, string, string, string, string, map[string][]string) (*cpaws.SetupResult, error) {
		return &cpaws.SetupResult{
			Valid:   true,
			Account: "123456789012",
			Owner:   "owner",
		}, nil
	}
	doctorValidateIAMCleanup = func(context.Context, *awssdk.Config, *cpaws.SetupResult, string) error {
		return nil
	}
	doctorValidateUserManagedEC2Key = func(_ context.Context, _ *awssdk.Config, account, region, keyPair, _, _ string) (*cpaws.UserManagedEC2KeyValidation, error) {
		return &cpaws.UserManagedEC2KeyValidation{
			KeyPair:         keyPair,
			Algorithm:       "ssh-ed25519",
			PublicKeySHA256: "SHA256:example",
			Account:         account,
			Region:          region,
		}, nil
	}
	doctorResolveRemoteImagePlan = func(context.Context, *awssdk.Config, string, string) (*remoteimage.Plan, error) {
		return &remoteimage.Plan{
			Family:       remoteimage.FamilyAmazonLinux2023,
			InstanceType: session.Host.EC2.InstanceType,
			Architecture: "x86_64",
			AMI:          "ami-12345678",
			Region:       session.Target.Region,
		}, nil
	}

	result := &Result{}
	CheckSessionIAM(result, session, &host.Session{Mode: config.HostModeRemoteEC2K3d}, nil)
	if result.UserManagedEC2Key == nil || result.UserManagedEC2Key.PublicKeySHA256 != "SHA256:example" {
		t.Fatalf("doctor key result = %#v", result.UserManagedEC2Key)
	}
	found := false
	for _, check := range result.Checks {
		if check.Name == "aws-key-pair" {
			found = true
			if check.Level != LevelPass || !strings.Contains(check.Message, "public_key_sha256=SHA256:example") {
				t.Fatalf("doctor key check = %#v", check)
			}
		}
	}
	if !found {
		t.Fatalf("doctor checks did not include aws-key-pair: %#v", result.Checks)
	}
}