package cluster

import (
	"context"
	"fmt"

	"github.com/mitchellh/mapstructure"
	"github.com/samber/lo"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	clusteragent "github.com/tidbcloud/eks-provider/controllers/utils/cluster_agent"
	"github.com/tidbcloud/eks-provider/controllers/workflow/action"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-agent/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
)

// setStandbyNode is used to manage standby node of cluster component without changing the component member's count.
// In other words, tc.spec.COMPONENT.replicas remains unchanged.
type setStandbyNode struct {
	cli client.Client
}

type setStandbyNodeInputs struct {
	Namespace     string                `mapstructure:"namespace" json:"namespace" validate:"required"`
	ClusterName   string                `mapstructure:"clusterName" json:"clusterName" validate:"required"`
	ComponentType cloudv1.ComponentType `mapstructure:"componentType" json:"componentType" validate:"required"` // for example: "pd", "tidb", etc.
	Zone          string                `mapstructure:"zone" json:"zone" validate:"required"`
	Count         uint                  `mapstructure:"count" json:"count"` // the desired node count, for example: 1

	// WaitNodeReady indicates whether to wait for nodes to be ready
	WaitNodeReady bool `mapstructure:"waitNodeReady" validate:"omitempty"`
}

var (
	_ action.Action = &setStandbyNode{}
)

func NewSetStandbyNode(cli client.Client) action.Action {
	return &setStandbyNode{cli: cli}
}

func (p *setStandbyNode) Name() cloudv1.WorkflowActionName {
	return consts.WfActionSetStandbyNodeAction
}

func (p *setStandbyNode) PreCheck(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep) error {
	inputs := &setStandbyNodeInputs{}
	if err := mapstructure.Decode(step.Inputs.Object, inputs); err != nil {
		return fmt.Errorf("decode inputs: %w", err)
	}
	if err := validate.Struct(inputs); err != nil {
		return fmt.Errorf("validate inputs: %w", err)
	}

	return nil
}

func (p *setStandbyNode) Handle(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep) error {
	inputs := &setStandbyNodeInputs{}
	_ = mapstructure.Decode(step.Inputs.Object, inputs)

	cluster, err := getCluster(ctx, p.cli, client.ObjectKey{Namespace: inputs.Namespace, Name: inputs.ClusterName})
	if err != nil {
		return err
	}

	err = cluster.SetStandbyNodeCnt(inputs.ComponentType, inputs.Zone, inputs.Count)
	if err != nil {
		return fmt.Errorf("set standby node count: %w", err)
	}
	if err := p.cli.Update(ctx, cluster); err != nil {
		return fmt.Errorf("update cluster: %w", err)
	}

	if inputs.WaitNodeReady {
		if err := p.waitNodesReady(ctx, cluster, inputs.ComponentType, inputs.Zone); err != nil {
			return fmt.Errorf("wait nodes ready: %w", err)
		}
	}

	return nil
}

func (p *setStandbyNode) waitNodesReady(ctx context.Context, cluster *cloudv1.Cluster, compType cloudv1.ComponentType, zone string) error {
	logger := log.FromContext(ctx)
	cli, err := clusteragent.NewClient(cluster)
	if err != nil {
		return fmt.Errorf("create cluster-agent client: %w", err)
	}

	resp, err := cli.ListNodes(ctx, cluster.Name)
	if err != nil {
		return fmt.Errorf("list nodes: %w", err)
	}
	component := cluster.GetComponent(compType)
	expectedCnt := component.GetReplicas() + int(cluster.GetStandbyNodeCnt(compType, zone))

	// hack: Since we create standby node for every tidb group rather than the target tidb group only,
	// and we have no easy way to know which tidb group node belongs to. In here, we just wait for all standby nodes
	// to be ready.
	// todo(chenfei): After we can create standby node per tidb group, change this logic.
	var nonEmptyTidbGroupCnt, replicasInTidbGroups int
	if compType == cloudv1.ComponentTiDB {
		for _, each := range cluster.GetExtraGroups() {
			replicasInTidbGroups += each.Spec.TiDB.Replicas
			if each.Spec.TiDB.Replicas != 0 {
				nonEmptyTidbGroupCnt++
			}
		}
		expectedCnt += nonEmptyTidbGroupCnt
		expectedCnt += replicasInTidbGroups
	}
	logger.Info("non-empty tidb group count",
		"nonEmptyTidbGroupCnt", nonEmptyTidbGroupCnt,
		"replicasInTidbGroup", replicasInTidbGroups,
		"expectedCnt", expectedCnt)

	compReadyNodes := lo.Filter(resp.Nodes, func(item *apiv1.Node, index int) bool {
		// todo(chenfei): filter out unschedulable nodes?
		return item.ComponentType == clusteragent.ConvertComponentTypeToAPI(compType) &&
			item.Ready
	})
	if expectedCnt != len(compReadyNodes) {
		return fmt.Errorf("expected cnt: %d, ready cnt: %d", expectedCnt, len(compReadyNodes))
	}
	return nil
}
