package backup

import (
	"context"
	"encoding/json"
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/iam"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	ststypes "github.com/aws/aws-sdk-go-v2/service/sts/types"
	"github.com/golang/mock/gomock"
	. "github.com/onsi/gomega"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"github.com/tidbcloud/config-service/sdk/api"
	"github.com/tidbcloud/config-service/sdk/mock"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/controllers/test"
	"github.com/tidbcloud/infra-api/pkg/consts"
	awsapi "github.com/tidbcloud/infra-api/pkg/providers/aws"
	provideraws "github.com/tidbcloud/infra-api/pkg/providers/aws"
	"github.com/tidbcloud/infra-api/pkg/providers/aws/fake"
	gcpfake "github.com/tidbcloud/infra-api/pkg/providers/gcp/fake"
)

func TestCreateRestore(t *testing.T) {
	g := NewGomegaWithT(t)

	mockCli := &test.MockClient{
		MockGet: test.NewMockGetFn(nil),
		MockUpdate: test.MockUpdateFn(func(ctx context.Context, obj client.Object, opts ...client.UpdateOption) error {
			mr, _ := obj.(*cloudv1.Restore)
			status := mr.Status.Phase
			g.Expect(status).Should(Equal(cloudv1.RestorePrepare))

			return nil
		}),
		MockCreate: test.MockCreateFn(func(ctx context.Context, obj client.Object, opts ...client.CreateOption) error {
			return nil
		}),
		MockSubresourceUpdate: test.NewMockSubresourceUpdateFn(nil),
		MockSubresourcePatch:  test.NewMockSubresourcePatchFn(nil),
	}

	mockClusterCli := &test.MockClient{
		MockGet: test.NewMockGetFn(nil),
		MockCreate: func(ctx context.Context, obj client.Object, opts ...client.CreateOption) error {
			tidbRestore, _ := obj.(*v1alpha1.Restore)

			g.Expect(tidbRestore.Spec.S3.Region).Should(Equal("us-west-2"))
			g.Expect(tidbRestore.Spec.S3.Bucket).Should(Equal("bucket1"))
			g.Expect(tidbRestore.Spec.S3.Prefix).Should(Equal("backup1"))
			g.Expect(tidbRestore.ObjectMeta.Namespace).Should(Equal("tidb1"))
			g.Expect(tidbRestore.Spec.Mode).Should(Equal(v1alpha1.RestoreModeSnapshot))

			return nil
		},
	}

	status := &cloudv1.GcpClusterStatus{
		Kubernetes: &cloudv1.KubernetesStatus{
			Namespace: "test",
		},
	}
	data, _ := json.Marshal(status)
	mode := cloudv1.BackupModeSnapshot

	cr := &clusterRestore{
		cli:        mockCli,
		clusterCli: mockClusterCli,
		network: &cloudv1.Network{
			Spec: cloudv1.NetworkSpec{
				Region: "us-west-2",
			},
			Status: cloudv1.NetworkStatus{
				FromProvider: runtime.RawExtension{
					Object: &cloudv1.GardenerNetworkStatus{
						ShootName: "shoot",
					},
				},
			},
		},
		cluster: &cloudv1.Cluster{
			ObjectMeta: metav1.ObjectMeta{
				Labels: map[string]string{
					cloudv1.ProviderTypeLabel: consts.ProviderTypeGCP,
				},
			},
			Status: cloudv1.ClusterStatus{
				ObjectStore: &cloudv1.ObjectStoreStatus{
					S3: &cloudv1.Path{
						Path: "test-restore",
					},
				},
				FromProvider: runtime.RawExtension{
					Raw: data,
				},
			},
		},
		provider: &cloudv1.Provider{
			Spec: cloudv1.ProviderSpec{
				ForProvider: &runtime.RawExtension{
					Object: &cloudv1.GardenerProviderSpec{
						GCP: cloudv1.GcpConfig{
							ProjectID: "projectid",
						},
					},
				},
			},
		},
		backup: &cloudv1.Backup{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "backup1",
				Namespace: "infra-test-1",
			},
			Spec: cloudv1.BackupSpec{
				Mode:        &mode,
				ClusterName: "tidb1",
			},
			Status: cloudv1.BackupStatus{
				BucketName: "bucket1",
			},
		},
		k8sStatus: status.Kubernetes,
		configSvc: &mock.MockClient{
			RestoreConfigOutput: &api.DescribeBRConfigOutput{
				Version: "v4.0.9",
				Config:  map[string]interface{}{"concurrency": float64(2)},
				Image: api.Image{
					BaseImage: "gcr.io/pingcap-public/dbaas/br",
					Tag:       pointer.String("v4.0.9-hotfix"),
				},
			},
		},
	}

	restore := &cloudv1.Restore{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "restore1",
			Namespace: "infra-test-1",
		},
		Spec: cloudv1.RestoreSpec{
			ClusterName: "tidb1",
			From: cloudv1.From{
				Backup: "backup1",
			},
		},
	}
	f := restore.GetFinalizers()
	restore.SetFinalizers(append(f, RestoreFinalizer))

	err := cr.Create(context.TODO(), restore)
	g.Expect(err).Should(Succeed())
	// TODO: should we add more assertions here?
}

func TestCreatePiTR(t *testing.T) {
	g := NewGomegaWithT(t)
	const (
		region         = "us-west-2"
		backupBucket   = "bucket1"
		backupName     = "backup1"
		logBackupName  = "log_backup1"
		restoredTS     = "ts1"
		tidbNamespace  = "tidb1"
		infraNamespace = "infra-test-1"
		clusterName    = "tidb1"
	)

	mockCli := &test.MockClient{
		MockGet: test.NewMockGetFn(nil),
		MockUpdate: test.MockUpdateFn(func(ctx context.Context, obj client.Object, opts ...client.UpdateOption) error {
			mr, ok := obj.(*cloudv1.Restore)
			g.Expect(ok).Should(BeTrue())
			status := mr.Status.Phase
			g.Expect(status).Should(Equal(cloudv1.RestorePrepare))

			return nil
		}),
		MockCreate: test.MockCreateFn(func(ctx context.Context, obj client.Object, opts ...client.CreateOption) error {
			return nil
		}),
		MockSubresourceUpdate: test.NewMockSubresourceUpdateFn(nil),
		MockSubresourcePatch:  test.NewMockSubresourcePatchFn(nil),
	}

	mockClusterCli := &test.MockClient{
		MockGet: test.NewMockGetFn(nil),
		MockCreate: func(ctx context.Context, obj client.Object, opts ...client.CreateOption) error {
			tidbRestore, ok := obj.(*v1alpha1.Restore)
			g.Expect(ok).Should(BeTrue())
			g.Expect(tidbRestore.Spec.Mode).Should(Equal(v1alpha1.RestoreModePiTR))
			g.Expect(tidbRestore.Spec.S3.Region).Should(Equal(region))
			g.Expect(tidbRestore.Spec.S3.Bucket).Should(Equal(backupBucket))
			g.Expect(tidbRestore.Spec.S3.Prefix).Should(Equal(logBackupName))
			g.Expect(tidbRestore.ObjectMeta.Namespace).Should(Equal(tidbNamespace))

			g.Expect(tidbRestore.Spec.PitrFullBackupStorageProvider.S3.Region).Should(Equal(region))
			g.Expect(tidbRestore.Spec.PitrFullBackupStorageProvider.S3.Bucket).Should(Equal(backupBucket))
			g.Expect(tidbRestore.Spec.PitrFullBackupStorageProvider.S3.Prefix).Should(Equal(backupName))
			g.Expect(tidbRestore.Spec.PitrRestoredTs).Should(Equal(restoredTS))
			return nil
		},
	}

	status := &cloudv1.GcpClusterStatus{
		Kubernetes: &cloudv1.KubernetesStatus{
			Namespace: "test",
		},
	}
	data, _ := json.Marshal(status)
	mode := cloudv1.BackupModeSnapshot

	cr := &clusterRestore{
		cli:        mockCli,
		clusterCli: mockClusterCli,
		network: &cloudv1.Network{
			Spec: cloudv1.NetworkSpec{
				Region: region,
			},
			Status: cloudv1.NetworkStatus{
				FromProvider: runtime.RawExtension{
					Object: &cloudv1.GardenerNetworkStatus{
						ShootName: "shoot",
					},
				},
			},
		},
		cluster: &cloudv1.Cluster{
			ObjectMeta: metav1.ObjectMeta{
				Labels: map[string]string{
					cloudv1.ProviderTypeLabel: consts.ProviderTypeGCP,
				},
			},
			Status: cloudv1.ClusterStatus{
				ObjectStore: &cloudv1.ObjectStoreStatus{
					S3: &cloudv1.Path{
						Path: "test-restore",
					},
				},
				FromProvider: runtime.RawExtension{
					Raw: data,
				},
			},
		},
		provider: &cloudv1.Provider{
			Spec: cloudv1.ProviderSpec{
				ForProvider: &runtime.RawExtension{
					Object: &cloudv1.GardenerProviderSpec{
						GCP: cloudv1.GcpConfig{
							ProjectID: "projectid",
						},
					},
				},
			},
		},
		backup: &cloudv1.Backup{
			ObjectMeta: metav1.ObjectMeta{
				Name:      backupName,
				Namespace: infraNamespace,
			},
			Spec: cloudv1.BackupSpec{
				Mode:        &mode,
				ClusterName: clusterName,
			},
			Status: cloudv1.BackupStatus{
				BucketName: backupBucket,
			},
		},
		logBackup: &cloudv1.Backup{
			ObjectMeta: metav1.ObjectMeta{
				Name:      logBackupName,
				Namespace: infraNamespace,
			},
			Spec: cloudv1.BackupSpec{
				ClusterName: clusterName,
			},
			Status: cloudv1.BackupStatus{
				BucketName: backupBucket,
			},
		},
		k8sStatus: status.Kubernetes,
		configSvc: &mock.MockClient{
			RestoreConfigOutput: &api.DescribeBRConfigOutput{
				Version: "v4.0.9",
				Config:  map[string]interface{}{"concurrency": float64(2)},
				Image: api.Image{
					BaseImage: "gcr.io/pingcap-public/dbaas/br",
					Tag:       pointer.String("v4.0.9-hotfix"),
				},
			},
		},
	}

	logBackupNameStr := logBackupName
	restoredTSStr := restoredTS
	restore := &cloudv1.Restore{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "restore1",
			Namespace: infraNamespace,
		},
		Spec: cloudv1.RestoreSpec{
			ClusterName: clusterName,
			From: cloudv1.From{
				Backup:         backupName,
				LogBackup:      &logBackupNameStr,
				PiTRRestoredTS: &restoredTSStr,
			},
		},
	}
	f := restore.GetFinalizers()
	restore.SetFinalizers(append(f, RestoreFinalizer))

	err := cr.Create(context.TODO(), restore)
	g.Expect(err).Should(Succeed())
}

func TestAttachS3BucketPermissionToNodeRole(t *testing.T) {
	g := NewGomegaWithT(t)
	awsCli := fake.NewMockClient(t)
	ctx := context.TODO()
	bucketName := "my-s3-bucket"
	roleName := "shoot--staging--12345678-nodes"

	pd := provideraws.PolicyDocument{
		Version: provideraws.PolicyVersion,
		Statement: []provideraws.PolicyStatement{
			{
				Effect:   "Allow",
				Action:   []string{"s3:*"},
				Resource: []string{fmt.Sprintf("arn:aws:s3:::%s", bucketName), fmt.Sprintf("arn:aws:s3:::%s/*", bucketName)},
			},
		},
	}
	pdStr, err := pd.ToString()
	g.Expect(err).ShouldNot(HaveOccurred())

	policyInput := &iam.PutRolePolicyInput{
		PolicyDocument: aws.String(pdStr),
		PolicyName:     aws.String(bucketName),
		RoleName:       aws.String(roleName),
	}

	awsCli.MockIAMClient.EXPECT().PutRolePolicy(ctx, policyInput).Times(1).Return(nil, nil)
	err = attachS3BucketPermissionToNodeRole(ctx, awsCli, "my-s3-bucket", fmt.Sprintf("arn:aws:iam::12345678:role/%s", roleName))
	g.Expect(err).ShouldNot(HaveOccurred())
}

func TestAttachGCPStorageAdminPermissionToSA(t *testing.T) {
	g := NewGomegaWithT(t)
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()

	gcpapi := gcpfake.NewMockGCPAPI(mockCtrl)
	nodeRole := "gcp-test@pingcap-gardener.iam.gserviceaccount.com"
	bucket := "dbaas-staging-us-west1-12345678"

	gcpapi.EXPECT().AddBucketIAMMember(bucket, storageAdminRoleName, fmt.Sprintf("serviceAccount:%s", nodeRole))
	err := attachGCPStorageAdminPermissionToSA(gcpapi, nodeRole, bucket)
	g.Expect(err).ShouldNot(HaveOccurred())
}

func TestUpdateNodeRoleAssumePolicy(t *testing.T) {
	var (
		g        = NewGomegaWithT(t)
		awsCli   = fake.NewMockClient(t)
		ctx      = context.TODO()
		nodeRole = "my-role"
		trusts   = []string{"arn:aws:iam::123456:user/tom"}
	)

	pd := &awsapi.PolicyDocument{
		Version: "2012-10-17",
		Statement: []awsapi.PolicyStatement{
			{
				Effect: awsapi.StatementEffectAllow,
				Principal: &awsapi.PolicyPrincipal{
					Service: "ec2.amazonaws.com",
					AWS:     trusts,
				},
				Action: []string{"sts:AssumeRole"},
			},
		},
	}
	pdStr, err := pd.ToString()
	g.Expect(err).ShouldNot(HaveOccurred())

	policyInput := &iam.UpdateAssumeRolePolicyInput{
		PolicyDocument: aws.String(pdStr),
		RoleName:       aws.String(nodeRole),
	}

	awsCli.MockIAMClient.EXPECT().UpdateAssumeRolePolicy(ctx, policyInput).Times(1).Return(nil, nil)
	err = updateNodeRoleAssumePolicy(ctx, awsCli, nodeRole, trusts)
	g.Expect(err).ShouldNot(HaveOccurred())
}

func TestAssumeRole(t *testing.T) {
	var (
		g          = NewGomegaWithT(t)
		awsCli     = fake.NewMockClient(t)
		ctx        = context.TODO()
		nodeRole   = "arn:aws:iam::123456:role/nodes"
		bucketName = "dbaas-nightly-us-west-2-1"
	)

	sessionName := fmt.Sprintf("cross-region-%s", bucketName)
	input := &sts.AssumeRoleInput{
		RoleArn:         &nodeRole,
		RoleSessionName: &sessionName,
	}
	cred := &ststypes.Credentials{
		AccessKeyId: aws.String("test"),
	}

	output := &sts.AssumeRoleOutput{
		Credentials: cred,
	}
	awsCli.MockStsClient.EXPECT().AssumeRole(ctx, input).Times(1).Return(output, nil)
	res, err := assumeRole(ctx, awsCli, bucketName, nodeRole)
	g.Expect(err).ShouldNot(HaveOccurred())
	g.Expect(*res.AccessKeyId).Should(Equal("test"))
}