package common

import (
	"context"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"go.uber.org/multierr"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	"github.com/tidbcloud/eks-provider/controllers/utils"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/pkg/namer"
	apiutil "github.com/tidbcloud/infra-provider/pkg/util"
)

// Resolve ... TODO(aylei): we can make the observation more responsive via registering a controller to shoot api-server
func (t *TiDBGroup) ObserveExtraGroupClusterStatus(ctx context.Context, tc *v1alpha1.TidbCluster, extraGroup *cloudv1.ExtraGroupCluster) error {
	var errs error
	tidb, err := t.computeTiDBStatus(ctx, tc, extraGroup)
	if err != nil {
		errs = multierr.Append(errs, err)
	} else if tidb == nil {
		errs = multierr.Append(errs, fmt.Errorf("extra tidb status is nil"))
		return errs
	}
	tiproxy, err := t.computeTiProxyStatus(ctx, tc, extraGroup)
	if err != nil {
		errs = multierr.Append(errs, err)
	} else if tiproxy == nil {
		errs = multierr.Append(errs, fmt.Errorf("extra tiproxy status is nil"))
		return errs
	}
	if errs != nil {
		return errs
	}

	// secondaryStatus
	secondaryStatus := cloudv1.DBSecondaryStatusHealthy

	if tiproxy.Status == cloudv1.ComponentUnhealthy ||
		tidb.Status == cloudv1.ComponentUnhealthy {
		secondaryStatus = cloudv1.DBSecondaryStatusUnhealthy
	}

	// Phase
	var phase cloudv1.ClusterPhase
	switch {
	// cluster is considered as unavailable when one of TiKV, PD, TiDB, TiProxy is unavailable
	case tidb.Status == cloudv1.ComponentUnavailable || tiproxy.Status == cloudv1.ComponentUnavailable:
		if extraGroup.Status.Phase != cloudv1.ClusterPhaseSuspend {
			phase = cloudv1.ClusterPhaseUnavailable
			secondaryStatus = cloudv1.DBSecondaryStatusUnhealthy
		} else {
			phase = cloudv1.ClusterPhaseSuspend
		}
	// if both upgrading and scaling are ongoing (e.g. failover when upgrading), show Upgrading
	case tidb.Status == cloudv1.ComponentUpgrading || tiproxy.Status == cloudv1.ComponentUpgrading:
		phase = cloudv1.ClusterPhaseUpgrading
	case tidb.Status == cloudv1.ComponentScaling || tiproxy.Status == cloudv1.ComponentScaling:
		phase = cloudv1.ClusterPhaseScaling
	default:
		phase = cloudv1.ClusterPhaseRunning
	}

	// Resons
	reason := fmt.Sprintf("tidb: %s, tiproxy: %s", tidb.Status, tiproxy.Status)

	// Condition
	if extraGroup.Status.Conditions == nil {
		extraGroup.Status.Conditions = []cloudv1.ClusterCondition{}
	}
	cloudv1.SetStatusCondition(&extraGroup.Status.Conditions, t.buildSuspendCondition(tidb, tiproxy))

	if t.Cluster.Suspending() || t.Cluster.Suspended() {
		phase = cloudv1.ClusterPhaseSuspend
	}

	extraGroup.Status.Phase = phase
	extraGroup.Status.Reason = reason
	extraGroup.Status.TiDB = tidb
	extraGroup.Status.TiProxy = tiproxy
	extraGroup.Status.SecondaryStatus = secondaryStatus
	extraGroup.Status.Version = tc.Spec.Version

	if extraGroup.Status.Name == "" {
		extraGroup.Status.Name = extraGroup.Spec.Name
	}

	return errs
}

func (t *TiDBGroup) buildSuspendCondition(tidbStatus *cloudv1.TiDBStatus, tiproxyStatus *cloudv1.TiProxyStatus) cloudv1.ClusterCondition {
	suspendCondition := cloudv1.ClusterCondition{
		Type:    cloudv1.ClusterSuspended,
		Status:  cloudv1.ConditionTrue,
		Reason:  "ClusterSuspended",
		Message: "The cluster has been suspended",
	}
	var msg []string
	if tidbStatus == nil {
		msg = append(msg, "extragroup tidb status is nil")
	} else if !meta.IsStatusConditionTrue(tidbStatus.Conditions,
		cloudv1.ComponentSuspended) {
		msg = append(msg, "TiDB has not been suspended")
	}
	if tiproxyStatus == nil {
		msg = append(msg, "extragroup tiproxy status is nil")
	} else if !meta.IsStatusConditionTrue(tiproxyStatus.Conditions,
		cloudv1.ComponentSuspended) {
		msg = append(msg, "TiProxy has not been suspended")
	}
	if len(msg) > 0 {
		suspendCondition = cloudv1.ClusterCondition{
			Type:    cloudv1.ClusterSuspended,
			Status:  cloudv1.ConditionFalse,
			Reason:  "ClusterNotSuspended",
			Message: strings.Join(msg, ";"),
		}
	}
	return suspendCondition
}

func (t *TiDBGroup) computeComponentConditions(ctx context.Context, cType cloudv1.ComponentType, nodes []corev1.Node, tc *v1alpha1.TidbCluster, extraGroup *cloudv1.ExtraGroupCluster) []metav1.Condition {
	var conditions []metav1.Condition
	var suspendCondition metav1.Condition
	var needSuspend bool
	var component cloudv1.ClusterComponent
	memberType := utils.GetMemberType(cType)

	switch cType {
	case cloudv1.ComponentTiDB:
		if extraGroup.Status.TiDB != nil {
			conditions = extraGroup.Status.TiDB.Conditions
		}
		component = &extraGroup.Spec.TiDB
	case cloudv1.ComponentTiProxy:
		if extraGroup.Status.TiProxy != nil {
			conditions = extraGroup.Status.TiProxy.Conditions
		}
		component = &extraGroup.Spec.TiProxy
	default:
		return nil
	}

	if conditions == nil {
		conditions = []metav1.Condition{}
	}
	needSuspend = component.NeedSuspend(t.Cluster.Spec.Cluster.SuspendAction)
	// The flow for the condition update is as below:
	// - If `suspendAction` is set in the spec, add the `ComponentSuspended`
	//   condition first to indicate that we're going to suspend the cluster
	// - Update TC with the `suspendAction` in the next reconcile,
	//   in case that the user may resume the cluster immediately
	//   after pausing the cluster
	// - Update `ComponentSuspended` condition to `True` after the cluster is suspended
	// - Remove the `ComponentSuspended` condition when the user resumes the cluster
	memberSuspending := tc.ComponentIsSuspending(memberType)
	if !needSuspend && !memberSuspending {
		meta.RemoveStatusCondition(&conditions, cloudv1.ComponentSuspended)
		return conditions
	}
	memberSuspended := tc.ComponentIsSuspended(memberType)
	nodeCount := len(nodes)

	if memberSuspended && nodeCount == 0 {
		suspendCondition = metav1.Condition{
			Type:    cloudv1.ComponentSuspended,
			Status:  metav1.ConditionTrue,
			Reason:  "ResourcesSuspended",
			Message: "All Pods and nodes are suspended",
		}
	} else {
		var msg []string
		if !memberSuspended {
			msg = append(msg, "Pods not suspended")
		}
		if nodeCount != 0 {
			msg = append(msg, fmt.Sprintf("%d Nodes not suspended", nodeCount))
		}
		suspendCondition = metav1.Condition{
			Type:    cloudv1.ComponentSuspended,
			Status:  metav1.ConditionFalse,
			Reason:  "ResourcesNotSuspended",
			Message: strings.Join(msg, ";"),
		}
	}
	meta.SetStatusCondition(&conditions, suspendCondition)
	return conditions
}

// selectComponentNodes selects nodes for the given component type
// For AWS self-managed nodegroup, after pausing the cluster, the nodes will still be schedulable for a few minutes.
// Refer to https://docs.aws.amazon.com/autoscaling/ec2/userguide/lifecycle-hooks-overview.html, the `Terminating:Wait` state will last for a few minutes.
// So we need to make sure that the nodes are deleted before we update the `ComponentSuspended` condition to `True`.
func (t *TiDBGroup) selectComponentNodes(ctx context.Context, componentType cloudv1.ComponentType, nodesOfPod []corev1.Node, tc *v1alpha1.TidbCluster) ([]corev1.Node, error) {
	// If the cluster is not AWS self-managed nodegroup, use nodes of pods directly.
	if apiutil.GetAnnotation(t.Cluster, consts.UseSelfManagedNodegroupAnnotation) != consts.AnnoValTrue {
		return nodesOfPod, nil
	}

	nodes, err := k8s.SelectNodes(ctx, t.EnvProvider.GetClient(), selectorOfTiDBClusterPods(tc, string(componentType)))
	if err != nil {
		return nil, err
	}

	nodesMap := make(map[string]struct{})
	for _, node := range nodes {
		nodesMap[node.Name] = struct{}{}
	}

	for _, node := range nodesOfPod {
		if _, ok := nodesMap[node.Name]; !ok {
			nodes = append(nodes, node)
		}
	}

	return nodes, nil
}

func (t *TiDBGroup) computeTiDBStatus(ctx context.Context, tc *v1alpha1.TidbCluster, extraGroup *cloudv1.ExtraGroupCluster) (*cloudv1.TiDBStatus, error) {
	logger := log.FromContext(ctx)
	// Continue to refresh status if TiDB is suspended
	suspend := tc.ComponentIsSuspended(v1alpha1.TiDBMemberType)
	// Skip refresh status if status in tc is outdated
	// like STS recreate.
	if !suspend &&
		len(tc.Status.TiDB.Members) == 0 &&
		extraGroup.Status.TiDB != nil {
		return extraGroup.Status.TiDB, nil
	}

	var members []cloudv1.TiDBMember
	instanceSpec, err := t.Provider.FindInstanceType(extraGroup.Spec.TiDB.GetInstanceType())
	if err != nil {
		return nil, err
	}
	tidbNodeScaling, err := t.componentNodeScaling(ctx, cloudv1.ComponentTiDB, tc)
	if err != nil {
		return nil, err
	}
	healthyCount := 0

	pods, err := k8s.SelectPods(ctx, t.EnvProvider.GetClient(), tc.Namespace, selectorOfTiDBClusterPods(tc, string(cloudv1.ComponentTiDB)))
	if err != nil {
		return nil, err
	}

	nodesOfPod := make([]corev1.Node, 0)
	for _, p := range pods {
		if len(p.Spec.NodeName) < 1 {
			continue
		}
		node := corev1.Node{}
		err := t.EnvProvider.GetClient().Get(ctx, client.ObjectKey{
			Name: p.Spec.NodeName,
		}, &node)
		if err != nil {
			return nil, err
		}
		nodesOfPod = append(nodesOfPod, node)
	}

	// For suspend condition update, consider both
	nodes, err := t.selectComponentNodes(ctx, cloudv1.ComponentTiDB, nodesOfPod, tc)
	if err != nil {
		return nil, err
	}
	conditions := t.computeComponentConditions(ctx, cloudv1.ComponentTiDB, nodes, tc, extraGroup)

	tidbNodeCountScaling := t.isNodeScalingCheckEnabled() && extraGroup.Spec.TiDB.Replicas != len(nodes)
	if tidbNodeCountScaling {
		logger.Info("TiDB node scaling", "expected", extraGroup.Spec.TiDB.Replicas, "current", len(nodes))
	}

	var nameToPod = make(map[string]corev1.Pod, len(pods))
	for _, pod := range pods {
		nameToPod[pod.Name] = pod
	}

	var previousInstanceTypeMap map[string]string
	var previousMemberStatusMap map[string]cloudv1.MemberStatus
	if extraGroup.Status.TiDB != nil {
		previousInstanceTypeMap = tidbMemberToInstanceType(extraGroup.Status.TiDB.Members)
		previousMemberStatusMap = tidbMemberToMemberStatus(extraGroup.Status.TiDB.Members)
	}

	for name, status := range tc.Status.TiDB.Members {
		pod, podOk := nameToPod[name]
		if !podOk {
			continue
		}

		var healthAfterPodCreated bool
		memberStatus := cloudv1.MemberStatusUnavailable
		switch {
		case pod.DeletionTimestamp != nil:
			memberStatus = cloudv1.MemberStatusTerminating
		case !k8s.IsPodReady2(pod):
			memberStatus = cloudv1.MemberStatusCreating
		case !status.Health && previousMemberStatusMap[name] == cloudv1.MemberStatusCreating:
			memberStatus = cloudv1.MemberStatusCreating
		case status.Health && k8s.IsPodReady2(pod):
			memberStatus = cloudv1.MemberStatusNormal

			healthyCount++
			healthAfterPodCreated = status.LastTransitionTime.After(pod.CreationTimestamp.Time) ||
				time.Now().After(pod.CreationTimestamp.Add(podCreationTimeSkew))
		}

		var zone string
		var instanceType = previousInstanceTypeMap[name]
		for _, node := range nodesOfPod {
			if node.Name == status.NodeName {
				zone = namer.GetNodeZoneFromLabels(node.Labels)
				if instanceType == "" || healthAfterPodCreated {
					instanceType = namer.GetNodeInstanceTypeFromLabels(node.Labels)
				}
				break
			}
		}

		if instanceType != "" {
			instanceSpec, err = t.Provider.FindInstanceType(instanceType)
			if err != nil {
				return nil, err
			}
		}

		members = append(members, cloudv1.TiDBMember{
			Name: name,
			Labels: map[string]string{
				cloudv1.MemberLabelAvailableZone: zone,
			},
			Resources: cloudv1.Resources{
				CPU:    instanceSpec.CPU,
				Memory: instanceSpec.Memory,
			},
			Status:       memberStatus,
			InstanceType: instanceType,
		})
	}

	var instanceTypeScalingPhase cloudv1.InstanceTypeScalePhase
	if extraGroup.Status.TiDB != nil {
		instanceTypeScalingPhase = extraGroup.Status.TiDB.InstanceTypeScalePhase
	}

	itScaling := tidbMembersHaveDisallowedInstanceType(members, func(instanceType string) bool {
		return t.Cluster.TiDBMemberInstanceTypeAllowedForGroup(instanceType, extraGroup.Name, extraGroup.Spec.TiDB.GetInstanceType())
	}) || instanceTypeScalingPhase != cloudv1.InstanceTypeScalePhaseNone

	var status cloudv1.ComponentStatus
	switch {
	case healthyCount == 0:
		status = cloudv1.ComponentUnavailable
	case tc.TiDBUpgrading() && !itScaling:
		status = cloudv1.ComponentUpgrading
	case tc.TiDBScaling() || tidbNodeScaling || tidbNodeCountScaling || itScaling:
		status = cloudv1.ComponentScaling
	case !tc.TiDBAllMembersReady() || len(tc.Status.TiDB.Members) != healthyCount:
		status = cloudv1.ComponentUnhealthy
	default:
		status = cloudv1.ComponentNormal
	}

	var updatedReplicas int32
	if tc.Status.TiDB.StatefulSet != nil {
		updatedReplicas = tc.Status.TiDB.StatefulSet.UpdatedReplicas
	}

	sort.Slice(members, func(i, j int) bool {
		return members[i].Name < members[j].Name
	})

	return &cloudv1.TiDBStatus{
		Members:                members,
		Status:                 status,
		UpdatedReplicas:        updatedReplicas,
		InstanceTypeScalePhase: instanceTypeScalingPhase,
		Conditions:             conditions,
		TokenBasedAuthEnabled:  pointer.BoolDeref(tc.Spec.TiDB.TokenBasedAuthEnabled, false),
	}, nil
}

func (t *TiDBGroup) computeTiProxyStatus(ctx context.Context, tc *v1alpha1.TidbCluster, extraGroup *cloudv1.ExtraGroupCluster) (*cloudv1.TiProxyStatus, error) {
	logger := log.FromContext(ctx).WithValues("task", "TiDBGroup")

	// compute members and conditions
	var members []cloudv1.TiProxyMember
	var nodesOfPod []corev1.Node
	if tc.TiProxyStsActualReplicas() != 0 {
		// get pods
		pods, err := k8s.SelectPods(ctx, t.EnvProvider.GetClient(), tc.Namespace, selectorOfTiDBClusterPods(tc, string(cloudv1.ComponentTiProxy)))
		if err != nil {
			return nil, err
		}
		var podMap = make(map[string]corev1.Pod, len(pods))
		for _, pod := range pods {
			podMap[pod.Name] = pod
		}

		// get nodesOfPod
		for _, p := range pods {
			if len(p.Spec.NodeName) < 1 {
				continue
			}
			node := corev1.Node{}
			err := t.EnvProvider.GetClient().Get(ctx, client.ObjectKey{
				Name: p.Spec.NodeName,
			}, &node)
			if err != nil {
				return nil, err
			}
			nodesOfPod = append(nodesOfPod, node)
		}

		// previous status
		var previousInstanceTypeMap map[string]string
		if extraGroup.Status.TiProxy != nil {
			previousInstanceTypeMap = proxyMemberToInstanceType(extraGroup.Status.TiProxy.Members)
		}

		// compute member status
		for name, status := range tc.Status.TiProxy.Members {
			member := cloudv1.TiProxyMember{
				Name:   name,
				Status: cloudv1.MemberStatusUnavailable,
				Labels: map[string]string{},
			}

			pod, podOk := podMap[name]
			if podOk {
				switch {
				case pod.DeletionTimestamp != nil:
					member.Status = cloudv1.MemberStatusTerminating
				case pod.Status.Phase == corev1.PodPending:
					member.Status = cloudv1.MemberStatusCreating
				case status.Health:
					member.Status = cloudv1.MemberStatusNormal
				}

				var podInstanceType string
				if pod.Spec.NodeName != "" {
					for _, node := range nodesOfPod {
						if node.Name == pod.Spec.NodeName {
							podInstanceType = namer.GetNodeInstanceTypeFromLabels(node.Labels)
							break
						}
					}
				}

				var instanceType string
				if instanceType == "" {
					// if there is no instanceType, try use the one from pod
					instanceType = podInstanceType
				} else if status.LastTransitionTime.After(pod.CreationTimestamp.Time) ||
					time.Now().After(pod.CreationTimestamp.Add(podCreationTimeSkew)) {
					// if it is up after the current pod, use the pod type
					instanceType = podInstanceType
				} else if instanceType != podInstanceType {
					// use previous instance-type's CPU/Memory when scaling up to a new instanc-type before ready
					instanceType = previousInstanceTypeMap[name]
				}

				if instanceType != "" {
					instanceSpec, err := t.Provider.FindInstanceType(instanceType)
					if err != nil {
						return nil, err
					}
					member.Resources.CPU = instanceSpec.CPU
					member.Resources.Memory = instanceSpec.Memory
					member.InstanceType = instanceType
				}

				// find zone info
				var zone string
				for _, node := range nodesOfPod {
					if node.Name == pod.Spec.NodeName {
						zone = namer.GetNodeZoneFromLabels(node.Labels)
						break
					}
				}
				member.Labels[cloudv1.MemberLabelAvailableZone] = zone
			}

			members = append(members, member)
		}
	}

	scaling, err := t.componentNodeScaling(ctx, cloudv1.ComponentTiProxy, tc)
	if err != nil {
		return nil, err
	}
	// compute scaling
	var instanceTypeScalingPhase cloudv1.InstanceTypeScalePhase
	if extraGroup.Status.TiProxy != nil {
		instanceTypeScalingPhase = extraGroup.Status.TiProxy.InstanceTypeScalePhase
	}
	itScaling := instanceTypeScalingFromTiProxyMembers(members, extraGroup.Spec.TiProxy.GetInstanceType()) || instanceTypeScalingPhase != cloudv1.InstanceTypeScalePhaseNone
	nodes, err := t.selectComponentNodes(ctx, cloudv1.ComponentTiProxy, nodesOfPod, tc)
	if err != nil {
		return nil, err
	}
	var status cloudv1.ComponentStatus

	count := 0
	for _, member := range members {
		if member.Status == cloudv1.MemberStatusNormal {
			count++
		}
	}
	switch {
	case extraGroup.Spec.TiProxy.Replicas != 0 && count == 0:
		status = cloudv1.ComponentUnavailable
	case tc.Status.TiProxy.Phase == v1alpha1.UpgradePhase && !itScaling:
		status = cloudv1.ComponentUpgrading
	case tc.Status.TiProxy.Phase == v1alpha1.ScalePhase || scaling || itScaling:
		status = cloudv1.ComponentScaling
	case extraGroup.Spec.TiProxy.Replicas != len(nodes):
		// mark as scaling when the number of nodes is not equal to the number of replicas.
		// For example, scale in/out/up/down TiProxy.
		logger.Info("TiProxy node scaling", "expected", extraGroup.Spec.TiProxy.Replicas, "current", len(nodes))
		status = cloudv1.ComponentScaling

	case count != len(members):
		status = cloudv1.ComponentUnhealthy
	default:
		status = cloudv1.ComponentNormal
	}

	var updatedReplicas int32
	if tc.Status.TiProxy.StatefulSet != nil {
		updatedReplicas = tc.Status.TiProxy.StatefulSet.UpdatedReplicas
	}

	sort.Slice(members, func(i, j int) bool {
		return members[i].Name < members[j].Name
	})

	conditions := t.computeComponentConditions(ctx, cloudv1.ComponentTiProxy, nodes, tc, extraGroup)

	// compute conditions
	return &cloudv1.TiProxyStatus{
		Members:                members,
		Status:                 status,
		UpdatedReplicas:        updatedReplicas,
		InstanceTypeScalePhase: instanceTypeScalingPhase,
		Conditions:             conditions,
	}, nil
}

func (t *TiDBGroup) componentNodeScaling(ctx context.Context, compType cloudv1.ComponentType, tc *v1alpha1.TidbCluster) (bool, error) {
	pods, err := k8s.SelectPods(ctx, t.EnvProvider.GetClient(), tc.Namespace, selectorOfTiDBClusterPods(tc, string(compType)))
	if err != nil {
		return false, err
	}
	scheduledPodCount := 0
	for _, pod := range pods {
		if len(pod.Spec.NodeName) > 0 {
			scheduledPodCount += 1
		}
	}

	podReplicas := 0
	switch compType {
	case cloudv1.ComponentTiDB:
		podReplicas = int(tc.TiDBStsActualReplicas())
	case cloudv1.ComponentTiProxy:
		podReplicas = int(tc.TiProxyStsActualReplicas())
	default:
		// other components does not consider node scaling
		return false, nil
	}

	if scheduledPodCount < podReplicas {
		return true, nil
	}
	return false, nil
}

func (t *TiDBGroup) isNodeScalingCheckEnabled() bool {
	// Node scaling check (replicas != len(nodes)) is only for AWS self-managed nodegroup (e.g. Terminating:Wait).
	return t.Cluster.FeatureEnabled(consts.ClusterFeatureNodeScalingCheck) &&
		apiutil.GetAnnotation(t.Cluster, consts.UseSelfManagedNodegroupAnnotation) == consts.AnnoValTrue
}

func selectorOfTiDBClusterPods(tc *v1alpha1.TidbCluster, componentName string) map[string]string {
	return map[string]string{
		consts.ComponentLabelKey: componentName,
		consts.InstanceLabelKey:  tc.GetName(),
	}
}

func resetExtraGroupStatusForTiDBCluster(c *cloudv1.Cluster, extraGroup *cloudv1.ExtraGroupCluster) {
	extraGroup.Status.TiDB = nil
	extraGroup.Status.TiProxy = nil
	extraGroup.Status.Phase = ""
	extraGroup.Status.Reason = ""
	extraGroup.Status.SecondaryStatus = ""
	extraGroup.Status.Version = ""
	extraGroup.Status.Conditions = nil
}

func isExtraGroupStatusExsisted(extraGroup *cloudv1.ExtraGroupCluster) bool {
	return extraGroup.Status.TiDB != nil
}