package aks

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"time"

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

type adminNodePool struct {
	*task.BaseTask
}

// NewAdminNodePool creates a task that manages additional admin node pools and the
// optional dedicated o11y node pool for the AKS cluster.
//
//   - Admin: one KubernetesClusterNodePool with mode=System and minCount=0 per
//     AdminNodeConfig.AdditionalInstanceTypes entry.
//   - O11y: when O11yNodeConfig.Enabled, one System-mode pool per pools[] entry with dedicated=o11y taint.
//
// This task must run after the AKS task so that KubernetesClusterName is set in APPBox status.
func NewAdminNodePool(ctx *task.Context) deploy.Task {
	return &adminNodePool{BaseTask: &task.BaseTask{Context: ctx}}
}

func (t *adminNodePool) Name() string {
	return "AdminNodePool"
}

func (t *adminNodePool) Create(ctx context.Context) deploy.Result {
	t.BeginAt = ptr.To(time.Now())

	adminConfig := config.GetAdminNodeConfig(
		t.APPBox.Spec.KubernetesConfig.AdminNodeConfig,
		t.Config.APPBoxConfig.AdminNodeConfig,
		DefaultInstanceType,
		DefaultNodeMinSize,
		DefaultNodeMaxSize,
		DefaultSystemVolumeSize,
	)
	o11yPools, err := config.GetO11yNodePools(
		t.Config.APPBoxConfig.O11yNodeConfig,
		0, // default minSize=0 for o11y
		DefaultNodeMaxSize,
		DefaultSystemVolumeSize,
	)
	if err != nil {
		return deploy.Error(err)
	}

	// Always reconcile so Enabled=false / empty pools can clean up previously created
	// o11y pools (and so removing AdditionalInstanceTypes cleans stale admin pools).
	aksName := t.APPBox.Status.KubernetesClusterName
	if aksName == "" {
		return deploy.Error(fmt.Errorf("kubernetes cluster name is empty"))
	}

	nodeSubnets, err := listSubnets(ctx, t.Local, t.VPC.Name, v1alpha1.NodeSubnetType)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	podSubnets, err := listSubnets(ctx, t.Local, t.VPC.Name, v1alpha1.PodSubnetType)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	nodeSubnet := nodeSubnets.Items[0]
	podSubnet := podSubnets.Items[0]

	tags := controllers.ConvertToCloudTags(t.APPBox.Labels, consts.ProviderTypeAKS)
	tags = util.MergeMaps(tags, t.Config.AdditionalCloudTags)
	// Intentionally exclude namer.AppBoxSelectorKey label so these pools are not picked up
	// by the workerNodePool task's list selector.
	metaLabels := util.MergeMaps(t.APPBox.Labels, labels.NewOwnerLabels(t.APPBox))

	// desired must include both additional admin and o11y names (when enabled) so
	// owner-label cleanup does not delete the other role's pools.
	desired := map[string]struct{}{}
	for _, instanceType := range adminConfig.AdditionalInstanceTypes {
		desired[additionalAdminNodePoolName(t.APPBox.Name, instanceType)] = struct{}{}
	}
	for _, pool := range o11yPools {
		desired[o11yNodePoolName(t.APPBox.Name, pool.InstanceType)] = struct{}{}
	}

	for _, instanceType := range adminConfig.AdditionalInstanceTypes {
		npName := additionalAdminNodePoolName(t.APPBox.Name, instanceType)
		np := t.buildSystemNodePool(aksName, npName, instanceType, 0, adminConfig.MaxSize, adminConfig.SystemVolumeSize,
			consts.LabelValComponentShared, nil, nil, metaLabels, tags, nodeSubnet.Name, podSubnet.Name)
		if err := k8s.Apply(ctx, t.Local, np); err != nil {
			return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
		}
	}

	for _, pool := range o11yPools {
		npName := o11yNodePoolName(t.APPBox.Name, pool.InstanceType)
		o11yTaint := ptr.To(fmt.Sprintf("%s=%s:%s", consts.TaintKeyDedicated, consts.TaintValueO11y, consts.TaintEffectNoScheduleForAKS))
		np := t.buildSystemNodePool(aksName, npName, pool.InstanceType, pool.MinSize, pool.MaxSize, pool.SystemVolumeSize,
			consts.LabelValComponentO11y, []*string{o11yTaint}, map[string]string{
				consts.LabelKeyNodeSpec: nodepool.NodeSpecLabel(uint32(pool.CPU), uint32(pool.MemoryGi)),
			}, metaLabels, tags, nodeSubnet.Name, podSubnet.Name)
		if err := k8s.Apply(ctx, t.Local, np); err != nil {
			return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
		}
	}

	if err := t.cleanupStaleAdminNodePools(ctx, desired); err != nil {
		return deploy.Error(fmt.Errorf("failed to clean up stale admin/o11y node pools: %w", err))
	}

	return deploy.Done()
}

func (t *adminNodePool) buildSystemNodePool(
	aksName, npName, instanceType string,
	minCount, maxCount int,
	systemVolumeSize int64,
	component string,
	nodeTaints []*string,
	extraNodeLabels map[string]string,
	metaLabels map[string]string,
	tags map[string]string,
	nodeSubnetName, podSubnetName string,
) *v1beta1.KubernetesClusterNodePool {
	nodeLabels := map[string]*string{
		consts.LabelKeyComponent:  ptr.To(component),
		consts.K8sVersionLabelKey: ptr.To(t.APPBox.Spec.KubernetesConfig.Version),
	}
	for k, v := range extraNodeLabels {
		nodeLabels[k] = ptr.To(v)
	}
	return &v1beta1.KubernetesClusterNodePool{
		ObjectMeta: metav1.ObjectMeta{
			Name:   npName,
			Labels: metaLabels,
		},
		Spec: v1beta1.KubernetesClusterNodePoolSpec{
			ForProvider: v1beta1.KubernetesClusterNodePoolParameters{
				KubernetesClusterIDRef: &xpv1.Reference{Name: aksName},
				VMSize:                 ptr.To(instanceType),
				EnableAutoScaling:      ptr.To(true),
				MinCount:               ptr.To(float64(minCount)),
				MaxCount:               ptr.To(float64(maxCount)),
				OsDiskSizeGb:           ptr.To(float64(systemVolumeSize)),
				OsSku:                  ptr.To("AzureLinux"),
				Mode:                   ptr.To("System"),
				VnetSubnetIDRef:        &xpv1.Reference{Name: nodeSubnetName},
				PodSubnetIDRef:         &xpv1.Reference{Name: podSubnetName},
				Tags:                   util.ToMapPointer(tags),
				NodeTaints:             nodeTaints,
				NodeLabels:             nodeLabels,
				UpgradeSettings: []v1beta1.KubernetesClusterNodePoolUpgradeSettingsParameters{
					{MaxSurge: ptr.To("10%")},
				},
			},
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{Name: t.ProviderConfigName()},
				ManagementPolicies:      controllers.GetResourceManagementPolicies(),
			},
		},
	}
}

// cleanupStaleAdminNodePools deletes additional admin / o11y node pools that are NOT in the desired set.
// If desired is nil, all such pools owned by this APPBox are deleted.
func (t *adminNodePool) cleanupStaleAdminNodePools(ctx context.Context, desired map[string]struct{}) error {
	nps := &v1beta1.KubernetesClusterNodePoolList{}
	if err := t.Local.List(ctx, nps, client.MatchingLabels(labels.NewOwnerLabels(t.APPBox))); err != nil {
		return err
	}
	for i := range nps.Items {
		np := &nps.Items[i]
		if desired != nil {
			if _, ok := desired[np.Name]; ok {
				continue
			}
		}
		if err := controllers.DeleteUpboundResource(ctx, t.Local, np); err != nil {
			return err
		}
	}
	return nil
}

func (t *adminNodePool) PostCreate(ctx context.Context) deploy.Result {
	adminConfig := config.GetAdminNodeConfig(
		t.APPBox.Spec.KubernetesConfig.AdminNodeConfig,
		t.Config.APPBoxConfig.AdminNodeConfig,
		DefaultInstanceType, DefaultNodeMinSize, DefaultNodeMaxSize, DefaultSystemVolumeSize,
	)
	for _, instanceType := range adminConfig.AdditionalInstanceTypes {
		if err := t.waitNodePoolReady(ctx, additionalAdminNodePoolName(t.APPBox.Name, instanceType), "AdminNodePoolNotReady"); err != nil {
			return deploy.Error(err)
		}
	}

	o11yPools, err := config.GetO11yNodePools(
		t.Config.APPBoxConfig.O11yNodeConfig,
		0, DefaultNodeMaxSize, DefaultSystemVolumeSize,
	)
	if err != nil {
		return deploy.Error(err)
	}
	for _, pool := range o11yPools {
		if err := t.waitNodePoolReady(ctx, o11yNodePoolName(t.APPBox.Name, pool.InstanceType), "O11yNodePoolNotReady"); err != nil {
			return deploy.Error(err)
		}
	}

	t.APPBox.Status.Conditions.SetCondition(metav1.Condition{
		Type:               v1alpha1.APPBoxConditionK8SReady,
		Status:             metav1.ConditionTrue,
		Reason:             "NodePoolReady",
		Message:            "AKS admin/o11y node pools are ready",
		LastTransitionTime: metav1.Now(),
	})
	return deploy.Done()
}

func (t *adminNodePool) waitNodePoolReady(ctx context.Context, name, reason string) error {
	np := &v1beta1.KubernetesClusterNodePool{}
	if err := t.Local.Get(ctx, types.NamespacedName{Name: name}, np); err != nil {
		return err
	}
	condition := np.Status.GetCondition(xpv1.TypeReady)
	if condition.Status != corev1.ConditionTrue {
		t.APPBox.Status.Conditions.SetCondition(metav1.Condition{
			Type:               v1alpha1.APPBoxConditionK8SReady,
			Status:             metav1.ConditionFalse,
			Reason:             reason,
			Message:            fmt.Sprintf("node pool %s is not ready", np.GetName()),
			LastTransitionTime: metav1.Now(),
		})
		return fmt.Errorf("waiting for node pool %s to be ready", np.GetName())
	}
	return nil
}

func (t *adminNodePool) Delete(ctx context.Context) deploy.Result {
	// Delete all additional admin / o11y node pools by label selector (not by names derived from config).
	if err := t.cleanupStaleAdminNodePools(ctx, nil); err != nil {
		return deploy.Panic(err)
	}
	return deploy.Done()
}

func (t *adminNodePool) PostDelete(ctx context.Context) deploy.Result {
	nps := &v1beta1.KubernetesClusterNodePoolList{}
	if err := t.Local.List(ctx, nps, client.MatchingLabels(labels.NewOwnerLabels(t.APPBox))); err != nil {
		return deploy.PanicOrDone(err)
	}
	if len(nps.Items) > 0 {
		return deploy.PanicOrDone(fmt.Errorf("admin/o11y node pools still exist: %d remaining", len(nps.Items)))
	}
	return deploy.Done()
}

// additionalAdminNodePoolName generates a unique AKS node pool name for an additional admin node pool.
// AKS pool names must be ≤12 chars, lowercase alphanumeric only.
// Format: a{hash9} — 'a' prefix (admin) + first 9 hex chars of MD5(appBoxName+instanceType).
func additionalAdminNodePoolName(appBoxName, instanceType string) string {
	hash := md5.Sum([]byte(appBoxName + instanceType))
	return fmt.Sprintf("a%s", hex.EncodeToString(hash[:])[:9])
}

// o11yNodePoolName generates a unique AKS node pool name for the dedicated o11y pool.
// Format: o{hash9} — 'o' prefix (o11y) + first 9 hex chars of MD5(appBoxName+instanceType).
// Hash input matches additionalAdminNodePoolName; the 'o' vs 'a' prefix avoids collisions.
func o11yNodePoolName(appBoxName, instanceType string) string {
	hash := md5.Sum([]byte(appBoxName + instanceType))
	return fmt.Sprintf("o%s", hex.EncodeToString(hash[:])[:9])
}
