package az

import (
	"context"
	"fmt"

	"github.com/mitchellh/mapstructure"
	"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"
	cloudaws "github.com/tidbcloud/infra-provider/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/dataplane"
)

const (
	recoverSubStepDiscover        = "discover"
	recoverSubStepUncordonNodes   = "uncordon-nodes"
	recoverSubStepWaitPodsRunning = "wait-pods-running"
	recoverSubStepVerify          = "verify"
)

type recoverAZ struct {
	cli            client.Client
	cachedProvider dataplane.EnvProviderManager
	awsCli         cloudaws.Client
	roleCfg        asgRoleConfig
}

type recoverAZInputs struct {
	Namespace        string   `mapstructure:"namespace" json:"namespace" validate:"required"`
	ClusterName      string   `mapstructure:"clusterName" json:"clusterName" validate:"required"`
	AvailabilityZone string   `mapstructure:"availabilityZone" json:"availabilityZone" validate:"required"`
	GroupIDs         []string `mapstructure:"groupIds" json:"groupIds"`
}

var _ action.Action = &recoverAZ{}

func NewRecoverAZ(cli client.Client, cachedProvider dataplane.EnvProviderManager, awsCli cloudaws.Client, assumeRoleName, externalID string) action.Action {
	return &recoverAZ{
		cli:            cli,
		cachedProvider: cachedProvider,
		awsCli:         awsCli,
		roleCfg:        asgRoleConfig{assumeRoleName: assumeRoleName, externalID: externalID},
	}
}

func (a *recoverAZ) Name() cloudv1.WorkflowActionName {
	return consts.WfActionRecoverAZAction
}

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

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

	logger := log.FromContext(ctx).WithValues(
		"clusterName", inputs.ClusterName,
		"availabilityZone", inputs.AvailabilityZone,
		"stepName", step.Name,
		"groupIds", inputs.GroupIDs,
	)
	ctx = log.IntoContext(ctx, logger)

	cluster, err := getCluster(ctx, a.cli, inputs.Namespace, inputs.ClusterName)
	if err != nil {
		return err
	}
	groups, err := resolveGroups(cluster, inputs.GroupIDs)
	if err != nil {
		return err
	}
	logger.Info("recover AZ groups resolved", "groups", groupIDsOf(groups))

	caClient, err := clusteragent.NewClient(cluster)
	if err != nil {
		return fmt.Errorf("create cluster-agent client: %w", err)
	}
	remoteCli, err := a.cachedProvider.GetClient(ctx, cluster)
	if err != nil {
		return fmt.Errorf("get dataplane cli: %w", err)
	}
	envProvider, err := a.cachedProvider.GetEnvProvider(ctx, cluster)
	if err != nil {
		return fmt.Errorf("get env provider: %w", err)
	}

	stepStatus, _ := wf.StepStatus(ctx, step.Name)
	current := stepStatus.Outputs[lastSubStepKey]
	if current == "" {
		current = recoverSubStepDiscover
	}

	return runSubSteps(ctx, a.cli, wf, step, current, func(ctx context.Context, subStep string) (subStepResult, error) {
		switch subStep {
		case recoverSubStepDiscover:
			return a.discover(ctx, caClient, cluster, groups, inputs)
		case recoverSubStepUnpauseASGs:
			return a.unpause(ctx, cluster, envProvider, groups, inputs)
		case recoverSubStepUncordonNodes:
			return a.uncordon(ctx, caClient, remoteCli, inputs, groups)
		case recoverSubStepWaitPodsRunning:
			return a.waitPods(ctx, caClient, inputs, groups)
		case recoverSubStepVerify:
			return a.verify(ctx, remoteCli, cluster, stepStatus, groups)
		default:
			return subStepResult{}, cloudv1.NewWFNonRetriableError(fmt.Errorf("unknown sub-step %s", subStep))
		}
	})
}

func (a *recoverAZ) discover(ctx context.Context, caClient clusteragent.Client, cluster *cloudv1.Cluster, groups []Group, inputs *recoverAZInputs) (subStepResult, error) {
	logger := log.FromContext(ctx)
	logger.Info("discovering AZ resources for recovery")
	totalNodes, totalPods, totalASGs := 0, 0, 0
	for _, g := range groups {
		resources, err := discoverAZResources(ctx, caClient, inputs.ClusterName, inputs.AvailabilityZone, g)
		if err != nil {
			return subStepResult{}, err
		}
		asgs, err := listGroupASGs(ctx, a.cli, g, inputs.AvailabilityZone)
		if err != nil {
			return subStepResult{}, err
		}
		logger.Info("discovered AZ resources for group",
			"groupId", g.ID(),
			"nodeCount", len(resources.NodeNames),
			"podCount", len(resources.Pods),
			"asgCount", len(asgs),
		)
		totalNodes += len(resources.NodeNames)
		totalPods += len(resources.Pods)
		totalASGs += len(asgs)
	}
	if isAWSProvider(cluster) {
		// Nodes may be NotReady or already gone after AZ fault; recovery still needs
		// to unpause ASGs so replacements can be launched. Require ASGs on AWS.
		if totalASGs == 0 {
			return subStepResult{}, cloudv1.NewWFNonRetriableError(fmt.Errorf(
				"no ASGs found in AZ %s for requested groups", inputs.AvailabilityZone))
		}
		if totalNodes == 0 {
			logger.Info("no nodes found in AZ; continuing to unpause ASGs for instance replacement",
				"availabilityZone", inputs.AvailabilityZone, "asgCount", totalASGs, "podCount", totalPods)
		}
	} else if totalNodes == 0 && totalPods == 0 {
		return subStepResult{}, cloudv1.NewWFNonRetriableError(fmt.Errorf(
			"no nodes or pods found in AZ %s for requested groups", inputs.AvailabilityZone))
	}
	logger.Info("discovered AZ resources for recovery results",
		"nodeCount", totalNodes, "podCount", totalPods, "asgCount", totalASGs)
	return advanceTo(recoverSubStepUnpauseASGs, nil)
}

func (a *recoverAZ) unpause(ctx context.Context, cluster *cloudv1.Cluster, envProvider dataplane.EnvProvider, groups []Group, inputs *recoverAZInputs) (subStepResult, error) {
	logger := log.FromContext(ctx)
	logger.Info("unpausing ASGs in AZ")
	var unpaused []string
	for _, g := range groups {
		names, err := unpauseASGs(ctx, a.cli, a.awsCli, cluster, envProvider, a.roleCfg, g, inputs.AvailabilityZone)
		if err != nil {
			return subStepResult{}, err
		}
		unpaused = append(unpaused, names...)
	}
	logger.Info("unpaused ASGs", "count", len(unpaused), "asgs", unpaused)
	return advanceTo(recoverSubStepUncordonNodes, map[string]string{
		recoverUnpausedASGsKey: joinNames(unpaused),
	})
}

func (a *recoverAZ) uncordon(ctx context.Context, caClient clusteragent.Client, remoteCli client.Client, inputs *recoverAZInputs, groups []Group) (subStepResult, error) {
	logger := log.FromContext(ctx)
	logger.Info("uncordoning AZ nodes")
	var uncordonedNodes []string
	for _, g := range groups {
		resources, err := discoverAZResources(ctx, caClient, inputs.ClusterName, inputs.AvailabilityZone, g)
		if err != nil {
			return subStepResult{}, err
		}
		names, err := uncordonNodes(ctx, remoteCli, resources.NodeNames)
		if err != nil {
			return subStepResult{}, err
		}
		uncordonedNodes = append(uncordonedNodes, names...)
	}
	logger.Info("uncordoned AZ nodes", "nodeCount", len(uncordonedNodes))
	return advanceTo(recoverSubStepWaitPodsRunning, map[string]string{
		recoverUncordonedNodesKey: joinNames(uncordonedNodes),
	})
}

func (a *recoverAZ) waitPods(ctx context.Context, caClient clusteragent.Client, inputs *recoverAZInputs, groups []Group) (subStepResult, error) {
	logger := log.FromContext(ctx)
	logger.Info("waiting for component pods to become running")
	var runningPods []string
	for _, g := range groups {
		names, err := waitPodsRunning(ctx, caClient, inputs.ClusterName, g)
		if err != nil {
			return subStepResult{}, err
		}
		runningPods = append(runningPods, names...)
	}
	logger.Info("wait pods running completed", "podCount", len(runningPods))
	return advanceTo(recoverSubStepVerify, nil)
}

func (a *recoverAZ) verify(ctx context.Context, remoteCli client.Client, cluster *cloudv1.Cluster, stepStatus *cloudv1.WorkflowStepStatus, groups []Group) (subStepResult, error) {
	logger := log.FromContext(ctx)
	namespace := cluster.PrefixedId()
	for _, g := range groups {
		logger.Info("verifying tidbcluster ready", "groupId", g.ID(), "tcName", g.Namer().GetTiDBName())
		if err := verifyGroupTCReady(ctx, remoteCli, namespace, g); err != nil {
			return subStepResult{}, err
		}
	}
	logger.Info("recover AZ completed",
		"uncordonedNodes", stepStatus.Outputs[recoverUncordonedNodesKey],
		"unpausedASGs", stepStatus.Outputs[recoverUnpausedASGsKey],
	)
	return subStepDone()
}