package aks

import (
	"context"
	"fmt"

	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/internet/task"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/deploy"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/k8s"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/namer"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/util"
	"github.com/upbound/provider-azure/apis/network/v1beta1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	labelsutil "k8s.io/apimachinery/pkg/labels"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

type NATGateway struct {
	*task.BaseTask
}

func NewNATGateway(ctx *task.Context) *NATGateway {
	return &NATGateway{BaseTask: &task.BaseTask{Context: ctx}}
}

func (t *NATGateway) Name() string {
	return "NATGateway"
}

func (t *NATGateway) Create(ctx context.Context) deploy.Result {
	tags := cloudv1.TagsFromLabels(t.Internet.Labels)
	labels := map[string]string{
		v1alpha1.InternetNameKey: t.Internet.Name,
	}

	nat := t.newNATGateway(labels, tags)
	if err := k8s.Apply(ctx, t.Local, nat); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}

	if err := t.createPublicIPAssociations(ctx, labels, nat); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}

	if err := t.createSubnetAssociations(ctx, labels, nat); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}

	return deploy.Done()
}

func (t *NATGateway) createSubnetAssociations(ctx context.Context, labels map[string]string, nat *v1beta1.NATGateway) error {
	// Associate NAT Gateway with both node subnets and pod subnets
	subnetTypes := []string{v1alpha1.NodeSubnetType, v1alpha1.PodSubnetType}

	for _, subnetType := range subnetTypes {
		subnets := &v1beta1.SubnetList{}
		if err := t.Local.List(ctx, subnets, &client.ListOptions{
			LabelSelector: labelsutil.SelectorFromSet(map[string]string{
				v1alpha1.VPCNameKey:    t.VPC.Name,
				v1alpha1.SubnetTypeKey: subnetType,
			}),
		}); err != nil {
			return err
		}
		if len(subnets.Items) == 0 {
			return fmt.Errorf("%s subnet not found", subnetType)
		}

		for _, subnet := range subnets.Items {
			// use subnet name as association directly
			name := subnet.Name
			association := t.newSubnetAssociation(name, subnet.Name, labels, nat)
			if err := k8s.Apply(ctx, t.Local, association); err != nil {
				return err
			}
		}
	}
	return nil
}

func (t *NATGateway) newSubnetAssociation(name string, subnetName string, labels map[string]string, nat *v1beta1.NATGateway) *v1beta1.SubnetNATGatewayAssociation {
	return &v1beta1.SubnetNATGatewayAssociation{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: labels,
		},
		Spec: v1beta1.SubnetNATGatewayAssociationSpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: t.ProviderConfigName(),
				},
				ManagementPolicies: controllers.GetResourceManagementPolicies(),
			},
			ForProvider: v1beta1.SubnetNATGatewayAssociationParameters{
				NATGatewayIDRef: &xpv1.Reference{
					Name: nat.Name,
				},
				SubnetIDRef: &xpv1.Reference{
					Name: subnetName,
				},
			},
		},
	}
}

func (t *NATGateway) createPublicIPAssociations(ctx context.Context, labels map[string]string, nat *v1beta1.NATGateway) error {
	ips := &v1beta1.PublicIPList{}
	if err := t.Local.List(ctx, ips, &client.ListOptions{
		LabelSelector: labelsutil.SelectorFromSet(map[string]string{
			v1alpha1.InternetNameKey: t.Internet.Name,
		}),
	}); err != nil {
		return err
	}
	if len(ips.Items) == 0 {
		return fmt.Errorf("public IP not found")
	}
	for _, ip := range ips.Items {
		// use PublicIP name as the ip association name directly
		name := ip.Name
		ipAssociation := t.newPublicIPAssociation(name, ip.Name, labels, nat)
		if err := k8s.Apply(ctx, t.Local, ipAssociation); err != nil {
			return err
		}
	}

	return nil
}

func (t *NATGateway) newNATGateway(labels map[string]string, tags map[string]string) *v1beta1.NATGateway {
	return &v1beta1.NATGateway{
		ObjectMeta: metav1.ObjectMeta{
			Name:   getNATGatewayName(t.Internet),
			Labels: labels,
		},
		Spec: v1beta1.NATGatewaySpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: t.ProviderConfigName(),
				},
				ManagementPolicies: controllers.GetResourceManagementPolicies(),
			},
			ForProvider: v1beta1.NATGatewayParameters{
				ResourceGroupName: ptr.To(controllers.GetResourceGroupName(t.VPC, t.Config.ResourceNamePrefix)),
				Location:          ptr.To(t.Internet.Spec.Region),
				SkuName:           ptr.To(standardSKU),
				Tags:              util.ToMapPointer(tags),
			},
		},
	}
}

func (t *NATGateway) newPublicIPAssociation(name string, ipAddressName string, labels map[string]string, nat *v1beta1.NATGateway) *v1beta1.NATGatewayPublicIPAssociation {
	return &v1beta1.NATGatewayPublicIPAssociation{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: labels,
		},
		Spec: v1beta1.NATGatewayPublicIPAssociationSpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: t.ProviderConfigName(),
				},
				ManagementPolicies: controllers.GetResourceManagementPolicies(),
			},
			ForProvider: v1beta1.NATGatewayPublicIPAssociationParameters{
				PublicIPAddressIDRef: &xpv1.Reference{
					Name: ipAddressName,
				},
				NATGatewayIDRef: &xpv1.Reference{
					Name: nat.Name,
				},
			},
		},
	}
}

func getNATGatewayName(internet *v1alpha1.Internet) string {
	return namer.UniqueIdentifier(internet.ObjectMeta, namer.WithPrefix("nat"))
}

func (t *NATGateway) PostCreate(ctx context.Context) deploy.Result {
	if err := t.checkNATGatewayStatus(ctx); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	if err := t.checkPublicIPAssociationStatus(ctx); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	if err := t.checkSubnetAssociationStatus(ctx); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	return deploy.Done()
}

func (t *NATGateway) checkNATGatewayStatus(ctx context.Context) error {
	// check natgateway status
	natName := getNATGatewayName(t.Internet)
	nat := &v1beta1.NATGateway{}
	if err := t.Local.Get(ctx, client.ObjectKey{Name: natName}, nat); err != nil {
		return err
	}
	condition := nat.Status.GetCondition(xpv1.TypeReady)
	if condition.Status != corev1.ConditionTrue {
		return fmt.Errorf("NATGateway %s is not ready", natName)
	}
	natStatus := v1alpha1.NAT{Name: nat.Name}
	if nat.Status.AtProvider.ID != nil {
		natStatus.NATGatewayID = *nat.Status.AtProvider.ID
	}
	t.Internet.Status.NATs = []v1alpha1.NAT{natStatus}
	return t.InternetRW.UpdateStatus(ctx, t.Internet)
}

func (t *NATGateway) checkPublicIPAssociationStatus(ctx context.Context) error {
	ipAssociations := &v1beta1.NATGatewayPublicIPAssociationList{}
	if err := t.Local.List(ctx, ipAssociations, &client.ListOptions{
		LabelSelector: labelsutil.SelectorFromSet(map[string]string{
			v1alpha1.InternetNameKey: t.Internet.Name,
		}),
	}); err != nil {
		return err
	}
	if len(ipAssociations.Items) == 0 {
		return fmt.Errorf("public IP association not found")
	}
	associations := make([]v1alpha1.IPAssociation, 0, len(ipAssociations.Items))
	externalIPs := make([]string, 0, len(ipAssociations.Items))
	for _, association := range ipAssociations.Items {
		condition := association.Status.GetCondition(xpv1.TypeReady)
		if condition.Status != corev1.ConditionTrue {
			return fmt.Errorf("PublicIPAssociation %s is not ready", association.Name)
		}
		associationStatus := v1alpha1.IPAssociation{Name: association.Name}
		if association.Status.AtProvider.ID != nil {
			associationStatus.IPAssociationID = *association.Status.AtProvider.ID
		}
		associations = append(associations, associationStatus)
		ip := &v1beta1.PublicIP{}
		if err := t.Local.Get(ctx, client.ObjectKey{Name: association.Spec.ForProvider.PublicIPAddressIDRef.Name}, ip); err != nil {
			return err
		}
		if ip.Status.AtProvider.IPAddress != nil {
			externalIPs = append(externalIPs, *ip.Status.AtProvider.IPAddress)
		}
	}
	t.Internet.Status.IPAssociations = associations
	t.Internet.Status.ExternalIPs = externalIPs
	return t.InternetRW.UpdateStatus(ctx, t.Internet)
}

func (t *NATGateway) checkSubnetAssociationStatus(ctx context.Context) error {
	subnetAssociations := &v1beta1.SubnetNATGatewayAssociationList{}
	if err := t.Local.List(ctx, subnetAssociations, &client.ListOptions{
		LabelSelector: labelsutil.SelectorFromSet(map[string]string{
			v1alpha1.InternetNameKey: t.Internet.Name,
		}),
	}); err != nil {
		return err
	}
	if len(subnetAssociations.Items) == 0 {
		return fmt.Errorf("subnet association not found")
	}
	associations := make([]v1alpha1.SubnetAssociation, 0, len(subnetAssociations.Items))
	for _, association := range subnetAssociations.Items {
		condition := association.Status.GetCondition(xpv1.TypeReady)
		if condition.Status != corev1.ConditionTrue {
			return fmt.Errorf("SubnetAssociation %s is not ready", association.Name)
		}
		associations = append(associations, v1alpha1.SubnetAssociation{Name: association.Name})
	}
	t.Internet.Status.SubnetAssociations = associations
	return t.InternetRW.UpdateStatus(ctx, t.Internet)
}

func (t *NATGateway) Delete(ctx context.Context) deploy.Result {
	subnetAssociations := &v1beta1.SubnetNATGatewayAssociationList{}
	if err := t.Local.List(ctx, subnetAssociations, &client.ListOptions{
		LabelSelector: labelsutil.SelectorFromSet(map[string]string{
			v1alpha1.InternetNameKey: t.Internet.Name,
		}),
	}); err != nil {
		return deploy.Error(err)
	}

	for _, association := range subnetAssociations.Items {
		if err := controllers.DeleteUpboundResource(ctx, t.Local, ptr.To(association)); err != nil {
			return deploy.Error(err)
		}
	}

	ipAssociation := &v1beta1.NATGatewayPublicIPAssociationList{}
	if err := t.Local.List(ctx, ipAssociation, &client.ListOptions{
		LabelSelector: labelsutil.SelectorFromSet(map[string]string{
			v1alpha1.InternetNameKey: t.Internet.Name,
		}),
	}); err != nil {
		return deploy.Error(err)
	}
	for _, association := range ipAssociation.Items {
		if err := controllers.DeleteUpboundResource(ctx, t.Local, ptr.To(association)); err != nil {
			return deploy.Error(err)
		}
	}

	return deploy.ErrorOrDone(controllers.DeleteUpboundResource(ctx, t.Local, &v1beta1.NATGateway{
		ObjectMeta: metav1.ObjectMeta{
			Name: getNATGatewayName(t.Internet),
		},
	}))
}