package aks

import (
	"context"
	"fmt"
	"slices"
	"time"

	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	"github.com/elliotchance/pie/v2"
	"github.com/upbound/provider-azure/apis/containerservice/v1beta1"
	networkv1beta1 "github.com/upbound/provider-azure/apis/network/v1beta1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/cloud/azure"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/common/cidr"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/common/utils"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/config"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/consts"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/appbox/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/util"
)

type AKS struct {
	*task.BaseTask
	AzureCli azure.Client
}

const (
	AKSNameLimit               = 40
	DefaultNodeMaxSize         = 10
	DefaultNodeMinSize         = 2
	DefaultInstanceType        = "Standard_D2s_v6"
	NetworkPlugin              = "azure"
	networkPolicyEngine        = "calico"
	SystemAssignedIdentityType = "SystemAssigned"
	UserAssignedIdentityType   = "UserAssigned"
	// defaultRequeuInterval is the default requeue interval for the task.
	defaultRequeuInterval   = time.Second * 30
	adminNodePoolName       = "adminpool"
	adminNodePoolTempName   = "tempool"
	nodeResourceGroupSuffix = "-node-rg"
	userAssignedNATGateway  = "userAssignedNATGateway"
	DefaultSystemVolumeSize = 50
)

func NewAKS(ctx *task.Context, azureCli azure.Client) *AKS {
	return &AKS{
		BaseTask: &task.BaseTask{Context: ctx},
		AzureCli: azureCli,
	}
}

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

func (t *AKS) Create(ctx context.Context) deploy.Result {
	t.BeginAt = util.Pointer(time.Now())
	aksName := t.getAKSName()
	tags := controllers.ConvertToCloudTags(t.APPBox.Labels, consts.ProviderTypeAKS)
	tags = util.MergeMaps(tags, t.Config.AdditionalCloudTags)

	existingAKS, existingAKSFound, err := t.getExistingAKS(ctx, aksName)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	enableAPIServerVNetIntegration := t.shouldEnableAPIServerVNetIntegration(existingAKS, existingAKSFound)
	useControllerManagedIdentity := t.shouldUseControllerManagedIdentity(enableAPIServerVNetIntegration)

	controlPlaneIdentityResourceID, kubeletIdentityResourceID, useBYOCManagedIdentity, err := getBYOCManagedIdentityResourceIDs(t.Context)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	if useControllerManagedIdentity && !useBYOCManagedIdentity {
		controlPlaneIdentityResourceID, err = aksControlPlaneManagedIdentityResourceID(t.Context)
		if err != nil {
			return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
		}
	}

	// Extract admin node configuration using common utility function
	adminConfig := config.GetAdminNodeConfig(
		t.APPBox.Spec.KubernetesConfig.AdminNodeConfig,
		t.Config.APPBoxConfig.AdminNodeConfig,
		DefaultInstanceType,
		DefaultNodeMinSize,
		DefaultNodeMaxSize,
		DefaultSystemVolumeSize,
	)
	cidrDispatcher, err := cidr.NewAKSCIDRDispatcher(t.VPC.Spec.CIDR)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	if err := cidrDispatcher.Allocate(); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}

	// if no subnet was found, `t.getSubnet` will return an error.
	nodeSubnets, err := t.getSubnet(ctx, v1alpha1.NodeSubnetType)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	nodeSubnet := nodeSubnets.Items[0]

	podSubnets, err := t.getSubnet(ctx, v1alpha1.PodSubnetType)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	podSubnet := podSubnets.Items[0]

	var apiServerSubnet *networkv1beta1.Subnet
	if enableAPIServerVNetIntegration {
		apiServerSubnets, err := t.getSubnet(ctx, v1alpha1.APIServerSubnetType)
		if err != nil {
			return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
		}
		apiServerSubnet = &apiServerSubnets.Items[0]
	}
	zoneOrdinals := pie.Map(t.VPC.Spec.Zones, getZoneOrdinal)

	aks := &v1beta1.KubernetesCluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:   aksName,
			Labels: util.MergeMaps(t.Namer.Labels(), t.APPBox.Labels, tags),
		},
		Spec: v1beta1.KubernetesClusterSpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: t.ProviderConfigName(),
				},
				ManagementPolicies: controllers.GetResourceManagementPolicies(),
			},
			ForProvider: v1beta1.KubernetesClusterParameters{
				KubernetesVersion:       ptr.To(t.APPBox.Spec.KubernetesConfig.Version),
				Location:                ptr.To(t.APPBox.Spec.Region),
				OidcIssuerEnabled:       ptr.To(true),
				WorkloadIdentityEnabled: ptr.To(true),
				ResourceGroupName:       ptr.To(controllers.GetResourceGroupName(t.VPC, t.Config.ResourceNamePrefix)),
				// NodeResourceGroupName will autogenerated syntax of MC_resourceGroupName_resourceName_AzureRegion must be no greater than 80
				// So we need to set it customly to avoid the name too long
				// https://learn.microsoft.com/en-us/troubleshoot/azure/azure-kubernetes/create-upgrade-delete/aks-common-issues-faq
				NodeResourceGroup:    ptr.To(t.getNodeResourceGroupName()),
				NodeOsChannelUpgrade: ptr.To("None"),
				Identity:             []v1beta1.IdentityParameters{{Type: ptr.To(SystemAssignedIdentityType)}},
				NetworkProfile: []v1beta1.NetworkProfileParameters{
					{
						NetworkPlugin: ptr.To(NetworkPlugin),
						ServiceCidr:   ptr.To(cidrDispatcher.ServiceIPv4CIDRBlock()),
						DNSServiceIP:  ptr.To(cidrDispatcher.DNSServiceIP()),
						OutboundType:  ptr.To(userAssignedNATGateway),
						// Use Calico instead of Azure Network Policy Manager (NPM) or Cilium
						//   - Because NPM has some limitations, and it will be retired on September 30, 2028.
						//   - If using Cilium, we have to enable Cilium dataplane.
						// See: https://learn.microsoft.com/en-us/azure/aks/use-network-policies#azure-network-policy-manager-limitations-linux
						NetworkPolicy: ptr.To(networkPolicyEngine),
					},
				},
				DNSPrefix: ptr.To(aksName),
				AutoScalerProfile: []v1beta1.AutoScalerProfileParameters{
					{
						BalanceSimilarNodeGroups: ptr.To(true),
					},
				},
				DefaultNodePool: []v1beta1.DefaultNodePoolParameters{
					{
						// specify the temporary name for the node pool rotation
						TemporaryNameForRotation: ptr.To(adminNodePoolTempName),
						EnableAutoScaling:        ptr.To(true),
						MaxCount:                 ptr.To(float64(adminConfig.MaxSize)),
						MinCount:                 ptr.To(float64(adminConfig.MinSize)),
						// use the aks name as the admin node pool name
						OsSku:           ptr.To("AzureLinux"),
						Name:            ptr.To(adminNodePoolName),
						VMSize:          ptr.To(adminConfig.InstanceType),
						OsDiskSizeGb:    ptr.To(float64(adminConfig.SystemVolumeSize)),
						VnetSubnetIDRef: &xpv1.Reference{Name: nodeSubnet.Name},
						PodSubnetIDRef:  &xpv1.Reference{Name: podSubnet.Name},
						UpgradeSettings: []v1beta1.UpgradeSettingsParameters{
							{
								// if it is not set, the default value is 10%, and aks will be frequently update by crossplane
								// because it is not set by crossplane lateinitialize
								// https://github.com/crossplane-contrib/provider-upjet-azure/issues/686
								MaxSurge: ptr.To(string("10%")),
							},
						},
						Tags:  util.ToMapPointer(tags),
						Zones: util.ToSlicePointer(zoneOrdinals),
					},
				},
				SkuTier: ptr.To("Standard"),
				Tags:    util.ToMapPointer(tags),
				AzureActiveDirectoryRoleBasedAccessControl: []v1beta1.AzureActiveDirectoryRoleBasedAccessControlParameters{
					{
						Managed: ptr.To(true),
					},
				},
			},
		},
	}
	if useBYOCManagedIdentity {
		aks.Spec.ForProvider.Identity = []v1beta1.IdentityParameters{
			{
				Type:        ptr.To(UserAssignedIdentityType),
				IdentityIds: []*string{ptr.To(controlPlaneIdentityResourceID)},
			},
		}
		aks.Spec.ForProvider.KubeletIdentity = []v1beta1.KubeletIdentityParameters{
			{
				UserAssignedIdentityID: ptr.To(kubeletIdentityResourceID),
			},
		}
	} else if useControllerManagedIdentity {
		aks.Spec.ForProvider.Identity = []v1beta1.IdentityParameters{
			{
				Type:        ptr.To(UserAssignedIdentityType),
				IdentityIds: []*string{ptr.To(controlPlaneIdentityResourceID)},
			},
		}
	}
	adminGroupIDs := make([]*string, 0, len(t.Config.APPBoxConfig.AdminGroupObjectIDs))
	for _, id := range t.Config.APPBoxConfig.AdminGroupObjectIDs {
		adminGroupIDs = append(adminGroupIDs, ptr.To(id))
	}
	if len(adminGroupIDs) > 0 {
		aks.Spec.ForProvider.AzureActiveDirectoryRoleBasedAccessControl[0].AdminGroupObjectIds = adminGroupIDs
	}
	if enableAPIServerVNetIntegration {
		util.SetAnnotation(aks, v1alpha1.AnnotationKeyAKSAPIServerVNetIntegrationRequested, "true")
	}

	if t.IsBYOC() {
		aks.Spec.ForProvider.APIServerAccessProfile = []v1beta1.APIServerAccessProfileParameters{
			{
				SubnetIDRef:            &xpv1.Reference{Name: apiServerSubnet.Name},
				VnetIntegrationEnabled: ptr.To(true),
			},
		}
		aks.Spec.ForProvider.PrivateClusterEnabled = ptr.To(true)
		aks.Spec.ForProvider.PrivateDNSZoneID = ptr.To("System")
		aks.Spec.ForProvider.PrivateClusterPublicFqdnEnabled = ptr.To(false)
		aks.Spec.ForProvider.PublicNetworkAccessEnabled = ptr.To(false)
	} else {
		accessProfile := v1beta1.APIServerAccessProfileParameters{}
		hasAccessProfile := false

		if enableAPIServerVNetIntegration {
			accessProfile.SubnetIDRef = &xpv1.Reference{Name: apiServerSubnet.Name}
			accessProfile.VnetIntegrationEnabled = ptr.To(true)
			hasAccessProfile = true
		}

		allowCidrs := t.Config.APPBoxConfig.ApiserverPublicAccessCidrs
		if len(allowCidrs) > 0 && !utils.IsIncludeString(allowCidrs, "0.0.0.0/0") {
			internet, err := t.getActiveInternet()
			if err != nil {
				return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
			}
			ips, err := t.getNATIPs(internet)
			if err != nil {
				return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
			}
			allowCidrs = append(slices.Clone(allowCidrs), ips...)
		}

		if len(allowCidrs) > 0 {
			accessProfile.AuthorizedIPRanges = util.ToSlicePointer(allowCidrs)
			hasAccessProfile = true
		}

		if hasAccessProfile {
			aks.Spec.ForProvider.APIServerAccessProfile = []v1beta1.APIServerAccessProfileParameters{accessProfile}
		}
	}

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

	return deploy.Done()
}

func (t *AKS) getExistingAKS(ctx context.Context, name string) (*v1beta1.KubernetesCluster, bool, error) {
	cluster := &v1beta1.KubernetesCluster{}
	if err := t.Local.Get(ctx, types.NamespacedName{Name: name}, cluster); err != nil {
		if apierrors.IsNotFound(err) {
			return nil, false, nil
		}
		return nil, false, err
	}
	return cluster, true, nil
}

func (t *AKS) shouldEnableAPIServerVNetIntegration(cluster *v1beta1.KubernetesCluster, clusterExists bool) bool {
	if t.IsBYOC() {
		return true
	}
	if !clusterExists {
		return true
	}
	if !isAKSProvisioned(cluster) {
		return true
	}
	if t.APPBox.Annotations[v1alpha1.AnnotationKeyEnableAKSAPIServerVNetIntegration] == "true" {
		return true
	}
	if isAKSAPIServerVNetIntegrationObservedEnabled(cluster) {
		return true
	}
	if cluster.Annotations[v1alpha1.AnnotationKeyAKSAPIServerVNetIntegrationRequested] == "true" {
		return true
	}
	return false
}

func (t *AKS) shouldUseControllerManagedIdentity(enableAPIServerVNetIntegration bool) bool {
	return enableAPIServerVNetIntegration && !t.IsBYOC()
}

func isAKSProvisioned(cluster *v1beta1.KubernetesCluster) bool {
	return cluster != nil && ptr.Deref(cluster.Status.AtProvider.ID, "") != ""
}

func isAKSAPIServerVNetIntegrationObservedEnabled(cluster *v1beta1.KubernetesCluster) bool {
	for _, profile := range cluster.Status.AtProvider.APIServerAccessProfile {
		if ptr.Deref(profile.VnetIntegrationEnabled, false) {
			return true
		}
	}
	return false
}

func (t *AKS) getNATIPs(internet *v1alpha1.Internet) ([]string, error) {
	if len(internet.Status.ExternalIPs) == 0 {
		return nil, fmt.Errorf("no external ip found in internet %s", internet.Name)
	}
	ips := make([]string, 0, len(internet.Status.ExternalIPs))
	for _, ip := range internet.Status.ExternalIPs {
		ips = append(ips, fmt.Sprintf("%s/32", ip))
	}
	return ips, nil
}

func (t *AKS) getActiveInternet() (*v1alpha1.Internet, error) {
	internets := &v1alpha1.InternetList{}
	ops := &client.ListOptions{
		LabelSelector: labels.SelectorFromSet(map[string]string{
			v1alpha1.VPCNameKey: t.VPC.Name,
		}),
	}
	if err := t.InternetR.List(context.Background(), internets, ops); err != nil {
		return nil, err
	}
	if len(internets.Items) == 0 {
		return nil, fmt.Errorf("no internet found for vpc %s", t.VPC.Name)
	}
	for _, internet := range internets.Items {
		if internet.Status.Phase == v1alpha1.InternetPhaseActive {
			return ptr.To(internet), nil
		}
	}
	return nil, fmt.Errorf("no active internet found for vpc %s", t.VPC.Name)
}

func (t *AKS) getAKSName() string {
	return aksClusterName(t.APPBox)
}

func (t *AKS) getNodeResourceGroupName() string {
	return aksNodeResourceGroupName(t.APPBox)
}

func (t *AKS) getSubnet(ctx context.Context, subnetType string) (*networkv1beta1.SubnetList, error) {
	return listSubnets(ctx, t.Local, t.VPC.Name, subnetType)
}

// listSubnets lists Azure subnets by VPC name and subnet type.
func listSubnets(ctx context.Context, cli client.Client, vpcName, subnetType string) (*networkv1beta1.SubnetList, error) {
	subnets := &networkv1beta1.SubnetList{}
	if err := cli.List(ctx, subnets, &client.ListOptions{
		LabelSelector: labels.SelectorFromSet(map[string]string{
			v1alpha1.VPCNameKey:    vpcName,
			v1alpha1.SubnetTypeKey: subnetType,
		}),
	}); err != nil {
		return nil, err
	}
	if len(subnets.Items) == 0 {
		return nil, fmt.Errorf("no subnet found for vpc %s and type %s", vpcName, subnetType)
	}
	return subnets, nil
}

func (t *AKS) PostCreate(ctx context.Context) (res deploy.Result) {
	key := types.NamespacedName{Name: t.getAKSName()}
	aks := &v1beta1.KubernetesCluster{}
	err := t.Local.Get(ctx, key, aks)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	t.AKSCluster = aks
	t.Cluster = task.NewAKSClusterIdentifier(aks)
	t.APPBox.Status.KubernetesClusterName = aks.Name
	t.APPBox.Status.KubernetesClusterID = ptr.Deref(aks.Status.AtProvider.ID, "")
	condition := aks.Status.GetCondition(xpv1.TypeReady)
	// If no Ready condition exists or it's not True, cluster is still creating
	if condition.Type != xpv1.TypeReady || condition.Status != corev1.ConditionTrue {
		t.APPBox.Status.Conditions.SetCondition(metav1.Condition{
			Type:               v1alpha1.APPBoxConditionK8SReady,
			Status:             metav1.ConditionFalse,
			Reason:             "AKSCreating",
			Message:            "AKS Cluster is creating",
			LastTransitionTime: metav1.Now(),
		})
		if err := t.APPBoxRW.UpdateStatus(ctx, t.APPBox); err != nil {
			return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
		}
		return deploy.Panic(fmt.Errorf("AKS Cluster is creating")).WithRequeuAfter(defaultRequeuInterval)
	}

	secret := &corev1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:      t.getAKSName(),
			Namespace: t.APPBox.Namespace,
			Labels: map[string]string{
				v1alpha1.APPBoxNameKey: t.APPBox.Name,
			},
		},
	}
	if t.IsBYOC() {
		secretExists, err := t.kubeconfigSecretExists(ctx, secret)
		if err != nil {
			return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
		}
		if !secretExists || t.APPBox.Status.KubeConfigSecretRef == nil {
			t.APPBox.Status.Conditions.SetCondition(metav1.Condition{
				Type:               v1alpha1.APPBoxConditionK8SReady,
				Status:             metav1.ConditionFalse,
				Reason:             "WaitingForPrivateEndpoint",
				Message:            "Waiting for the AKS API server private endpoint",
				LastTransitionTime: metav1.Now(),
			})
		} else {
			t.APPBox.Status.Conditions.SetCondition(metav1.Condition{
				Type:               v1alpha1.APPBoxConditionK8SReady,
				Status:             metav1.ConditionTrue,
				Reason:             "AKSReady",
				Message:            "AKS cluster is ready",
				LastTransitionTime: metav1.Now(),
			})
		}
	} else {
		t.APPBox.Status.Conditions.SetCondition(metav1.Condition{
			Type:               v1alpha1.APPBoxConditionK8SReady,
			Status:             metav1.ConditionTrue,
			Reason:             "AKSReady",
			Message:            "AKS cluster is ready",
			LastTransitionTime: metav1.Now(),
		})

		if err := t.ensureKubeconfigSecret(ctx, secret, aks); err != nil {
			return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
		}

		t.APPBox.Status.KubeConfigSecretRef = &corev1.SecretReference{
			Name:      secret.Name,
			Namespace: secret.Namespace,
		}
	}

	// Set the kubelet identity client ID for Azure
	if len(aks.Status.AtProvider.KubeletIdentity) > 0 {
		if aks.Status.AtProvider.KubeletIdentity[0].ClientID != nil {
			t.APPBox.Status.IAM.KubeletIdentityClientID = *aks.Status.AtProvider.KubeletIdentity[0].ClientID
		}
		if aks.Status.AtProvider.KubeletIdentity[0].ObjectID != nil {
			t.APPBox.Status.IAM.KubeletIdentityObjectID = *aks.Status.AtProvider.KubeletIdentity[0].ObjectID
		}
	}

	return deploy.PanicOrDone(t.APPBoxRW.UpdateStatus(ctx, t.APPBox))
}

func (t *AKS) kubeconfigSecretExists(ctx context.Context, secret *corev1.Secret) (bool, error) {
	if err := t.Local.Get(ctx, client.ObjectKeyFromObject(secret), secret); err != nil {
		if apierrors.IsNotFound(err) {
			return false, nil
		}
		return false, fmt.Errorf("failed to get secret: %w", err)
	}
	return true, nil
}

func (t *AKS) ensureKubeconfigSecret(ctx context.Context, secret *corev1.Secret, cluster *v1beta1.KubernetesCluster) error {
	return ensureAKSKubeconfigSecret(ctx, t.BaseTask, t.AzureCli, secret, cluster)
}

func (t *AKS) createKubeconfigSecret(ctx context.Context, secret *corev1.Secret, cluster *v1beta1.KubernetesCluster) error {
	return createAKSKubeconfigSecret(ctx, t.BaseTask, t.AzureCli, secret, cluster)
}

func (t *AKS) Delete(ctx context.Context) deploy.Result {
	aks := &v1beta1.KubernetesCluster{
		ObjectMeta: metav1.ObjectMeta{
			Name: t.getAKSName(),
		},
	}
	if err := controllers.DeleteUpboundResource(ctx, t.Local, aks); err != nil {
		return deploy.Panic(err)
	}
	return deploy.PanicOrDone(k8s.DeleteObject(ctx, t.Local, &corev1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:      t.getAKSName(),
			Namespace: t.APPBox.Namespace,
		},
	}))
}

func (t *AKS) PostDelete(ctx context.Context) deploy.Result {
	if err := k8s.GetDeletedObject(ctx, t.Local, types.NamespacedName{
		Name: t.getAKSName(),
	}, &v1beta1.KubernetesCluster{}); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	if err := k8s.GetDeletedObject(ctx, t.Local, types.NamespacedName{
		Name:      t.getAKSName(),
		Namespace: t.APPBox.Namespace,
	}, &corev1.Secret{}); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	return deploy.Done()
}