package e2e

import (
	"context"
	"database/sql"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	s3types "github.com/aws/aws-sdk-go-v2/service/s3/types"
	_ "github.com/go-sql-driver/mysql"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
)

var _ = PDescribe("Import Backup and Restore test", func() {
	interval := 10 * time.Second
	instanceType := "t3.medium"
	It("Should take import backup and restore for AWS provider successfully", func() {
		By("create old cluster")
		oldCluster := newCommonCluster(e2eEnv.AwsNetwork.Name, instanceType)
		oldCluster.Spec.Cluster.PD.Replicas = 1
		oldCluster.Spec.Cluster.TiDB.Replicas = 1
		oldCluster.Spec.Cluster.TiFlash.Replicas = 0
		Expect(infraCli.Create(ctx, oldCluster)).Should(Succeed())

		By("create new cluster")
		newCluster := newCommonCluster(e2eEnv.AwsNetwork.Name, instanceType)
		newCluster.Spec.Cluster.PD.Replicas = 1
		newCluster.Spec.Cluster.TiDB.Replicas = 1
		newCluster.Spec.Cluster.TiFlash.Replicas = 0
		Expect(infraCli.Create(ctx, newCluster)).Should(Succeed())

		By("wait network ready")
		networkStatus, err := cloudv1.GardenerNetworkStatusFromNetwork(e2eEnv.AwsNetwork)
		Expect(err).Should(Succeed())

		By("wait old cluster ready")
		waitClusterReady(ctx, oldCluster, 25*time.Minute, interval)
		Expect(oldCluster.Status.SqlEndpoints.PublicEndpoint).NotTo(BeNil())

		By("wait new cluster ready")
		waitClusterReady(ctx, newCluster, 25*time.Minute, interval)
		Expect(newCluster.Status.SqlEndpoints.PublicEndpoint).NotTo(BeNil())

		By("wait old TiDB ready") // should wait TiDB to be ready in order to execute SQL statements
		waitTiDBReady(ctx, oldCluster, 20*time.Minute, interval)

		By("wait new TiDB ready")
		waitTiDBReady(ctx, newCluster, 20*time.Minute, interval)

		By("prepare data for import")
		region := e2eEnv.AwsNetwork.Spec.Region
		bucket := networkStatus.BucketName
		keyPrefix := "import-sample-data"
		_, currFile, _, _ := runtime.Caller(0)
		dataDir := filepath.Join(filepath.Dir(currFile), "assets", "import-sample-data")
		// Using the SDK's default configuration, loading additional config
		// and credentials values from the environment variables, shared
		// credentials, and shared configuration files
		cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
		Expect(err).Should(Succeed())
		s3Cli := s3.NewFromConfig(cfg)
		objs, err := uploadImportSampleDataToS3(ctx, s3Cli, bucket, keyPrefix, dataDir)
		defer func() {
			// delete any uploaded objects
			Expect(deleteImportSampleDataInS3(ctx, s3Cli, bucket, objs)).Should(Succeed())
		}()
		Expect(err).Should(Succeed())

		By("create import")
		cloudImport := &cloudv1.Import{
			ObjectMeta: metav1.ObjectMeta{
				Namespace: e2eEnv.Namespace,
				Name:      cloudv1.RandomString(6),
			},
			Spec: cloudv1.ImportSpec{
				ClusterName:    oldCluster.Name,
				StorageRequest: "200Mi",
				Source: cloudv1.Source{
					Format:   cloudv1.DataFormatSQL,
					UserName: "cloud_admin",
					Storage: cloudv1.Storage{
						S3: &cloudv1.S3{
							Region: region,
							URL:    fmt.Sprintf("s3://%s/%s/", bucket, keyPrefix),
						},
					},
				},
			},
		}
		Expect(infraCli.Create(ctx, cloudImport)).Should(Succeed())

		By("wait import complete")
		Eventually(func() error {
			err := infraCli.Get(ctx, types.NamespacedName{Namespace: cloudImport.Namespace, Name: cloudImport.Name}, cloudImport)
			if err != nil {
				logger.Errorw("error inspecting import status", "import", cloudImport.Name, "error", err.Error())
				return err
			}
			logger.Infow("dump import status", "status", cloudImport.Status)
			if cloudImport.Status.Phase != cloudv1.ImportCompleted {
				logger.Infow("wait import complete", "import", cloudImport.Name, "current phase", cloudImport.Status.Phase)
				return errWait
			}
			return nil
		}, 30*time.Minute, interval).Should(Succeed(), "Import should complete")
		Expect(cloudImport.Status.CompletedTables).Should(Equal(uint32(1)))
		Expect(cloudImport.Status.PendingTables).Should(Equal(uint32(0)))
		Expect(cloudImport.Status.TotalSize > uint64(0)).To(BeTrue())
		Expect(cloudImport.Status.CurrentTables).Should(BeEmpty())
		Expect(cloudImport.Status.CompletedPercent > uint32(0)).To(BeTrue()) // this may not be 100 but 96, I don't know the reason
		Expect(*cloudImport.Status.TotalFiles).Should(Equal(uint32(1)))

		By("Verify data in old cluster for import")
		verifyImportSampleData(*oldCluster.Status.SqlEndpoints.PublicEndpoint)

		By("delete import")
		Expect(infraCli.Delete(ctx, cloudImport)).Should(Succeed())

		By("wait import to be deleted")
		waitImportDeleted(ctx, cloudImport, 10*time.Minute, interval)

		By("create backup")
		backup := &cloudv1.Backup{
			ObjectMeta: metav1.ObjectMeta{
				Namespace: e2eEnv.Namespace,
				Name:      cloudv1.RandomString(6),
			},
			Spec: cloudv1.BackupSpec{
				ClusterName: oldCluster.Name,
			},
		}
		Expect(infraCli.Create(ctx, backup)).Should(Succeed())

		By("wait backup complete")
		Eventually(func() error {
			err := infraCli.Get(ctx, types.NamespacedName{Namespace: backup.Namespace, Name: backup.Name}, backup)
			if err != nil {
				logger.Errorw("error inspecting backup status", "backup", backup.Name, "error", err.Error())
				return err
			}
			if backup.Status.Phase != cloudv1.BackupComplete {
				logger.Infow("wait backup complete", "backup", backup.Name, "current phase", backup.Status.Phase)
				return errWait
			}
			return nil
		}, 15*time.Minute, interval).Should(Succeed(), "Backup should complete")
		Expect(backup.Status.BackupSize > int64(0)).To(BeTrue())
		Expect(backup.Status.BackupTS).NotTo(BeEmpty())
		Expect(backup.Status.BucketName).NotTo(BeEmpty())

		By("create restore")
		restore := &cloudv1.Restore{
			ObjectMeta: metav1.ObjectMeta{
				Namespace: e2eEnv.Namespace,
				Name:      cloudv1.RandomString(6),
			},
			Spec: cloudv1.RestoreSpec{
				ClusterName: newCluster.Name,
				From: cloudv1.From{
					Backup: backup.Name,
				},
			},
		}
		Expect(infraCli.Create(ctx, restore)).Should(Succeed())

		By("wait restore complete")
		Eventually(func() error {
			err := infraCli.Get(ctx, types.NamespacedName{Namespace: restore.Namespace, Name: restore.Name}, restore)
			if err != nil {
				logger.Errorw("error inspecting restore status", "restore", restore.Name, "error", err.Error())
				return err
			}
			if restore.Status.Phase != cloudv1.RestoreComplete {
				logger.Infow("wait restore complete", "restore", restore.Name, "current phase", restore.Status.Phase)
				return errWait
			}
			return nil
		}, 15*time.Minute, interval).Should(Succeed(), "Restore should complete")

		By("Verify data in new cluster for restore")
		verifyImportSampleData(*newCluster.Status.SqlEndpoints.PublicEndpoint)

		By("delete backup and restore")
		Expect(infraCli.Delete(ctx, backup)).Should(Succeed())
		Expect(infraCli.Delete(ctx, restore)).Should(Succeed())

		By("wait backup and restore to be deleted")
		waitBackupDeleted(ctx, backup, 5*time.Minute, interval)
		waitRestoreDeleted(ctx, restore, 5*time.Minute, interval)

		By("delete old and new clusters")
		Expect(infraCli.Delete(ctx, newCluster)).Should(Succeed())
		Expect(infraCli.Delete(ctx, oldCluster)).Should(Succeed())

		By("wait old and new cluster to be deleted")
		waitClusterDeleted(ctx, oldCluster, 5*time.Minute, interval)
		waitClusterDeleted(ctx, newCluster, 5*time.Minute, interval)
	})
})

// uploadImportSampleDataToS3 uploads the import sample data into S3 and return the imported objects.
func uploadImportSampleDataToS3(ctx context.Context, s3Cli *s3.Client, bucket, keyPrefix, dataDir string) ([]string, error) {
	var objs []string
	logger.Infow("uploading import sample data to S3", "bucket", bucket, "prefix", keyPrefix)

	fileInfos, err := ioutil.ReadDir(dataDir)
	if err != nil {
		return objs, err
	}

	for _, info := range fileInfos {
		if info.IsDir() {
			continue
		}

		path := filepath.Join(dataDir, info.Name())
		f, err2 := os.Open(path)
		if err2 != nil {
			return objs, err2
		}

		logger.Infow("uploading data to S3", "file", path)
		obj := filepath.Join(keyPrefix, info.Name())
		_, err2 = s3Cli.PutObject(ctx, &s3.PutObjectInput{
			Bucket: &bucket,
			Key:    &obj,
			Body:   f,
		})
		_ = f.Close()
		if err2 != nil {
			return objs, err2
		}
		objs = append(objs, obj)
	}
	return objs, nil
}

// deleteImportSampleDataInS3 deletes the data in S3 which uploaded by `uploadImportSampleDataToS3`.
func deleteImportSampleDataInS3(ctx context.Context, s3Cli *s3.Client, bucket string, objs []string) error {
	if len(objs) == 0 {
		return nil
	}

	logger.Infow("deleting import sample data in S3", "bucket", bucket, "objects", objs)
	var ids []s3types.ObjectIdentifier
	for _, obj := range objs {
		obj2 := obj
		ids = append(ids, s3types.ObjectIdentifier{
			Key: &obj2,
		})
	}
	_, err := s3Cli.DeleteObjects(ctx, &s3.DeleteObjectsInput{
		Bucket: &bucket,
		Delete: &s3types.Delete{
			Objects: ids,
		},
	})
	return err
}

// verifyImportSampleData verify import sample data for a TidbCluster.
func verifyImportSampleData(endpoint cloudv1.SqlEndpointStatus) {
	dsn := DSN(endpoint)
	db, err := sql.Open("mysql", dsn)
	Expect(err).ShouldNot(HaveOccurred())
	row := db.QueryRowContext(ctx, "SELECT `c1` FROM `e2e`.`import`")
	var got int
	Expect(row.Scan(&got)).Should(Succeed())
	Expect(got).Should(Equal(123)) // column value in sample data is `123`
	Expect(db.Close()).Should(Succeed())
}