package eks

import (
	"context"
	"fmt"
	"sort"
	"time"

	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/vpc/task"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/deploy"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/namer"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/util"
	"github.com/upbound/provider-aws/apis/ec2/v1beta1"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
)

type Subnet struct {
	*task.BaseTask
	nodeSubnets         []*v1beta1.Subnet
	publicSubnets       []*v1beta1.Subnet
	privateSubnets      []*v1beta1.Subnet
	controlPlaneSubnets []*v1beta1.Subnet
	endpointSubnets     []*v1beta1.Subnet
}

func NewSubnet(ctx *task.Context) *Subnet {
	return &Subnet{
		BaseTask:            &task.BaseTask{Context: ctx},
		nodeSubnets:         []*v1beta1.Subnet{},
		publicSubnets:       []*v1beta1.Subnet{},
		privateSubnets:      []*v1beta1.Subnet{},
		controlPlaneSubnets: []*v1beta1.Subnet{},
		endpointSubnets:     []*v1beta1.Subnet{},
	}
}

// Name implements Task interface.
// Return the name of the task.
func (s *Subnet) Name() string {
	return "Subnet"
}

func (s *Subnet) getVPCName() string {
	return s.VPC.Status.VPC.Name
}

// Create implements Task interface.
// for each zone create 3 subnets:
//   - nodes-<zone>: used for node group
//   - private-<zone>: used for private LB
//   - public-<zone>: used for public LB and accessing internet(via NAT now).
func (s *Subnet) Create(ctx context.Context) deploy.Result {
	s.BeginAt = util.Pointer(time.Now())
	tags := controllers.ConvertToCloudTags(s.VPC.Labels, cloudv1.ProviderTypeEKS)
	tags = util.MergeMaps(tags, s.Config.AdditionalCloudTags)
	key := types.NamespacedName{
		Name: s.getVPCName(),
	}
	vpcRes := &v1beta1.VPC{}
	err := s.Local.Get(ctx, key, vpcRes)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(time.Second * 30)
	}

	zones := s.GetZones()
	nodeZones := s.GetNodeZones()
	configuredEndpointZones := s.VPC.Spec.PrivateLinkEndpointZones
	useEndpointSubnets := util.ShouldUsePrivateLinkEndpointSubnets(configuredEndpointZones, nodeZones)
	endpointZones := []string(nil)
	if useEndpointSubnets {
		endpointZones = configuredEndpointZones
	}

	if shouldUseOptimizedCIDR(nodeZones, zones) {
		if err := s.createOptimizedSubnets(ctx, vpcRes, tags, nodeZones, zones, endpointZones); err != nil {
			return deploy.Panic(err)
		}

		return deploy.Done()
	}

	if useEndpointSubnets {
		return deploy.Panic(fmt.Errorf(
			"vpc.spec.privateLinkEndpointZones requires optimized CIDR layout, nodeZones=%v, zones=%v",
			nodeZones,
			zones,
		))
	}

	workers, public, private, err := util.GetAwsSubnetCIDRs(s.VPC.Spec.CIDR, len(zones))
	if err != nil {
		return deploy.Panic(err)
	}

	for i := 0; i < len(workers); i++ {
		zone := zones[i]

		// node subnet
		nodeSubnet, err := s.createNodeSubnet(ctx, zone, workers[i], vpcRes, tags)
		if err != nil {
			return deploy.Panic(err)
		}
		s.nodeSubnets = append(s.nodeSubnets, nodeSubnet)

		// public subnet
		publicSubnet, err := s.createPublicSubnet(ctx, zone, public[i], vpcRes, tags)
		if err != nil {
			return deploy.Panic(err)
		}
		s.publicSubnets = append(s.publicSubnets, publicSubnet)

		// private subnet
		privateSubnet, err := s.createPrivateSubnet(ctx, zone, private[i], vpcRes, tags)
		if err != nil {
			return deploy.Panic(err)
		}
		s.privateSubnets = append(s.privateSubnets, privateSubnet)
	}

	return deploy.Done()
}

func (s *Subnet) createOptimizedSubnets(ctx context.Context, vpcRes *v1beta1.VPC, tags map[string]string, nodeZones []string, zones []string, endpointZones []string) error {
	workers, public, private, controlPlanes, err := util.GetAwsSubnetCIDRsOptimized(s.VPC.Spec.CIDR, nodeZones, zones)
	if err != nil {
		return err
	}
	endpoints := map[string]string{}
	if len(endpointZones) > 0 {
		endpoints, err = util.GetAwsPrivateLinkEndpointSubnetCIDRs(s.VPC.Spec.CIDR, endpointZones)
		if err != nil {
			return err
		}
	}

	for _, zone := range sortedZones(nodeZones) {
		nodeSubnet, createErr := s.createNodeSubnet(ctx, zone, workers[zone], vpcRes, tags)
		if createErr != nil {
			return createErr
		}
		s.nodeSubnets = append(s.nodeSubnets, nodeSubnet)

		publicSubnet, createErr := s.createPublicSubnet(ctx, zone, public[zone], vpcRes, tags)
		if createErr != nil {
			return createErr
		}
		s.publicSubnets = append(s.publicSubnets, publicSubnet)

		privateSubnet, createErr := s.createPrivateSubnet(ctx, zone, private[zone], vpcRes, tags)
		if createErr != nil {
			return createErr
		}
		s.privateSubnets = append(s.privateSubnets, privateSubnet)
	}

	for _, zone := range sortedControlPlaneZones(nodeZones, zones) {
		controlPlaneSubnet, createErr := s.createControlPlaneSubnet(ctx, zone, controlPlanes[zone], vpcRes, tags)
		if createErr != nil {
			return createErr
		}
		s.controlPlaneSubnets = append(s.controlPlaneSubnets, controlPlaneSubnet)
	}

	for _, zone := range sortedZones(endpointZones) {
		endpointSubnet, createErr := s.createPrivateLinkEndpointSubnet(ctx, zone, endpoints[zone], vpcRes, tags)
		if createErr != nil {
			return createErr
		}
		s.endpointSubnets = append(s.endpointSubnets, endpointSubnet)
	}

	return nil
}

func (s *Subnet) createControlPlaneSubnet(ctx context.Context, zone string, cidr string, vpcRes *v1beta1.VPC, tags map[string]string) (*v1beta1.Subnet, error) {
	controlPlaneTags := util.BuildAWSTagWithName(tags, namer.ControlPlaneSubnetName(vpcRes.Name, zone))
	labels := util.MergeMaps(s.VPC.Labels, map[string]string{v1alpha1.VPCNameKey: s.VPC.Name})
	controlPlaneSubnet := &v1beta1.Subnet{
		ObjectMeta: metav1.ObjectMeta{
			Name:   namer.ControlPlaneSubnetName(vpcRes.Name, zone),
			Labels: labels,
		},
		Spec: v1beta1.SubnetSpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{Name: s.ProviderConfigName()},
				ManagementPolicies:      controllers.GetResourceManagementPolicies(),
			},
			ForProvider: v1beta1.SubnetParameters_2{
				Region:           vpcRes.Spec.ForProvider.Region,
				CidrBlock:        pointer.String(cidr),
				AvailabilityZone: pointer.String(zone),
				Tags:             util.ToMapPointer(controlPlaneTags),
				VPCIDRef: &xpv1.Reference{
					Name: vpcRes.Name,
				},
			},
		},
	}
	err := s.createOrUpdateSubnet(ctx, vpcRes, controlPlaneSubnet, labels, controlPlaneTags)
	return controlPlaneSubnet, err
}

func (s *Subnet) createNodeSubnet(ctx context.Context, zone string, cidr string, vpcRes *v1beta1.VPC, tags map[string]string) (*v1beta1.Subnet, error) {
	nodeTags := util.BuildAWSTagWithName(tags, namer.NodesSubnetName(vpcRes.Name, zone))
	labels := util.MergeMaps(s.VPC.Labels, map[string]string{v1alpha1.VPCNameKey: s.VPC.Name})
	nodeSubnet := &v1beta1.Subnet{
		ObjectMeta: metav1.ObjectMeta{
			Name:   namer.NodesSubnetName(vpcRes.Name, zone),
			Labels: labels,
		},
		Spec: v1beta1.SubnetSpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{Name: s.ProviderConfigName()},
				ManagementPolicies:      controllers.GetResourceManagementPolicies(),
			},
			ForProvider: v1beta1.SubnetParameters_2{
				Region:           vpcRes.Spec.ForProvider.Region,
				CidrBlock:        pointer.String(cidr),
				AvailabilityZone: pointer.String(zone),
				Tags:             util.ToMapPointer(nodeTags),
				VPCIDRef: &xpv1.Reference{
					Name: vpcRes.Name,
				},
			},
		},
	}
	err := s.createOrUpdateSubnet(ctx, vpcRes, nodeSubnet, labels, nodeTags)
	return nodeSubnet, err
}

func (s *Subnet) createPublicSubnet(ctx context.Context, zone string, cidr string, vpcRes *v1beta1.VPC, tags map[string]string) (*v1beta1.Subnet, error) {
	subnetTags := util.BuildAWSTagWithName(tags, namer.PublicSubnetName(vpcRes.Name, zone))
	labels := util.MergeMaps(s.VPC.Labels, map[string]string{v1alpha1.VPCNameKey: s.VPC.Name})
	publicSubnet := &v1beta1.Subnet{
		ObjectMeta: metav1.ObjectMeta{
			Name:   namer.PublicSubnetName(vpcRes.Name, zone),
			Labels: labels,
		},
		Spec: v1beta1.SubnetSpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{Name: s.ProviderConfigName()},
				ManagementPolicies:      controllers.GetResourceManagementPolicies(),
			},
			ForProvider: v1beta1.SubnetParameters_2{
				Region:           vpcRes.Spec.ForProvider.Region,
				CidrBlock:        pointer.String(cidr),
				AvailabilityZone: pointer.String(zone),
				Tags:             util.ToMapPointer(subnetTags),
				VPCIDRef: &xpv1.Reference{
					Name: vpcRes.Name,
				},
			},
		},
	}
	// for migrated subenet, we need to update it's owner reference
	err := s.createOrUpdateSubnet(ctx, vpcRes, publicSubnet, labels, subnetTags)
	if err != nil {
		return nil, err
	}
	return publicSubnet, nil
}

func (s *Subnet) createPrivateSubnet(ctx context.Context, zone string, cidr string, vpcRes *v1beta1.VPC, tags map[string]string) (*v1beta1.Subnet, error) {
	privateTags := util.BuildAWSTagWithName(tags, namer.PrivateSubnetName(vpcRes.Name, zone))
	labels := util.MergeMaps(s.VPC.Labels, map[string]string{v1alpha1.VPCNameKey: s.VPC.Name})
	privateSubnet := &v1beta1.Subnet{
		ObjectMeta: metav1.ObjectMeta{
			Name:   namer.PrivateSubnetName(vpcRes.Name, zone),
			Labels: labels,
		},
		Spec: v1beta1.SubnetSpec{
			ResourceSpec: xpv1.ResourceSpec{
				// ProviderConfigReference: &xpv1.Reference{Name: name},
				ProviderConfigReference: &xpv1.Reference{Name: s.ProviderConfigName()},
				ManagementPolicies:      controllers.GetResourceManagementPolicies(),
			},
			ForProvider: v1beta1.SubnetParameters_2{
				Region:           vpcRes.Spec.ForProvider.Region,
				CidrBlock:        pointer.String(cidr),
				AvailabilityZone: pointer.String(zone),
				Tags:             util.ToMapPointer(privateTags),
				VPCIDRef: &xpv1.Reference{
					Name: vpcRes.Name,
				},
			},
		},
	}
	// for migrated subenet, we need to update it's owner reference
	err := s.createOrUpdateSubnet(ctx, vpcRes, privateSubnet, labels, privateTags)
	if err != nil {
		return nil, err
	}
	return privateSubnet, nil
}

func (s *Subnet) createPrivateLinkEndpointSubnet(ctx context.Context, zone string, cidr string, vpcRes *v1beta1.VPC, tags map[string]string) (*v1beta1.Subnet, error) {
	endpointTags := util.BuildAWSTagWithName(tags, namer.PrivateLinkEndpointSubnetName(vpcRes.Name, zone))
	labels := util.MergeMaps(s.VPC.Labels, map[string]string{v1alpha1.VPCNameKey: s.VPC.Name})
	endpointSubnet := &v1beta1.Subnet{
		ObjectMeta: metav1.ObjectMeta{
			Name:   namer.PrivateLinkEndpointSubnetName(vpcRes.Name, zone),
			Labels: labels,
		},
		Spec: v1beta1.SubnetSpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{Name: s.ProviderConfigName()},
				ManagementPolicies:      controllers.GetResourceManagementPolicies(),
			},
			ForProvider: v1beta1.SubnetParameters_2{
				Region:           vpcRes.Spec.ForProvider.Region,
				CidrBlock:        pointer.String(cidr),
				AvailabilityZone: pointer.String(zone),
				Tags:             util.ToMapPointer(endpointTags),
				VPCIDRef: &xpv1.Reference{
					Name: vpcRes.Name,
				},
			},
		},
	}
	err := s.createOrUpdateSubnet(ctx, vpcRes, endpointSubnet, labels, endpointTags)
	return endpointSubnet, err
}

func (s *Subnet) createOrUpdateSubnet(
	ctx context.Context,
	vpcRes *v1beta1.VPC,
	subnet *v1beta1.Subnet,
	desiredLabels map[string]string,
	desiredTags map[string]string,
) error {
	_, err := controllerutil.CreateOrUpdate(ctx, s.Local, subnet, func() error {
		// CreateOrUpdate reloads an existing object before invoking this callback.
		// Reapply inherited labels and tags here while preserving tags owned by
		// other controllers, such as EKS cluster and load balancer role tags.
		subnet.Labels = util.MergeMaps(subnet.Labels, desiredLabels)
		subnet.Spec.ForProvider.Tags = util.MergeMaps(
			subnet.Spec.ForProvider.Tags,
			util.ToMapPointer(desiredTags),
		)
		return controllerutil.SetControllerReference(vpcRes, subnet, s.Local.Scheme())
	})
	return err
}

func (s *Subnet) PostCreate(ctx context.Context) deploy.Result {
	rangeSubnet := func(subnets []*v1beta1.Subnet) ([]v1alpha1.Subnet, error) {
		subnetRecords := make([]v1alpha1.Subnet, 0, len(subnets))
		for _, subnet := range subnets {
			if err := s.Local.Get(ctx, types.NamespacedName{Name: subnet.Name}, subnet); err != nil {
				return nil, err
			}
			if !controllers.IsResourceReady(subnet.Status.ResourceStatus) {
				return nil, fmt.Errorf("subnet %s is not ready", subnet.Name)
			}
			subnetRecords = append(subnetRecords, v1alpha1.Subnet{
				SubnetID:           *subnet.Status.AtProvider.ID,
				Name:               subnet.Name,
				AvailabilityZone:   *subnet.Spec.ForProvider.AvailabilityZone,
				AvailabilityZoneID: *subnet.Status.AtProvider.AvailabilityZoneID,
				CIDRBlock:          *subnet.Spec.ForProvider.CidrBlock,
			})
		}
		return subnetRecords, nil
	}
	var err error
	s.VPC.Status.NodeSubnets, err = rangeSubnet(s.nodeSubnets)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(time.Second * 30)
	}
	s.VPC.Status.PublicSubnets, err = rangeSubnet(s.publicSubnets)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(time.Second * 30)
	}
	s.VPC.Status.PrivateSubnets, err = rangeSubnet(s.privateSubnets)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(time.Second * 30)
	}
	if len(s.controlPlaneSubnets) > 0 {
		s.VPC.Status.ControlPlaneSubnets, err = rangeSubnet(s.controlPlaneSubnets)
		if err != nil {
			return deploy.Panic(err).WithRequeuAfter(time.Second * 30)
		}
	} else {
		s.VPC.Status.ControlPlaneSubnets = nil
	}
	if len(s.endpointSubnets) > 0 {
		s.VPC.Status.PrivateLinkEndpointSubnets, err = rangeSubnet(s.endpointSubnets)
		if err != nil {
			return deploy.Panic(err).WithRequeuAfter(time.Second * 30)
		}
	} else {
		s.VPC.Status.PrivateLinkEndpointSubnets = nil
	}

	// this update won't trigger reconcile
	if err := s.VPCRW.UpdateStatus(ctx, s.VPC); err != nil {
		return deploy.Panic(err).WithRequeuAfter(time.Second * 30)
	}

	return deploy.Done()
}

func (s *Subnet) Delete(ctx context.Context) deploy.Result {
	subnets := &v1beta1.SubnetList{}
	err := s.Local.List(
		ctx,
		subnets,
		client.MatchingLabels{v1alpha1.VPCNameKey: s.VPC.Name},
	)
	if err != nil {
		return deploy.Error(err)
	}

	for i := range subnets.Items {
		subnet := &subnets.Items[i]
		if err := controllers.DeleteUpboundResource(ctx, s.Local, subnet); err != nil {
			return deploy.Error(err)
		}
	}
	return deploy.Done()
}

func sortedZones(zones []string) []string {
	result := append([]string(nil), zones...)
	sort.Strings(result)
	return result
}

func sortedControlPlaneZones(nodeZones []string, zones []string) []string {
	nodeSet := make(map[string]struct{}, len(nodeZones))
	for _, zone := range nodeZones {
		nodeSet[zone] = struct{}{}
	}
	cpZones := make([]string, 0, len(zones)-len(nodeZones))
	for _, zone := range zones {
		if _, ok := nodeSet[zone]; !ok {
			cpZones = append(cpZones, zone)
		}
	}
	sort.Strings(cpZones)
	return cpZones
}

func shouldUseOptimizedCIDR(nodeZones []string, zones []string) bool {
	return len(nodeZones) > 0 && len(nodeZones) <= 2 && len(nodeZones) < len(zones)
}

func (s *Subnet) PostDelete(ctx context.Context) deploy.Result {
	return deploy.Done()
}