package aks

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"maps"
	"strconv"
	"strings"
	"time"

	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	aksv1beta1 "github.com/upbound/provider-azure/apis/containerservice/v1beta1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"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/common/constants"
	"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/nodepool/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/labels"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/nodepool"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/util"
)

const (
	defaultNodeMinSize = 0
	defaultNodeMaxSize = 100
)

var nodePoolManagementPolicies = xpv1.ManagementPolicies{
	// Intentionally skip LateInitialize for autoscaled nodepools.
	// LateInitialize can write provider-observed fields (e.g. nodeCount)
	// back into spec and turn runtime autoscaler state into desired state,
	// causing repeated AKS agentPool update calls.
	xpv1.ManagementActionCreate,
	xpv1.ManagementActionUpdate,
	xpv1.ManagementActionObserve,
}

type aksNodePool struct {
	*task.BaseTask
}

func NewNodePool(ctx *task.Context) deploy.Task {
	return &aksNodePool{BaseTask: &task.BaseTask{Context: ctx}}
}

func (t *aksNodePool) Name() string {
	return "NodePool"
}

// Create provisions AKS nodepools for the NodePool CR.
// It creates one KubernetesClusterNodePool per zone configured in the VPC.
func (t *aksNodePool) Create(ctx context.Context) deploy.Result {
	t.BeginAt = ptr.To(time.Now())

	if t.APPBox.Status.KubernetesClusterName == "" {
		return deploy.Panic(fmt.Errorf("kubernetes cluster name is empty"))
	}

	nodeSubnets := t.VPC.Status.NodeSubnets
	podSubnets := t.VPC.Status.PodSubnets

	if len(nodeSubnets) == 0 {
		return deploy.Error(fmt.Errorf("no node subnets found in VPC status"))
	}
	if len(podSubnets) == 0 {
		return deploy.Error(fmt.Errorf("no pod subnets found in VPC status"))
	}

	metaLabels := util.MergeMaps(t.Namer.Labels(), t.APPBox.Labels)
	metaLabels = util.MergeMaps(metaLabels, labels.NewOwnerLabels(t.NodePool))

	nodeLabels := map[string]*string{
		consts.LabelKeyNodeSpec:   ptr.To(nodepool.NodeSpecLabel(t.NodePool.Spec.CpuCore, t.NodePool.Spec.MemoryGi)),
		consts.K8sVersionLabelKey: ptr.To(t.NodePool.Spec.KubernetesVersion),
	}
	for k, v := range t.NodePool.Spec.Labels {
		val := v
		nodeLabels[k] = &val
	}

	userTaints := task.EnsureNodeTypeLabel(t.NodePool.Spec.Taints, task.NormalizeNodeType(t.NodePool.Spec.NodeType))
	if t.NodePool.Spec.WithLocalDisk {
		userTaints = task.EnsureLocalDiskTaint(userTaints)
		task.EnsureLocalDiskLabels(nodeLabels, metaLabels, t.NodePool.Spec.LocalDiskSizeGi)
	}

	nodeTaints := make([]*string, 0, len(userTaints))
	for _, taint := range userTaints {
		effect := task.NormalizeTaintEffect(taint.Effect)
		nodeTaints = append(nodeTaints, ptr.To(fmt.Sprintf("%s=%s:%s", taint.Key, taint.Value, effect)))
	}

	baseTags := task.CloudTags(consts.ProviderTypeAKS, t.NodePool, t.DataPlane, t.Config)

	availableCPU := int(t.NodePool.Spec.AvailableCpuMilliCore / 1000)
	if availableCPU <= 0 {
		availableCPU = int(t.NodePool.Spec.CpuCore)
	}
	availableMemoryMi := uint64(t.NodePool.Spec.AvailableMemoryMi)
	if availableMemoryMi == 0 {
		availableMemoryMi = nodepool.ComputeAvailableMemory(uint64(t.NodePool.Spec.MemoryGi) * 1024)
	}

	targetZones := t.GetNodeZones()
	if len(targetZones) == 0 {
		seenZones := make(map[string]bool)
		for _, s := range nodeSubnets {
			if s.AvailabilityZone != "" && !seenZones[s.AvailabilityZone] {
				targetZones = append(targetZones, s.AvailabilityZone)
				seenZones[s.AvailabilityZone] = true
			}
		}
	}
	if len(targetZones) == 0 {
		targetZones = []string{""}
	}

	aksName := t.APPBox.Status.KubernetesClusterName
	systemDiskSize := t.Config.APPBoxConfig.WorkerNodeConfig.GetSystemVolumeSize()

	for _, zoneStr := range targetZones {
		zone := getZoneOrdinal(zoneStr)

		nodeSubnet := findSubnetByZoneOrdinal(nodeSubnets, zone)
		if nodeSubnet == nil {
			nodeSubnet = &nodeSubnets[0]
		}

		podSubnet := findSubnetByZoneOrdinal(podSubnets, zone)
		if podSubnet == nil {
			podSubnet = &podSubnets[0]
		}

		npName := t.nodePoolName(zone)

		poolTags := util.ToMapPointer(maps.Clone(baseTags))

		np := &aksv1beta1.KubernetesClusterNodePool{
			ObjectMeta: metav1.ObjectMeta{
				Name:   npName,
				Labels: maps.Clone(metaLabels),
			},
			Spec: aksv1beta1.KubernetesClusterNodePoolSpec{
				ForProvider: aksv1beta1.KubernetesClusterNodePoolParameters{
					KubernetesClusterIDRef: &xpv1.Reference{Name: aksName},
					VMSize:                 ptr.To(t.NodePool.Spec.InstanceType),
					EnableAutoScaling:      ptr.To(true),
					// NodeCount is intentionally not set for autoscaled pools.
					// MinCount/MaxCount define desired bounds; actual replicas
					// are controlled by the AKS cluster autoscaler.
					MinCount: ptr.To(float64(defaultNodeMinSize)),
					MaxCount: ptr.To(float64(defaultNodeMaxSize)),
					// MaxPods set to 16 to optimize Azure reserved memory calculation.
					// See: https://learn.microsoft.com/en-us/azure/aks/node-resource-reservations#memory-reservations
					MaxPods:         ptr.To(float64(16)),
					OsDiskSizeGb:    ptr.To(float64(systemDiskSize)),
					OsSku:           ptr.To("AzureLinux"),
					Mode:            ptr.To("User"), // User mode allows scaling to 0, unlike System mode
					VnetSubnetIDRef: &xpv1.Reference{Name: nodeSubnet.Name},
					PodSubnetIDRef:  &xpv1.Reference{Name: podSubnet.Name},
					Tags:            poolTags,
					NodeTaints:      nodeTaints,
					NodeLabels:      maps.Clone(nodeLabels),
					// Keep maxSurge explicit to avoid AKS defaulting drift.
					UpgradeSettings: []aksv1beta1.KubernetesClusterNodePoolUpgradeSettingsParameters{
						{MaxSurge: ptr.To("10%")},
					},
				},
				ResourceSpec: xpv1.ResourceSpec{
					ProviderConfigReference: &xpv1.Reference{Name: t.ProviderConfigName()},
					ManagementPolicies:      nodePoolManagementPolicies,
				},
			},
		}

		if zone != "" {
			np.Spec.ForProvider.Zones = []*string{ptr.To(zone)}
		}

		// Add capacity tags for monitoring and capacity planning.
		// Azure doesn't allow "/" in tag names, so we use ToAzureTagKey to convert them.
		np.Spec.ForProvider.Tags[controllers.ToAzureTagKey(constants.TagNodeGroupSpecCPU)] = ptr.To(strconv.Itoa(int(t.NodePool.Spec.CpuCore)))
		np.Spec.ForProvider.Tags[controllers.ToAzureTagKey(constants.TagNodeGroupSpecMemoryGi)] = ptr.To(strconv.Itoa(int(t.NodePool.Spec.MemoryGi)))
		np.Spec.ForProvider.Tags[controllers.ToAzureTagKey(constants.TagNodeGroupAvailableCPU)] = ptr.To(strconv.Itoa(availableCPU))
		np.Spec.ForProvider.Tags[controllers.ToAzureTagKey(constants.TagNodeGroupAvailableMemoryMi)] = ptr.To(strconv.FormatUint(availableMemoryMi, 10))

		if err := k8s.Apply(ctx, t.Local, np); err != nil {
			return deploy.Error(err)
		}
	}

	t.NodePool.Status.KubernetesName = aksName
	t.NodePool.Status.Zones = targetZones

	return deploy.Done()
}

// PostCreate verifies that all nodepools have been successfully provisioned by Azure.
func (t *aksNodePool) PostCreate(ctx context.Context) deploy.Result {
	pools := new(aksv1beta1.KubernetesClusterNodePoolList)
	selector := client.MatchingLabels(labels.NewOwnerLabels(t.NodePool))
	if err := t.Local.List(ctx, pools, selector); err != nil {
		return deploy.Error(err)
	}

	for _, pool := range pools.Items {
		condition := pool.Status.GetCondition(xpv1.TypeReady)
		if condition.Status != corev1.ConditionTrue {
			t.NodePool.Status.Conditions = []metav1.Condition{task.ToMetaCondition(condition)}
			return deploy.Error(fmt.Errorf("waiting for nodepool %s to be ready", pool.GetName())).WithRequeuAfter(deploy.DefaultRequeueAfter)
		}
		t.NodePool.Status.Conditions = []metav1.Condition{task.ToMetaCondition(condition)}
	}

	return deploy.Done()
}

// Delete initiates deletion of all nodepools associated with this NodePool CR.
func (t *aksNodePool) Delete(ctx context.Context) deploy.Result {
	pools := new(aksv1beta1.KubernetesClusterNodePoolList)
	selector := client.MatchingLabels(labels.NewOwnerLabels(t.NodePool))
	if err := t.Local.List(ctx, pools, selector); err != nil {
		return deploy.Error(err)
	}
	for i := range pools.Items {
		if err := controllers.DeleteUpboundResource(ctx, t.Local, &pools.Items[i]); err != nil {
			return deploy.Panic(err)
		}
	}
	return deploy.Done()
}

// PostDelete verifies that all nodepools have been completely removed.
func (t *aksNodePool) PostDelete(ctx context.Context) deploy.Result {
	pools := new(aksv1beta1.KubernetesClusterNodePoolList)
	selector := client.MatchingLabels(labels.NewOwnerLabels(t.NodePool))
	if err := t.Local.List(ctx, pools, selector); err != nil {
		return deploy.Panic(err)
	}
	if len(pools.Items) > 0 {
		return deploy.Panic(fmt.Errorf("kubernetesclusternodepools not deleted: %d remaining", len(pools.Items)))
	}
	return deploy.Done()
}

// nodePoolName generates a unique AKS nodepool name within naming constraints.
// AKS nodepool names must be 1-12 characters, lowercase letters and numbers only.
// Format: w{zoneOrdinal}{hash8chars}
//   - 'w' prefix indicates worker nodepool
//   - zone ordinal (1 char): identifies the availability zone
//   - hash (8 chars): MD5 of appBoxName+instanceType+k8sVersion for uniqueness
func (t *aksNodePool) nodePoolName(zone string) string {
	hashInput := t.APPBox.Name + t.NodePool.Spec.InstanceType + t.NodePool.Spec.KubernetesVersion
	if t.NodePool.Annotations[consts.AnnotationKeyBuiltinNodePool] != "true" {
		hashInput += fmt.Sprintf("%dc%dg", t.NodePool.Spec.CpuCore, t.NodePool.Spec.MemoryGi)
	}
	hash := md5.Sum([]byte(hashInput))
	shortHash := hex.EncodeToString(hash[:])[:8]
	return fmt.Sprintf("w%s%s", zone, shortHash)
}

// getZoneOrdinal extracts the zone ordinal from an availability zone string.
// Examples: "eastus-1" -> "1", "westus2-2" -> "2", "1" -> "1", "" -> ""
func getZoneOrdinal(zone string) string {
	if zone == "" {
		return ""
	}
	parts := strings.Split(zone, "-")
	return parts[len(parts)-1]
}

// findSubnetByZoneOrdinal returns the first subnet whose zone ordinal matches the given ordinal.
func findSubnetByZoneOrdinal(subnets []v1alpha1.Subnet, zoneOrdinal string) *v1alpha1.Subnet {
	for i := range subnets {
		if getZoneOrdinal(subnets[i].AvailabilityZone) == zoneOrdinal {
			return &subnets[i]
		}
	}
	return nil
}