package zonalshift

import (
	"context"
	"fmt"
	"strings"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/arczonalshift"
	arcztypes "github.com/aws/aws-sdk-go-v2/service/arczonalshift/types"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	elbv2 "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
	elbv2types "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
	"github.com/pingcap/errors"

	cloudaws "github.com/tidbcloud/infra-provider/pkg/cloud/aws"
)

const (
	attrNLBZonalShiftEnabled       = "zonal_shift.config.enabled"
	attrTGConnectionTermination    = "target_health_state.unhealthy.connection_termination.enabled"
	errNotManagedResourceSubstring = "is not a managed resource"
)

// privateNLB holds a private NLB and its target groups.
type privateNLB struct {
	ARN    string
	TGARNs []string
}

// activeZonalShift is a customer-initiated ARC zonal shift on an NLB.
type activeZonalShift struct {
	ID       string
	AwayFrom string
}

// AWSAPI abstracts AWS calls used by zonal shift actions for testability.
type AWSAPI interface {
	ResolveAZID(ctx context.Context, role *cloudaws.AssumedRole, region, azName string) (string, error)
	// DescribeNLB validates the NLB exists and returns it with target group ARNs.
	DescribeNLB(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) (privateNLB, error)
	EnsureTGConnectionTermination(ctx context.Context, role *cloudaws.AssumedRole, region string, tgARNs []string, enabled bool) error
	EnsureNLBZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string, enabled bool) error
	StartZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN, azID, expiresIn, comment string) (string, error)
	UpdateZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, shiftID, expiresIn string) error
	ListActiveZonalShifts(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) ([]activeZonalShift, error)
	CancelZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, shiftID string) error
}

type sdkAWSAPI struct {
	cli cloudaws.Client
}

func NewSDKAWSAPI(cli cloudaws.Client) AWSAPI {
	return &sdkAWSAPI{cli: cli}
}

func (a *sdkAWSAPI) ResolveAZID(ctx context.Context, role *cloudaws.AssumedRole, region, azName string) (string, error) {
	out, err := a.cli.EC2(role, cloudaws.WithRegion(region)).DescribeAvailabilityZones(ctx, &ec2.DescribeAvailabilityZonesInput{
		ZoneNames: []string{azName},
	})
	if err != nil {
		return "", errors.Wrapf(err, "describe availability zone %s", azName)
	}
	if len(out.AvailabilityZones) == 0 || out.AvailabilityZones[0].ZoneId == nil || *out.AvailabilityZones[0].ZoneId == "" {
		return "", fmt.Errorf("availability zone %s not found", azName)
	}
	return *out.AvailabilityZones[0].ZoneId, nil
}

func (a *sdkAWSAPI) DescribeNLB(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) (privateNLB, error) {
	elb := a.cli.ELBV2(role, cloudaws.WithRegion(region))
	lbOut, err := elb.DescribeLoadBalancers(ctx, &elbv2.DescribeLoadBalancersInput{
		LoadBalancerArns: []string{nlbARN},
	})
	if err != nil {
		return privateNLB{}, errors.Wrapf(err, "describe load balancer %s", nlbARN)
	}
	if len(lbOut.LoadBalancers) == 0 {
		return privateNLB{}, fmt.Errorf("load balancer %s not found", nlbARN)
	}

	tgOut, err := elb.DescribeTargetGroups(ctx, &elbv2.DescribeTargetGroupsInput{
		LoadBalancerArn: aws.String(nlbARN),
	})
	if err != nil {
		return privateNLB{}, errors.Wrapf(err, "describe target groups for NLB %s", nlbARN)
	}
	if len(tgOut.TargetGroups) == 0 {
		return privateNLB{}, fmt.Errorf("no target groups found for NLB %s", nlbARN)
	}
	tgARNs := make([]string, 0, len(tgOut.TargetGroups))
	for _, tg := range tgOut.TargetGroups {
		if tg.TargetGroupArn != nil && *tg.TargetGroupArn != "" {
			tgARNs = append(tgARNs, *tg.TargetGroupArn)
		}
	}
	if len(tgARNs) == 0 {
		return privateNLB{}, fmt.Errorf("no target group ARNs found for NLB %s", nlbARN)
	}
	return privateNLB{ARN: nlbARN, TGARNs: tgARNs}, nil
}

func (a *sdkAWSAPI) EnsureTGConnectionTermination(ctx context.Context, role *cloudaws.AssumedRole, region string, tgARNs []string, enabled bool) error {
	elb := a.cli.ELBV2(role, cloudaws.WithRegion(region))
	want := boolToAttr(enabled)
	for _, tgARN := range tgARNs {
		cur, err := getTargetGroupAttributes(ctx, elb, tgARN)
		if err != nil {
			return errors.Wrapf(err, "describe target group attributes %s", tgARN)
		}
		if cur[attrTGConnectionTermination] == want {
			continue
		}
		_, err = elb.ModifyTargetGroupAttributes(ctx, &elbv2.ModifyTargetGroupAttributesInput{
			TargetGroupArn: aws.String(tgARN),
			Attributes: []elbv2types.TargetGroupAttribute{
				{Key: aws.String(attrTGConnectionTermination), Value: aws.String(want)},
			},
		})
		if err != nil {
			return errors.Wrapf(err, "modify target group attributes %s", tgARN)
		}
	}
	return nil
}

func (a *sdkAWSAPI) EnsureNLBZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string, enabled bool) error {
	elb := a.cli.ELBV2(role, cloudaws.WithRegion(region))
	want := boolToAttr(enabled)
	cur, err := getLoadBalancerAttributes(ctx, elb, nlbARN)
	if err != nil {
		return errors.Wrapf(err, "describe load balancer attributes %s", nlbARN)
	}
	if cur[attrNLBZonalShiftEnabled] == want {
		return nil
	}
	_, err = elb.ModifyLoadBalancerAttributes(ctx, &elbv2.ModifyLoadBalancerAttributesInput{
		LoadBalancerArn: aws.String(nlbARN),
		Attributes: []elbv2types.LoadBalancerAttribute{
			{Key: aws.String(attrNLBZonalShiftEnabled), Value: aws.String(want)},
		},
	})
	if err != nil {
		return errors.Wrapf(err, "modify load balancer attributes %s", nlbARN)
	}
	if !enabled {
		verify, err := getLoadBalancerAttributes(ctx, elb, nlbARN)
		if err != nil {
			return errors.Wrapf(err, "verify load balancer attributes %s", nlbARN)
		}
		if verify[attrNLBZonalShiftEnabled] != "false" {
			return fmt.Errorf("verification failed: %s is %q, expected false", attrNLBZonalShiftEnabled, verify[attrNLBZonalShiftEnabled])
		}
	}
	return nil
}

func (a *sdkAWSAPI) StartZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN, azID, expiresIn, comment string) (string, error) {
	out, err := a.cli.ARCZonalShift(role, cloudaws.WithRegion(region)).StartZonalShift(ctx, &arczonalshift.StartZonalShiftInput{
		AwayFrom:           aws.String(azID),
		ExpiresIn:          aws.String(expiresIn),
		ResourceIdentifier: aws.String(nlbARN),
		Comment:            aws.String(comment),
	})
	if err != nil {
		return "", err
	}
	if out.ZonalShiftId == nil || *out.ZonalShiftId == "" {
		return "", fmt.Errorf("StartZonalShift returned empty zonal shift id")
	}
	return *out.ZonalShiftId, nil
}

func (a *sdkAWSAPI) UpdateZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, shiftID, expiresIn string) error {
	_, err := a.cli.ARCZonalShift(role, cloudaws.WithRegion(region)).UpdateZonalShift(ctx, &arczonalshift.UpdateZonalShiftInput{
		ZonalShiftId: aws.String(shiftID),
		ExpiresIn:    aws.String(expiresIn),
	})
	return errors.Wrapf(err, "update zonal shift %s", shiftID)
}

func (a *sdkAWSAPI) ListActiveZonalShifts(ctx context.Context, role *cloudaws.AssumedRole, region, nlbARN string) ([]activeZonalShift, error) {
	arc := a.cli.ARCZonalShift(role, cloudaws.WithRegion(region))
	var (
		shifts    []activeZonalShift
		nextToken *string
	)
	for {
		out, err := arc.ListZonalShifts(ctx, &arczonalshift.ListZonalShiftsInput{
			Status:    arcztypes.ZonalShiftStatusActive,
			NextToken: nextToken,
		})
		if err != nil {
			return nil, errors.Wrap(err, "list active zonal shifts")
		}
		for _, item := range out.Items {
			if item.ShiftType != arcztypes.ShiftTypeZonalShift {
				continue
			}
			if item.ResourceIdentifier == nil || *item.ResourceIdentifier != nlbARN || item.ZonalShiftId == nil {
				continue
			}
			awayFrom := ""
			if item.AwayFrom != nil {
				awayFrom = *item.AwayFrom
			}
			shifts = append(shifts, activeZonalShift{
				ID:       *item.ZonalShiftId,
				AwayFrom: awayFrom,
			})
		}
		if out.NextToken == nil || *out.NextToken == "" {
			break
		}
		nextToken = out.NextToken
	}
	return shifts, nil
}

func (a *sdkAWSAPI) CancelZonalShift(ctx context.Context, role *cloudaws.AssumedRole, region, shiftID string) error {
	_, err := a.cli.ARCZonalShift(role, cloudaws.WithRegion(region)).CancelZonalShift(ctx, &arczonalshift.CancelZonalShiftInput{
		ZonalShiftId: aws.String(shiftID),
	})
	return errors.Wrapf(err, "cancel zonal shift %s", shiftID)
}

func boolToAttr(v bool) string {
	if v {
		return "true"
	}
	return "false"
}

func getTargetGroupAttributes(ctx context.Context, elb *elbv2.Client, tgARN string) (map[string]string, error) {
	out, err := elb.DescribeTargetGroupAttributes(ctx, &elbv2.DescribeTargetGroupAttributesInput{
		TargetGroupArn: aws.String(tgARN),
	})
	if err != nil {
		return nil, err
	}
	attrs := map[string]string{}
	for _, a := range out.Attributes {
		if a.Key != nil && a.Value != nil {
			attrs[*a.Key] = *a.Value
		}
	}
	return attrs, nil
}

func getLoadBalancerAttributes(ctx context.Context, elb *elbv2.Client, nlbARN string) (map[string]string, error) {
	out, err := elb.DescribeLoadBalancerAttributes(ctx, &elbv2.DescribeLoadBalancerAttributesInput{
		LoadBalancerArn: aws.String(nlbARN),
	})
	if err != nil {
		return nil, err
	}
	attrs := map[string]string{}
	for _, a := range out.Attributes {
		if a.Key != nil && a.Value != nil {
			attrs[*a.Key] = *a.Value
		}
	}
	return attrs, nil
}

func isNotManagedResourceError(err error) bool {
	if err == nil {
		return false
	}
	return strings.Contains(err.Error(), errNotManagedResourceSubstring)
}