package e2e

import (
	"context"
	crand "crypto/rand"
	"encoding/json"
	"fmt"
	"math/rand"
	"time"

	_ "github.com/go-sql-driver/mysql"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	"github.com/pingcap/errors"
	branchapi "github.com/tidbcloud/aws-shared-provider/pkg/controllers/branch"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/pkg/consts"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/util/retry"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

var _ = Describe("branch test api in vip pool", func() {
	interval := 10 * time.Second
	ctx := context.Background()

	It("branch test depend on api only", func() {
		By("Create a new serverless cluster")
		cluster := newCommonCluster()
		salt, err := generateSalt()
		Expect(err).Should(Succeed())
		metadata := &utils.ClusterMetadata{
			ID:          0,
			Name:        "",
			TenantID:    0,
			ProjectID:   0,
			CreatedAt:   time.Now(),
			CreatedByID: 0,
			Salt:        salt,
		}
		j, err := json.Marshal(metadata)
		Expect(err).Should(Succeed())
		cluster.Annotations[utils.TiDBClusterMetadataAnnKey] = string(j)
		cluster.Spec.TiDB.PoolSelector = vipTidbPool
		Expect(infraClusterCli.Create(ctx, cluster)).Should(Succeed())
		waitClusterReady(ctx, cluster, 5*time.Minute, interval)

		By("Create a new branch")
		branch := newBranch()
		branch.Annotations[branchapi.ParentBranchIdAnnKey] = cluster.Name
		branch.Annotations[branchapi.ParentTypeAnnKey] = branchapi.ParentTypeCluster
		branch.Annotations[types.AnnoTidbPoolSelector] = vipTidbPool
		Expect(infraCli.Create(ctx, branch)).Should(Succeed())
		waitBranchReady(ctx, branch, 10*time.Minute, interval)
		branchStatus, err := branch.AwsFreeTierClusterStatusFromBranch()
		Expect(err).Should(Succeed(), branchStatus.KeyspaceToken)

		By("Create a new branch from branch")
		childBranch := newBranch()
		childBranch.Annotations[branchapi.ParentBranchIdAnnKey] = branch.Name
		childBranch.Annotations[branchapi.ParentTypeAnnKey] = branchapi.ParentTypeBranch
		Expect(infraCli.Create(ctx, childBranch)).Should(Succeed(), branchStatus.KeyspaceToken)
		waitBranchReady(ctx, childBranch, 10*time.Minute, interval)

		By("Change child branch to vip pool")
		childBranch.Annotations[types.AnnoTidbPoolSelector] = vipTidbPool
		Expect(infraCli.Update(ctx, childBranch)).Should(Succeed(), branchStatus.KeyspaceToken)
		waitBranchPoolChanged(ctx, childBranch, 1*time.Minute, interval)

		By("Test update tidb two replicas")
		err = retry.RetryOnConflict(retry.DefaultRetry, func() error {
			err := infraCli.Get(ctx, client.ObjectKey{Namespace: branch.Namespace, Name: branch.Name}, branch)
			if err != nil {
				return err
			}
			branch.Spec.Cluster.TiDB.Replicas = 2
			err = infraCli.Update(ctx, branch)
			return err
		})
		Expect(err).Should(Succeed(), branchStatus.KeyspaceToken)
		waitBranchSynced(ctx, branch, 1*time.Minute, interval)

		By("Check if manager assign two tidb pod eventually")
		Expect(endpointConnected(ctx, branchStatus.KeyspaceToken,
			SqlEndpointStatus2PublicStatus(*branch.Status.SqlEndpoints.PublicEndpoint))).
			Should(Succeed(), branchStatus.KeyspaceToken)
		Eventually(func() error {
			return checkTokenInfoFromManager(branchStatus.KeyspaceToken, vipTidbPool, func(t *tokenInfo) error {
				if t.Replicas != branch.Spec.Cluster.TiDB.Replicas {
					return errors.Errorf("branch %v replicas not equal", branch.Name)
				}
				if len(t.Endpoints) != branch.Spec.Cluster.TiDB.Replicas {
					return errors.Errorf("branch %v replicas not equal endpoints", branch.Name)
				}
				return nil
			})
		}, 1*time.Minute, 3*time.Second).Should(Succeed(), branchStatus.KeyspaceToken)

		By("Test update tidb scale-in to one replica")
		err = retry.RetryOnConflict(retry.DefaultRetry, func() error {
			err := infraCli.Get(ctx, client.ObjectKey{Namespace: branch.Namespace, Name: branch.Name}, branch)
			if err != nil {
				return err
			}
			branch.Spec.Cluster.TiDB.Replicas = 1
			err = infraCli.Update(ctx, branch)
			return err
		})
		Expect(err).Should(Succeed(), branchStatus.KeyspaceToken)
		waitBranchSynced(ctx, branch, 1*time.Minute, interval)

		By("Check if manager assign back to one tidb pod")
		Expect(endpointConnected(ctx, branchStatus.KeyspaceToken,
			SqlEndpointStatus2PublicStatus(*branch.Status.SqlEndpoints.PublicEndpoint))).
			Should(Succeed(), branchStatus.KeyspaceToken)
		Expect(checkTokenInfoFromManager(branchStatus.KeyspaceToken, vipTidbPool, func(t *tokenInfo) error {
			if t.Replicas != branch.Spec.Cluster.TiDB.Replicas {
				return errors.Errorf("branch %v replicas not equal", branch.Name)
			}
			if len(t.Endpoints) != branch.Spec.Cluster.TiDB.Replicas {
				return errors.Errorf("branch %v replicas not equal endpoints", branch.Name)
			}
			return nil
		})).Should(Succeed(), branchStatus.KeyspaceToken)

		By("wait assigned tidb pod to free")
		Eventually(func() error {
			return checkTokenInfoFromManager(branchStatus.KeyspaceToken, vipTidbPool, func(t *tokenInfo) error {
				if t.Replicas != branch.Spec.Cluster.TiDB.Replicas {
					return errors.Errorf("branch %v replicas not equal", branch.Name)
				}
				if len(t.Endpoints) > 0 {
					logger.Infow("endpoints not empty", "branch", branch.Name, "endpoints", t.Endpoints)
					return errors.Errorf("branch %v endpoints not empty", branch.Name)
				}
				return nil
			})
		}, 10*time.Minute, time.Minute).Should(Succeed(), branchStatus.KeyspaceToken)

		By("check mode is standard")
		Expect(branchStatus.RuLimitMode).Should(Equal("standard"))

		By("change ru limit mode to throttled")
		branch.Annotations["free-tier.tidbcloud.com/ru-limit-mode"] = throttledMode
		Expect(infraCli.Update(ctx, branch)).Should(Succeed(), branchStatus.KeyspaceToken)
		Eventually(func() error {
			key := client.ObjectKey{Namespace: branch.Namespace, Name: branch.Name}
			if err := infraCli.Get(ctx, key, branch); err != nil {
				return err
			}
			branchStatus, err = branch.AwsFreeTierClusterStatusFromBranch()
			if err != nil {
				return err
			}
			if branchStatus.RuLimitMode != throttledMode {
				return errors.Errorf("branch %s ru limit mode is not throttled", branch.Name)
			}
			return nil
		}).WithTimeout(5*time.Minute).WithPolling(interval).Should(Succeed(), branchStatus.KeyspaceToken)

		By("delete childBranch and wait it to be deleted")
		Expect(infraCli.Delete(ctx, childBranch)).Should(Succeed(), branchStatus.KeyspaceToken)
		waitBranchDeleted(ctx, childBranch, 5*time.Minute, interval)

		By("delete branch and wait it to be deleted")
		Expect(infraCli.Delete(ctx, branch)).Should(Succeed(), branchStatus.KeyspaceToken)
		waitBranchDeleted(ctx, branch, 5*time.Minute, interval)

		By("delete cluster and wait it to be deleted")
		Expect(infraClusterCli.Delete(ctx, cluster)).Should(Succeed(), branchStatus.KeyspaceToken)
		waitClusterDeleted(ctx, cluster, 5*time.Minute, interval)
	})
})

func newBranch() *cloudv1.Branch {
	instanceType := "aws_freetier_10"
	e2eLabels := e2eResourceLabels()
	e2eLabels[consts.TenantTagKey] = "e2e-tenant-id"
	e2eLabels[cloudv1.ProviderTypeLabel] = "aws-free-tier"
	e2eLabels[cloudv1.RegionLabel] = awsRegion
	meta := &utils.BranchMetadata{
		ID:          "",
		Name:        "",
		TenantID:    0,
		ProjectID:   0,
		CreatedAt:   time.Now(),
		CreatedByID: 0,
		Salt:        []byte("testsalt"),
	}
	metaJson, _ := json.Marshal(meta)

	return &cloudv1.Branch{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: namespace,
			Name:      fmt.Sprintf("e2e-branch-%s-%d", kubernetesLabel, rand.Int31()),
			Labels:    e2eLabels,
			Annotations: map[string]string{
				utils.TiDBBranchMetadataAnnKey:  string(metaJson),
				branchapi.ParentTimestampAnnKey: time.Now().UTC().Format(time.RFC3339),
			},
		},
		Spec: cloudv1.BranchSpec{
			SqlEndpoints: cloudv1.BranchSqlEndpointsSpec{
				Public: cloudv1.BranchPublicSpec{
					Service: cloudv1.Service{
						Enabled: ptr.To(true),
						// open for all access
						LoadBalancerSourceRanges: []string{
							"0.0.0.0/0",
						},
					},
					AuthorizedNetworks: []cloudv1.AuthorizedNetwork{
						{
							DisplayName:    "default",
							StartIpAddress: "0.0.0.0",
							EndIpAddress:   "255.255.255.255",
						},
					},
				},
			},
			Cluster: cloudv1.ClusterTopology{
				TiDB: cloudv1.TiDB{
					Replicas:     1,
					InstanceType: instanceType,
				},
				PD: cloudv1.PD{
					Replicas:     1,
					InstanceType: instanceType,
				},
				TiKV: cloudv1.TiKV{
					Replicas:     1,
					InstanceType: instanceType,
				},
				TiFlash: cloudv1.TiFlash{
					Replicas:     1,
					InstanceType: instanceType,
				},
			},
		},
	}
}

// waitBranchSynced wait status synced with spec when you update replicas.
func waitBranchSynced(ctx context.Context, branch *cloudv1.Branch, timeout, internal time.Duration) {
	time.Sleep(3 * time.Second)
	Eventually(func() error {
		err := infraCli.Get(ctx, client.ObjectKey{Namespace: branch.Namespace, Name: branch.Name}, branch)
		if err != nil {
			logger.Errorw("error inspecting branch", "branch", branch.Name, "error", err.Error())
			return err
		}
		if branch.Status.Phase != cloudv1.BranchPhaseRunning {
			return errWait
		}
		if branch.Status.DB == nil || branch.Status.DB.TiDB == nil ||
			branch.Status.DB.TiDB.UpdatedReplicas != int32(branch.Spec.Cluster.TiDB.Replicas) {
			return errWait
		}
		clusterStatus, err := branch.AwsFreeTierClusterStatusFromBranch()
		if err != nil {
			return err
		}
		if branch.Annotations[types.AnnoTidbPoolSelector] != "" &&
			branch.Annotations[types.AnnoTidbPoolSelector] != clusterStatus.TidbPool {
			return errWait
		}
		return nil
	}, timeout, internal).Should(
		Succeed(),
		"Branch should be synced anno: %v, spec: %v, status: %v", branch.Annotations, branch.Spec.Cluster, branch.Status.DB,
	)
}

func waitBranchPoolChanged(ctx context.Context, branch *cloudv1.Branch, timeout, internal time.Duration) {
	time.Sleep(3 * time.Second)
	Eventually(func() error {
		err := infraCli.Get(ctx, client.ObjectKey{Namespace: branch.Namespace, Name: branch.Name}, branch)
		if err != nil {
			logger.Errorw("error inspecting branch", "branch", branch.Name, "error", err.Error())
			return err
		}
		if branch.Annotations[types.AnnoKeyTidbPoolChanging] != "" {
			return errWait
		}
		clusterStatus, err := branch.AwsFreeTierClusterStatusFromBranch()
		if err != nil {
			return err
		}
		if branch.Annotations[types.AnnoTidbPoolSelector] != "" &&
			branch.Annotations[types.AnnoTidbPoolSelector] != clusterStatus.TidbPool {
			return errWait
		}
		return nil
	}, timeout, internal).Should(
		Succeed(),
		"Branch AnnoKeyTidbPoolChanging should be removed: %v, status: %v", branch.Annotations, branch.Status,
	)
}

func waitBranchReady(ctx context.Context, branch *cloudv1.Branch, timeout, interval time.Duration) {
	Eventually(func() error {
		err := infraCli.Get(ctx, client.ObjectKey{Namespace: branch.Namespace, Name: branch.Name}, branch)
		if err != nil {
			logger.Errorw("error inspecting branch status", "branch", branch.Name, "error", err.Error())
			return err
		}
		if branch.Status.Phase != cloudv1.BranchPhaseRunning {
			logger.Infow("wait branch ready",
				"branch", branch.Name, "current phase", branch.Status.Phase)
			return errWait
		}
		return nil
	}, timeout, interval).Should(Succeed(), "Branch should be ready")
}

func waitBranchDeleted(ctx context.Context, branch *cloudv1.Branch, timeout, interval time.Duration) {
	Eventually(func() error {
		err := infraCli.Get(ctx, client.ObjectKey{Namespace: branch.Namespace, Name: branch.Name}, branch)
		if err != nil {
			if apierrors.IsNotFound(err) {
				return nil
			}
			logger.Errorw("error inspecting branch status", "branch", branch.Name, "error", err.Error())
			return err
		}
		logger.Infow("wait branch to be cleaned", "branch", branch.Name)
		return errWait
	}, timeout, interval).Should(Succeed(), "Branch should be cleaned up")
}

func generateSalt() ([]byte, error) {
	salt := make([]byte, 8)
	if _, e := crand.Read(salt); e != nil {
		return nil, e
	}
	return salt, nil
}