package az

import (
	"context"
	"fmt"
	"time"

	"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 (
	offlineSubStepDiscover         = "discover"
	offlineSubStepTransferPDLeader = "transfer-pd-leader"
	offlineSubStepCordonNodes      = "cordon-nodes"
	offlineSubStepGracefulShutdown = "graceful-shutdown"
	offlineSubStepWaitGraceful     = "wait-graceful"
	offlineSubStepForceShutdown    = "force-shutdown"
	offlineSubStepVerify           = "verify"

	offlineGracefulStartedAtKey = "graceful-started-at"
)

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

type offlineAZInputs 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"`
	GracefulTimeout  string   `mapstructure:"gracefulTimeout" json:"gracefulTimeout"`
	GroupIDs         []string `mapstructure:"groupIds" json:"groupIds"`
}

var _ action.Action = &offlineAZ{}

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

func (a *offlineAZ) Name() cloudv1.WorkflowActionName {
	return consts.WfActionOfflineAZAction
}

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

func (a *offlineAZ) Handle(ctx context.Context, wf *cloudv1.Workflow, step cloudv1.WorkflowStep) error {
	inputs := &offlineAZInputs{}
	_ = 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("offline 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 = offlineSubStepDiscover
	}
	gracefulTimeout := parseGracefulTimeout(inputs.GracefulTimeout)
	namespace := cluster.PrefixedId()

	return runSubSteps(ctx, a.cli, wf, step, current, func(ctx context.Context, subStep string) (subStepResult, error) {
		switch subStep {
		case offlineSubStepDiscover:
			return a.discover(ctx, caClient, inputs, groups)
		case offlineSubStepPauseASGs:
			return a.pauseASGs(ctx, cluster, envProvider, groups, inputs)
		case offlineSubStepTransferPDLeader:
			return a.transferPDLeader(ctx, caClient, inputs, groups)
		case offlineSubStepCordonNodes:
			return a.cordon(ctx, caClient, remoteCli, inputs, groups)
		case offlineSubStepGracefulShutdown:
			return a.gracefulShutdown(ctx, caClient, remoteCli, stepStatus, inputs, namespace, groups)
		case offlineSubStepWaitGraceful:
			return a.waitGraceful(ctx, caClient, remoteCli, stepStatus, inputs, namespace, gracefulTimeout, groups)
		case offlineSubStepForceShutdown:
			return a.forceShutdown(ctx, caClient, remoteCli, stepStatus, inputs, namespace, groups)
		case offlineSubStepVerify:
			return a.verify(ctx, caClient, stepStatus, inputs, groups)
		default:
			return subStepResult{}, cloudv1.NewWFNonRetriableError(fmt.Errorf("unknown sub-step %s", subStep))
		}
	})
}

func (a *offlineAZ) discover(ctx context.Context, caClient clusteragent.Client, inputs *offlineAZInputs, groups []Group) (subStepResult, error) {
	logger := log.FromContext(ctx)
	logger.Info("discovering AZ resources for AZ", "availabilityZone", inputs.AvailabilityZone)
	totalNodes, totalPods := 0, 0
	for _, g := range groups {
		resources, err := discoverAZResources(ctx, caClient, inputs.ClusterName, inputs.AvailabilityZone, g)
		if err != nil {
			return subStepResult{}, err
		}
		logger.Info("discovered AZ resources for group", "groupId", g.ID(), "groupLabel", g.ClusterLabel(),
			"nodeCount", len(resources.NodeNames), "podCount", len(resources.Pods))
		totalNodes += len(resources.NodeNames)
		totalPods += len(resources.Pods)
	}
	if totalNodes == 0 && totalPods == 0 {
		return subStepResult{}, cloudv1.NewWFNonRetriableError(fmt.Errorf("no nodes or pods found in AZ %s", inputs.AvailabilityZone))
	}
	logger.Info("discovered AZ resources", "nodeCount", totalNodes, "podCount", totalPods)
	return advanceTo(offlineSubStepPauseASGs, nil)
}

func (a *offlineAZ) pauseASGs(ctx context.Context, cluster *cloudv1.Cluster, envProvider dataplane.EnvProvider, groups []Group, inputs *offlineAZInputs) (subStepResult, error) {
	logger := log.FromContext(ctx)
	logger.Info("pausing ASGs in fault AZ")
	var paused []string
	for _, g := range groups {
		names, err := pauseASGs(ctx, a.cli, a.awsCli, cluster, envProvider, a.roleCfg, g, inputs.AvailabilityZone)
		if err != nil {
			return subStepResult{}, err
		}
		paused = append(paused, names...)
	}
	logger.Info("paused ASGs", "count", len(paused), "asgs", paused)
	return advanceTo(offlineSubStepTransferPDLeader, map[string]string{
		offlinePausedASGsKey: joinNames(paused),
	})
}

func (a *offlineAZ) transferPDLeader(ctx context.Context, caClient clusteragent.Client, inputs *offlineAZInputs, groups []Group) (subStepResult, error) {
	logger := log.FromContext(ctx)
	dg := findDefaultGroup(groups)
	if dg == nil {
		logger.Info("skip PD leader transfer: default group not in scope")
		return advanceTo(offlineSubStepCordonNodes, nil)
	}
	logger.Info("transferring PD leader if needed", "groupId", dg.ID())
	resources, err := discoverAZResources(ctx, caClient, inputs.ClusterName, "", dg)
	if err != nil {
		return subStepResult{}, err
	}
	pdURL, err := getPDURL(ctx, caClient, inputs.ClusterName)
	if err != nil {
		return subStepResult{}, err
	}
	if err := transferPDLeaderIfNeeded(ctx, caClient, inputs.ClusterName, pdURL, inputs.AvailabilityZone, resources); err != nil {
		return subStepResult{}, err
	}
	return advanceTo(offlineSubStepCordonNodes, nil)
}

func (a *offlineAZ) cordon(ctx context.Context, caClient clusteragent.Client, remoteCli client.Client, inputs *offlineAZInputs, groups []Group) (subStepResult, error) {
	logger := log.FromContext(ctx)
	logger.Info("cordoning AZ nodes")
	var cordonedNodes []string
	for _, g := range groups {
		resources, err := discoverAZResources(ctx, caClient, inputs.ClusterName, inputs.AvailabilityZone, g)
		if err != nil {
			return subStepResult{}, err
		}
		names, err := cordonNodes(ctx, remoteCli, resources.NodeNames)
		if err != nil {
			return subStepResult{}, err
		}
		cordonedNodes = append(cordonedNodes, names...)
	}
	logger.Info("cordoned AZ nodes", "nodeCount", len(cordonedNodes))
	return advanceTo(offlineSubStepGracefulShutdown, map[string]string{
		offlineCordonedNodesKey: joinNames(cordonedNodes),
	})
}

func (a *offlineAZ) gracefulShutdown(
	ctx context.Context,
	caClient clusteragent.Client,
	remoteCli client.Client,
	stepStatus *cloudv1.WorkflowStepStatus,
	inputs *offlineAZInputs,
	namespace string,
	groups []Group,
) (subStepResult, error) {
	logger := log.FromContext(ctx)
	logger.Info("graceful shutting down AZ pods")
	var deletedPods []string
	for _, g := range groups {
		resources, err := discoverAZResources(ctx, caClient, inputs.ClusterName, inputs.AvailabilityZone, g)
		if err != nil {
			return subStepResult{}, err
		}
		names, err := gracefulShutdownPods(ctx, remoteCli, namespace, resources, g.Components())
		if err != nil {
			return subStepResult{}, err
		}
		deletedPods = append(deletedPods, names...)
	}
	logger.Info("graceful shutdown requested for pods", "podCount", len(deletedPods))
	return advanceTo(offlineSubStepWaitGraceful, map[string]string{
		offlineGracefulStartedAtKey: time.Now().UTC().Format(time.RFC3339),
		offlineDeletedPodsKey:       joinNames(mergeNames(stepStatus.Outputs[offlineDeletedPodsKey], deletedPods)),
	})
}

func (a *offlineAZ) waitGraceful(
	ctx context.Context,
	caClient clusteragent.Client,
	remoteCli client.Client,
	stepStatus *cloudv1.WorkflowStepStatus,
	inputs *offlineAZInputs,
	namespace string,
	gracefulTimeout time.Duration,
	groups []Group,
) (subStepResult, error) {
	logger := log.FromContext(ctx)
	logger.Info("waiting for graceful shutdown to complete")
	var remaining []string
	for _, g := range groups {
		resources, err := discoverAZResources(ctx, caClient, inputs.ClusterName, inputs.AvailabilityZone, g)
		if err != nil {
			return subStepResult{}, err
		}
		names, err := remainingShutdownPods(ctx, remoteCli, namespace, resources, g.Components())
		if err != nil {
			return subStepResult{}, err
		}
		remaining = append(remaining, names...)
	}
	if len(remaining) == 0 {
		logger.Info("all pods gracefully shut down")
		return advanceTo(offlineSubStepVerify, nil)
	}
	startedAt := time.Now().UTC()
	if raw := stepStatus.Outputs[offlineGracefulStartedAtKey]; raw != "" {
		if parsed, err := time.Parse(time.RFC3339, raw); err == nil {
			startedAt = parsed
		}
	}
	if time.Since(startedAt) >= gracefulTimeout {
		logger.Info("graceful timeout reached, force shutdown remaining pods", "remainingCount", len(remaining), "timeout", gracefulTimeout)
		return advanceTo(offlineSubStepForceShutdown, nil)
	}
	logger.Info("pods still shutting down", "remaining", remaining, "elapsed", time.Since(startedAt), "timeout", gracefulTimeout)
	return subStepResult{}, fmt.Errorf("waiting for %d pods to terminate in AZ %s", len(remaining), inputs.AvailabilityZone)
}

func (a *offlineAZ) forceShutdown(
	ctx context.Context,
	caClient clusteragent.Client,
	remoteCli client.Client,
	stepStatus *cloudv1.WorkflowStepStatus,
	inputs *offlineAZInputs,
	namespace string,
	groups []Group,
) (subStepResult, error) {
	logger := log.FromContext(ctx)
	logger.Info("force shutting down remaining pods")
	var remaining []string
	for _, g := range groups {
		resources, err := discoverAZResources(ctx, caClient, inputs.ClusterName, inputs.AvailabilityZone, g)
		if err != nil {
			return subStepResult{}, err
		}
		names, err := remainingShutdownPods(ctx, remoteCli, namespace, resources, g.Components())
		if err != nil {
			return subStepResult{}, err
		}
		remaining = append(remaining, names...)
	}
	forceDeletedPods, err := forceShutdownPods(ctx, remoteCli, namespace, remaining)
	if err != nil {
		return subStepResult{}, err
	}
	logger.Info("force shutdown requested for pods", "podCount", len(forceDeletedPods))
	return advanceTo(offlineSubStepVerify, map[string]string{
		offlineDeletedPodsKey: joinNames(mergeNames(stepStatus.Outputs[offlineDeletedPodsKey], forceDeletedPods)),
	})
}

func (a *offlineAZ) verify(
	ctx context.Context,
	caClient clusteragent.Client,
	stepStatus *cloudv1.WorkflowStepStatus,
	inputs *offlineAZInputs,
	groups []Group,
) (subStepResult, error) {
	logger := log.FromContext(ctx)
	dg := findDefaultGroup(groups)
	if dg == nil {
		logger.Info("skip store verify: default group not in scope")
		return subStepDone()
	}
	logger.Info("verifying stores offline in AZ", "groupId", dg.ID())
	pdURL, err := getPDURL(ctx, caClient, inputs.ClusterName)
	if err != nil {
		return subStepResult{}, err
	}
	if err := verifyStoresOffline(ctx, caClient, inputs.ClusterName, pdURL, inputs.AvailabilityZone); err != nil {
		return subStepResult{}, err
	}
	logger.Info("offline AZ completed",
		"cordonedNodes", stepStatus.Outputs[offlineCordonedNodesKey],
		"deletedPods", stepStatus.Outputs[offlineDeletedPodsKey],
		"pausedASGs", stepStatus.Outputs[offlinePausedASGsKey],
	)
	return subStepDone()
}