package zonalshift

import (
	"context"
	"fmt"
	"slices"
	"strings"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"sigs.k8s.io/controller-runtime/pkg/client"

	eksv1alpha1 "github.com/tidbcloud/eks-provider/api/v1alpha1"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	cloudaws "github.com/tidbcloud/infra-provider/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/pkg/dataplane"
	"github.com/tidbcloud/infra-provider/pkg/util"
)

const zonalShiftIDsOutputKey = "zonalShiftIds"

type roleConfig struct {
	assumeRoleName string
	externalID     string
}

func (c roleConfig) assumedRole(accountID string) *cloudaws.AssumedRole {
	ar := &cloudaws.AssumedRole{
		RoleARN: util.GetRoleARN(accountID, c.assumeRoleName),
	}
	if c.externalID != "" {
		ar.ExternalID = &c.externalID
	}
	return ar
}

func resolveAssumedRole(ctx context.Context, cli client.Client, cluster *cloudv1.Cluster, envProvider dataplane.EnvProvider, cfg roleConfig) (*cloudaws.AssumedRole, error) {
	// Legacy clusters (with NetworkName) store assume role in Network CR.
	// Migrated legacy clusters may no longer have Network; fall through to envProvider.
	if cluster.Spec.NetworkName != "" {
		var network cloudv1.Network
		if err := cli.Get(ctx, cluster.GetNetworkRef(), &network); err == nil {
			if spec, err := eksv1alpha1.EKSNetworkSpecForNetwork(&network); err == nil && spec != nil && spec.AssumeRoleArn != "" {
				return cloudaws.GetRoleFromEKSNetworkSpec(spec), nil
			}
		}
	}
	accountID := envProvider.GetAccount()
	if accountID == "" {
		return nil, fmt.Errorf("empty AWS account id from env provider")
	}
	if cfg.assumeRoleName == "" {
		return nil, fmt.Errorf("assume role name is not configured")
	}
	return cfg.assumedRole(accountID), nil
}

func getCluster(ctx context.Context, cli client.Client, namespace, name string) (*cloudv1.Cluster, error) {
	cluster := &cloudv1.Cluster{}
	if err := cli.Get(ctx, client.ObjectKey{Namespace: namespace, Name: name}, cluster); err != nil {
		if apierrors.IsNotFound(err) {
			return nil, cloudv1.NewWFNonRetriableError(fmt.Errorf("cluster %s/%s not found", namespace, name))
		}
		return nil, err
	}
	return cluster, nil
}

func ensureAWSCluster(cluster *cloudv1.Cluster) error {
	switch cluster.GetProviderType() {
	case cloudv1.ProviderTypeEKS, cloudv1.ProviderTypeAWS:
		return nil
	default:
		return cloudv1.NewWFNonRetriableError(fmt.Errorf("zonal shift is only supported on AWS clusters, got provider %s", cluster.GetProviderType()))
	}
}

// privateNLBARNsForGroups returns private NLB ARNs for the requested groups from Cluster status.
// Empty groupIDs means default group + all ExtraGroups. Order: default first, then ExtraGroups.
func privateNLBARNsForGroups(cluster *cloudv1.Cluster, groupIDs []string) ([]string, error) {
	includeAll := len(groupIDs) == 0
	includeDefault := includeAll || slices.Contains(groupIDs, cluster.GetDefaultGroupID())

	seen := map[string]struct{}{}
	arns := make([]string, 0)
	add := func(arn string) {
		if arn == "" {
			return
		}
		if _, ok := seen[arn]; ok {
			return
		}
		seen[arn] = struct{}{}
		arns = append(arns, arn)
	}

	if includeDefault {
		if ep := cluster.Status.SqlEndpoints.PrivateEndpoint; ep != nil {
			add(ep.ResourceID)
		}
	}
	for _, eg := range cluster.GetExtraGroups() {
		if eg.Deleted() {
			continue
		}
		if !includeAll && !slices.Contains(groupIDs, eg.Name) {
			continue
		}
		if eg.Status == nil {
			continue
		}
		if ep := eg.Status.SqlEndpointsStatus.PrivateEndpoint; ep != nil {
			add(ep.ResourceID)
		}
	}
	if len(arns) == 0 {
		return nil, cloudv1.NewWFNonRetriableError(fmt.Errorf("no private NLB ARNs found in cluster status for requested groups"))
	}
	return arns, nil
}

// discoverTargetNLBs selects private NLBs from Cluster status, then loads TG info via AWS.
func discoverTargetNLBs(ctx context.Context, api AWSAPI, role *cloudaws.AssumedRole, region string, cluster *cloudv1.Cluster, groupIDs []string) ([]privateNLB, error) {
	arns, err := privateNLBARNsForGroups(cluster, groupIDs)
	if err != nil {
		return nil, err
	}
	out := make([]privateNLB, 0, len(arns))
	for _, arn := range arns {
		nlb, err := api.DescribeNLB(ctx, role, region, arn)
		if err != nil {
			return nil, cloudv1.NewWFNonRetriableError(err)
		}
		out = append(out, nlb)
	}
	return out, nil
}

func findMatchingZonalShift(shifts []activeZonalShift, azID string) (string, bool) {
	for _, shift := range shifts {
		if shift.AwayFrom == azID {
			return shift.ID, true
		}
	}
	return "", false
}

// ensureZonalShift returns an active zonal shift ID for the NLB, starting one if needed.
// Idempotent: reuses an existing ACTIVE customer shift with the same awayFrom AZ.
func ensureZonalShift(
	ctx context.Context,
	api AWSAPI,
	role *cloudaws.AssumedRole,
	region, nlbARN, azID, expiresIn, comment string,
) (string, error) {
	shifts, err := api.ListActiveZonalShifts(ctx, role, region, nlbARN)
	if err != nil {
		return "", err
	}
	if shiftID, ok := findMatchingZonalShift(shifts, azID); ok {
		return shiftID, nil
	}
	if len(shifts) > 0 {
		return "", cloudv1.NewWFNonRetriableError(fmt.Errorf(
			"NLB %s already has active zonal shift %s away from %s, requested %s",
			nlbARN, shifts[0].ID, shifts[0].AwayFrom, azID))
	}

	shiftID, err := api.StartZonalShift(ctx, role, region, nlbARN, azID, expiresIn, comment)
	if err == nil {
		return shiftID, nil
	}
	if isNotManagedResourceError(err) {
		return "", fmt.Errorf("ARC has not managed the NLB yet, will retry: %w", err)
	}

	// Start may race or conflict if a shift became active between list and start; re-check.
	shifts, listErr := api.ListActiveZonalShifts(ctx, role, region, nlbARN)
	if listErr == nil {
		if shiftID, ok := findMatchingZonalShift(shifts, azID); ok {
			return shiftID, nil
		}
		if len(shifts) > 0 {
			return "", cloudv1.NewWFNonRetriableError(fmt.Errorf(
				"NLB %s already has active zonal shift %s away from %s, requested %s",
				nlbARN, shifts[0].ID, shifts[0].AwayFrom, azID))
		}
	}
	return "", cloudv1.NewWFNonRetriableError(err)
}

func persistZonalShiftIDs(ctx context.Context, wfCli client.Client, wf *cloudv1.Workflow, step cloudv1.WorkflowStep, shiftIDs []string) error {
	stepStatus, ok := wf.StepStatus(ctx, step.Name)
	if !ok {
		return fmt.Errorf("step status %s not found", step.Name)
	}
	if err := stepStatus.UpdateOutput(ctx, map[string]string{
		zonalShiftIDsOutputKey: strings.Join(shiftIDs, ","),
	}); err != nil {
		return err
	}
	return wfCli.Status().Update(ctx, wf)
}