//go:build integration

package cmd

import (
	"context"
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"testing"

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

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

func TestIntegration_UpResumeWithPersistedInstanceValidatesLocalKeyWithoutAWS(t *testing.T) {
	repo := newKeyValidationGitRepo(t)
	t.Chdir(repo)
	session := integrationRemoteKeyValidationSession("resume-with-host")
	stateDir := filepath.Join(repo, cluster.DefaultStateDir, session.Name)
	lifecycleState := cluster.NewStateForLifecycle(session.Name, session.Name, cluster.LifecycleV3)
	if err := cluster.SaveState(stateDir, lifecycleState); err != nil {
		t.Fatal(err)
	}
	if err := host.Save(stateDir, &host.Session{
		Mode:     config.HostModeRemoteEC2K3d,
		StateDir: stateDir,
		Remote:   &host.RemoteSession{InstanceID: "i-existing", Region: session.Target.Region},
	}); err != nil {
		t.Fatal(err)
	}

	var calls []string
	restoreKeyValidationCommandHooks(t, func() {
		cmdLoadAWSConfig = func(context.Context, string, string) (awssdk.Config, error) {
			t.Fatal("resume with persisted instance loaded AWS config")
			return awssdk.Config{}, nil
		}
		cmdRunV1Alpha2SetupIAMPreflight = func(context.Context, *awssdk.Config, string, string, string, string, map[string][]string) (*cpaws.SetupResult, awssdk.Config, error) {
			t.Fatal("resume with persisted instance ran IAM preflight")
			return nil, awssdk.Config{}, nil
		}
		cmdValidateUserManagedEC2Key = func(_ context.Context, cfg *awssdk.Config, account, region, keyPair, privateKeyPath, instanceID string) (*cpaws.UserManagedEC2KeyValidation, error) {
			calls = append(calls, "key")
			if cfg != nil || account != "" || region != session.Target.Region || keyPair != session.Host.EC2.KeyPair ||
				privateKeyPath != session.Host.EC2.SSHPrivateKey || instanceID != "i-existing" {
				t.Fatalf("persisted-instance validation args = cfg=%#v account=%q region=%q key=%q path=%q instance=%q", cfg, account, region, keyPair, privateKeyPath, instanceID)
			}
			return &cpaws.UserManagedEC2KeyValidation{KeyPair: keyPair, Region: region, EqualitySkipped: true}, nil
		}
	})

	inventory, err := runV1Alpha2UpSessionCollisionPreflight(context.Background(), session)
	if err != nil {
		t.Fatal(err)
	}
	if inventory.Status != "resume" || !reflect.DeepEqual(calls, []string{"key"}) {
		t.Fatalf("resume collision result = %#v, validation calls = %v", inventory, calls)
	}
}

func TestIntegration_ResumeStateWithoutHostStillRunsAWSKeyEquality(t *testing.T) {
	repo := newKeyValidationGitRepo(t)
	t.Chdir(repo)
	session := integrationRemoteKeyValidationSession("resume-without-host")
	stateDir := filepath.Join(repo, cluster.DefaultStateDir, session.Name)
	lifecycleState := cluster.NewStateForLifecycle(session.Name, session.Name, cluster.LifecycleV3)
	if err := cluster.SaveState(stateDir, lifecycleState); err != nil {
		t.Fatal(err)
	}

	var calls []string
	restoreKeyValidationCommandHooks(t, func() {
		cmdLoadAWSConfig = func(context.Context, string, string) (awssdk.Config, error) {
			calls = append(calls, "load")
			return awssdk.Config{Region: session.Target.Region}, nil
		}
		cmdRunV1Alpha2SetupIAMPreflight = func(_ context.Context, _ *awssdk.Config, _ string, _ string, _ string, _ string, _ map[string][]string) (*cpaws.SetupResult, awssdk.Config, error) {
			calls = append(calls, "iam")
			return &cpaws.SetupResult{Account: "123456789012", Owner: "owner"}, awssdk.Config{Region: session.Target.Region}, nil
		}
		cmdValidateUserManagedEC2Key = func(_ context.Context, cfg *awssdk.Config, account, region, _, _, instanceID string) (*cpaws.UserManagedEC2KeyValidation, error) {
			calls = append(calls, "key")
			if cfg == nil || account != "123456789012" || region != session.Target.Region || instanceID != "" {
				t.Fatalf("state-without-host validation args = cfg=%#v account=%q region=%q instance=%q", cfg, account, region, instanceID)
			}
			return &cpaws.UserManagedEC2KeyValidation{Account: account, Region: region}, nil
		}
	})

	inventory, err := runV1Alpha2UpSessionCollisionPreflight(context.Background(), session)
	if err != nil {
		t.Fatal(err)
	}
	if inventory.Status != "resume" || !reflect.DeepEqual(calls, []string{"load", "iam", "key"}) {
		t.Fatalf("state-without-host result = %#v, calls = %v", inventory, calls)
	}
}

func TestIntegration_FreshUpValidatesKeyBeforeCollisionScan(t *testing.T) {
	repo := newKeyValidationGitRepo(t)
	t.Chdir(repo)
	session := integrationRemoteKeyValidationSession("fresh-key-before-scan")
	var calls []string
	restoreKeyValidationCommandHooks(t, func() {
		cmdLoadAWSConfig = func(context.Context, string, string) (awssdk.Config, error) {
			calls = append(calls, "load")
			return awssdk.Config{Region: session.Target.Region}, nil
		}
		cmdRunV1Alpha2SetupIAMPreflight = func(_ context.Context, _ *awssdk.Config, _ string, _ string, _ string, _ string, _ map[string][]string) (*cpaws.SetupResult, awssdk.Config, error) {
			calls = append(calls, "iam")
			return &cpaws.SetupResult{Account: "123456789012", Owner: "owner"}, awssdk.Config{Region: session.Target.Region}, nil
		}
		cmdValidateUserManagedEC2Key = func(_ context.Context, cfg *awssdk.Config, account, region, _, _, instanceID string) (*cpaws.UserManagedEC2KeyValidation, error) {
			calls = append(calls, "key")
			if cfg == nil || account != "123456789012" || region != session.Target.Region || instanceID != "" {
				t.Fatalf("fresh validation args = cfg=%#v account=%q region=%q instance=%q", cfg, account, region, instanceID)
			}
			return &cpaws.UserManagedEC2KeyValidation{Account: account, Region: region}, nil
		}
		cmdRunV1Alpha2SessionCollisionPreflight = func(_ context.Context, _ *config.SessionConfig, _ string, _ bool) (*cpaws.SessionCollisionInventory, error) {
			calls = append(calls, "collision")
			return &cpaws.SessionCollisionInventory{Status: "clear"}, nil
		}
	})

	inventory, err := runV1Alpha2UpSessionCollisionPreflight(context.Background(), session)
	if err != nil {
		t.Fatal(err)
	}
	if inventory.Status != "clear" || !reflect.DeepEqual(calls, []string{"load", "iam", "key", "collision"}) {
		t.Fatalf("fresh result = %#v, calls = %v", inventory, calls)
	}
}

func TestIntegration_FreshUpKeyFailureStopsBeforeCollisionOrLifecycle(t *testing.T) {
	repo := newKeyValidationGitRepo(t)
	t.Chdir(repo)
	session := integrationRemoteKeyValidationSession("fresh-key-failure")
	var calls []string
	restoreKeyValidationCommandHooks(t, func() {
		cmdLoadAWSConfig = func(context.Context, string, string) (awssdk.Config, error) {
			calls = append(calls, "load")
			return awssdk.Config{Region: session.Target.Region}, nil
		}
		cmdRunV1Alpha2SetupIAMPreflight = func(_ context.Context, _ *awssdk.Config, _ string, _ string, _ string, _ string, _ map[string][]string) (*cpaws.SetupResult, awssdk.Config, error) {
			calls = append(calls, "iam")
			return &cpaws.SetupResult{Account: "123456789012", Owner: "owner"}, awssdk.Config{Region: session.Target.Region}, nil
		}
		cmdValidateUserManagedEC2Key = func(context.Context, *awssdk.Config, string, string, string, string, string) (*cpaws.UserManagedEC2KeyValidation, error) {
			calls = append(calls, "key")
			return nil, errors.New("caller-owned key validation failed: account=123456789012 region=us-west-2 key_pair=\"caller-key\" path=\"/absolute/private-key\"")
		}
		cmdRunV1Alpha2SessionCollisionPreflight = func(context.Context, *config.SessionConfig, string, bool) (*cpaws.SessionCollisionInventory, error) {
			t.Fatal("collision scan ran after key validation failed")
			return nil, nil
		}
	})

	_, err := runV1Alpha2UpSessionCollisionPreflight(context.Background(), session)
	if err == nil {
		t.Fatal("key validation unexpectedly succeeded")
	}
	if !reflect.DeepEqual(calls, []string{"load", "iam", "key"}) {
		t.Fatalf("key-failure call sequence = %v; collision/lifecycle must remain untouched", calls)
	}
}

func TestIntegration_ReadOnlyPreflightCarriesSafeKeyResult(t *testing.T) {
	session := integrationRemoteKeyValidationSession("read-only-result")
	var calls []string
	restoreKeyValidationCommandHooks(t, func() {
		cmdLoadAWSConfig = func(context.Context, string, string) (awssdk.Config, error) {
			calls = append(calls, "load")
			return awssdk.Config{Region: session.Target.Region}, nil
		}
		cmdRunV1Alpha2SetupIAMPreflight = func(_ context.Context, _ *awssdk.Config, _ string, _ string, _ string, _ string, _ map[string][]string) (*cpaws.SetupResult, awssdk.Config, error) {
			calls = append(calls, "iam")
			return &cpaws.SetupResult{Valid: true, Account: "123456789012", Owner: "owner"}, awssdk.Config{Region: session.Target.Region}, nil
		}
		cmdValidateIAMCleanup = func(context.Context, *awssdk.Config, *cpaws.SetupResult, string) error {
			calls = append(calls, "cleanup")
			return nil
		}
		cmdValidateUserManagedEC2Key = func(_ context.Context, _ *awssdk.Config, account, region, keyPair, _, _ string) (*cpaws.UserManagedEC2KeyValidation, error) {
			calls = append(calls, "key")
			return &cpaws.UserManagedEC2KeyValidation{
				KeyPair:         keyPair,
				Algorithm:       "ssh-ed25519",
				PublicKeySHA256: "SHA256:example",
				Account:         account,
				Region:          region,
			}, nil
		}
		cmdRunV1Alpha2SessionCollisionPreflight = func(_ context.Context, _ *config.SessionConfig, _ string, _ bool) (*cpaws.SessionCollisionInventory, error) {
			calls = append(calls, "collision")
			return &cpaws.SessionCollisionInventory{Status: "clear"}, nil
		}
		cmdResolveV1Alpha2RemoteImagePlan = func(context.Context, *awssdk.Config, *cpaws.SetupResult, *config.SessionConfig) error {
			calls = append(calls, "image")
			return nil
		}
	})

	result, err := runV1Alpha2ReadOnlyPreflight(context.Background(), session)
	if err != nil {
		t.Fatal(err)
	}
	if result.UserManagedEC2Key == nil || result.UserManagedEC2Key.PublicKeySHA256 != "SHA256:example" ||
		!reflect.DeepEqual(calls, []string{"load", "iam", "cleanup", "key", "collision", "image"}) {
		t.Fatalf("read-only result = %#v, calls = %v", result, calls)
	}
	setupText := result.FormatText()
	if !strings.Contains(setupText, "public_key_sha256=SHA256:example") ||
		!strings.Contains(setupText, "key_pair=caller-key") {
		t.Fatalf("setup text omitted safe key result: %s", setupText)
	}
	setupJSON, err := json.Marshal(result)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(setupJSON), `"user_managed_ec2_key"`) ||
		!strings.Contains(string(setupJSON), `"public_key_sha256":"SHA256:example"`) {
		t.Fatalf("setup JSON omitted safe key result: %s", setupJSON)
	}
	dryRun := &upDryRunResult{
		DryRun:            true,
		Session:           session.Name,
		UserManagedEC2Key: result.UserManagedEC2Key,
	}
	dryRunText := dryRun.FormatText()
	if !strings.Contains(dryRunText, "public_key_sha256=SHA256:example") ||
		!strings.Contains(dryRunText, "key_pair=caller-key") {
		t.Fatalf("dry-run text omitted safe key result: %s", dryRunText)
	}
	dryRunJSON, err := json.Marshal(dryRun)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(dryRunJSON), `"user_managed_ec2_key"`) ||
		!strings.Contains(string(dryRunJSON), `"public_key_sha256":"SHA256:example"`) {
		t.Fatalf("dry-run JSON omitted safe key result: %s", dryRunJSON)
	}
}

func restoreKeyValidationCommandHooks(t *testing.T, install func()) {
	t.Helper()
	oldLoad := cmdLoadAWSConfig
	oldSetup := cmdRunV1Alpha2SetupIAMPreflight
	oldKey := cmdValidateUserManagedEC2Key
	oldCollision := cmdRunV1Alpha2SessionCollisionPreflight
	oldImage := cmdResolveV1Alpha2RemoteImagePlan
	oldCleanup := cmdValidateIAMCleanup
	t.Cleanup(func() {
		cmdLoadAWSConfig = oldLoad
		cmdRunV1Alpha2SetupIAMPreflight = oldSetup
		cmdValidateUserManagedEC2Key = oldKey
		cmdRunV1Alpha2SessionCollisionPreflight = oldCollision
		cmdResolveV1Alpha2RemoteImagePlan = oldImage
		cmdValidateIAMCleanup = oldCleanup
	})
	install()
}

func integrationRemoteKeyValidationSession(name string) *config.SessionConfig {
	return &config.SessionConfig{
		APIVersion: config.APIVersionV1Alpha2,
		Name:       name,
		Host: config.SessionHostConfig{
			Mode: config.HostModeRemoteEC2K3d,
			EC2: config.HostEC2Config{
				KeyPair:       "caller-key",
				SSHPrivateKey: "/absolute/private-key",
			},
		},
		Target: config.SessionTargetConfig{
			Cloud:  "aws",
			Region: "us-west-2",
			AWS: config.SessionAWSConfig{IAM: config.SessionIAMConfig{
				SourceProfile: "source",
				RoleARN:       "arn:aws:iam::123456789012:role/cp9-developer",
			}},
		},
	}
}

func newKeyValidationGitRepo(t *testing.T) string {
	t.Helper()
	repo := t.TempDir()
	git(t, repo, "init")
	git(t, repo, "config", "user.name", "CP9 Test")
	git(t, repo, "config", "user.email", "cp9@example.com")
	if err := os.WriteFile(filepath.Join(repo, "README"), []byte("test\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	git(t, repo, "add", "README")
	git(t, repo, "commit", "-m", "test")
	return repo
}