package commongcp

import (
	"context"
	"fmt"
	"maps"
	"strconv"
	"strings"
	"time"

	gkecontainerv1beta1 "github.com/crossplane-contrib/provider-gcp/apis/container/v1beta1"
	gkecontainerv1beta2 "github.com/crossplane-contrib/provider-gcp/apis/container/v1beta2"
	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	"github.com/tidbcloud/observability/agent-control-server/pkg/agentbuild"
	"github.com/tidbcloud/observability/agent-control-server/pkg/config"
	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
	"github.com/tidbcloud/observability/agent-control-server/pkg/deploy"
	"github.com/tidbcloud/observability/agent-control-server/pkg/reconciler"
	"github.com/tidbcloud/observability/agent-control-server/pkg/util"
)

const (
	defaultMinSize     = 1
	defaultMaxSize     = 3
	defaultDesiredSize = 1
	defaultDiskSize    = 100 // in GB
	defaultSpecName    = "default"
)

// NodePoolTask manages the GCP NodePool for the observability agent.
//
// NodePoolTask creates, updates, and deletes the NodePool based on the desired specifications.
// It also handles scaling operations and ensures the NodePool is in the desired state.
// The task uses annotations on the agent resource to track the desired and previous specifications.
type NodePoolTask[T client.Object] struct {
	*reconciler.BaseTask[T]
	clusterName      string
	nameFunc         func(string) string
	specs            config.CustomerAgentSpecs
	taints           []*gkecontainerv1beta1.NodeTaint
	nodeMatchLabels  map[string]string
	nodeConfigLabels map[string]string
}

func NewNodePoolTask[T client.Object](
	ctx *reconciler.Context[T],
	clusterName string,
	nameFunc func(string) string,
	nodeMatchLabels map[string]string,
	nodeConfigLabels map[string]string,
	specs config.CustomerAgentSpecs,
	taints []*gkecontainerv1beta1.NodeTaint,
) *NodePoolTask[T] {
	return &NodePoolTask[T]{
		BaseTask:         &reconciler.BaseTask[T]{Context: ctx},
		clusterName:      clusterName,
		nameFunc:         nameFunc,
		specs:            specs,
		taints:           taints,
		nodeMatchLabels:  nodeMatchLabels,
		nodeConfigLabels: nodeConfigLabels,
	}
}

func (t *NodePoolTask[T]) Name() string {
	return "NodePool"
}

func (t *NodePoolTask[T]) Create(ctx context.Context) deploy.Result {
	return deploy.PanicOrDone(t.create(ctx))
}

func (t *NodePoolTask[T]) PostCreate(ctx context.Context) deploy.Result {
	desiredSpec := util.GetAnnotationWithDefault(t.Agent, consts.AnnotationKeyO11yDesiredSpec, defaultSpecName)
	name := t.nameFunc(desiredSpec)
	if err := agentbuild.CheckNodePoolReady(ctx, t.Local, name); err != nil {
		t.clearStalledExternalCreate(ctx, name)
		return deploy.Panic(err).WithRequeueAfter(time.Minute)
	}
	switch util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus) {
	case consts.SpecScaleStatusPendingScale:
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus, consts.SpecScaleStatusNewNodeGroupSynced)
	}
	return deploy.Done()
}

// clearStalledExternalCreate best-effort clears crossplane's stalled
// external-create-pending annotation on the nodepool so crossplane resumes
// reconciling it. See agentbuild.ClearStalledExternalCreate for why this is safe.
func (t *NodePoolTask[T]) clearStalledExternalCreate(ctx context.Context, name string) {
	cleared, err := agentbuild.ClearStalledExternalCreate(ctx, t.Local,
		&gkecontainerv1beta1.NodePool{ObjectMeta: metav1.ObjectMeta{Name: name}})
	switch {
	case err != nil:
		t.L.Error(err, "failed to clear stalled crossplane external-create-pending annotation", "name", name)
	case cleared:
		t.L.Info("cleared stalled crossplane external-create-pending annotation for crossplane to resume", "name", name)
	}
}

func (t *NodePoolTask[T]) Delete(ctx context.Context) deploy.Result {
	// capable with legacy
	if err := agentbuild.DeleteNodePool(ctx, t.Local, t.nameFunc(defaultSpecName)); err != nil {
		return deploy.Panic(err)
	}
	return deploy.PanicOrDone(agentbuild.DeleteNodePoolWithLabels(ctx, t.Local, t.nodeMatchLabels))
}

func (t *NodePoolTask[T]) PostDelete(ctx context.Context) deploy.Result {
	// capable with legacy
	name := t.nameFunc(defaultSpecName)
	if err := agentbuild.CheckNodePoolDeleted(ctx, t.Local, name); err != nil {
		t.clearStalledExternalCreate(ctx, name)
		return deploy.Panic(err).WithRequeueAfter(time.Minute)
	}
	if err := agentbuild.CheckNodePoolDeletedWithLabels(ctx, t.Local, t.nodeMatchLabels); err != nil {
		return deploy.Panic(err).WithRequeueAfter(time.Minute)
	}
	return deploy.Done()
}

func (t *NodePoolTask[T]) create(ctx context.Context) error {
	desireSpec := util.GetAnnotationWithDefault(t.Agent, consts.AnnotationKeyO11yDesiredSpec, defaultSpecName)
	spec := t.specs.Select(desireSpec)
	if spec == nil {
		return fmt.Errorf("desired spec not found: %s", desireSpec)
	}
	if err := t.reconcile(ctx, t.nameFunc(desireSpec), spec.InstanceType); err != nil {
		return err
	}

	switch util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus) {
	case consts.SpecScaleStatusPendingScale:
		// If we are in a pending-scale state, remember to reconcile the original spec
		// until the new spec is ready. This is to avoid a situation where the old spec
		// cannot be further edited when the new spec is delayed.
		previousSpec := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yPreviousSpec)
		if previousSpec == "" {
			return fmt.Errorf("unexpected empty previous spec name")
		}
		spec := t.specs.Select(previousSpec)
		if spec == nil {
			return fmt.Errorf("previous spec not found: %s", previousSpec)
		}
		if err := t.reconcile(ctx, t.nameFunc(previousSpec), spec.InstanceType); err != nil {
			return err
		}

	case consts.SpecScaleStatusResourceUpdated:
		previousSpec := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yPreviousSpec)
		if previousSpec == "" {
			return fmt.Errorf("unexpected empty previous spec name")
		}
		prevName := t.nameFunc(previousSpec)
		if err := agentbuild.DeleteNodePool(ctx, t.Local, prevName); err != nil {
			return fmt.Errorf("failed to delete previous nodepool: %w", err)
		}
		if err := agentbuild.CheckNodePoolDeleted(ctx, t.Local, prevName); err != nil {
			t.clearStalledExternalCreate(ctx, prevName)
			return deploy.NewRetriableError(err, ctrl.Result{Requeue: true, RequeueAfter: time.Minute})
		}
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus, consts.SpecScaleStatusNormal)
	}
	return nil
}

func (t *NodePoolTask[T]) updateSizeAnnotation(key string, val int) (int, error) {
	s := util.SetAnnotationIfEmpty(t.Agent, key, strconv.Itoa(val))
	i, err := strconv.Atoi(s)
	if err != nil {
		return 0, fmt.Errorf("invalid key '%s': %w", key, err)
	}
	return i, nil
}

// reconcile ensures that a NodePool with the given name and instance type exists
// and is configured according to the specifications stored in the agent's annotations.
// It retrieves necessary configuration details from annotations or falls back to
// an existing admin NodePool if the annotations are missing. It also updates size-related
// annotations on the agent as needed.
func (t *NodePoolTask[T]) reconcile(ctx context.Context, name, instanceType string) error {
	cluster := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeGKECluster)
	clusterRef := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeGKEClusterRef)
	nodeSA := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeSA)
	providerConfigReference := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeProviderConfigReference)
	imageType := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeImageType)
	version := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeVersion)
	zone := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeZone)
	if cluster == "" || clusterRef == "" || nodeSA == "" || providerConfigReference == "" || imageType == "" || version == "" || zone == "" {
		var nodepools gkecontainerv1beta1.NodePoolList
		if err := t.Local.List(ctx, &nodepools, client.MatchingFields{
			".spec.forProvider.clusterRef.name": t.clusterName,
		}); err != nil {
			return fmt.Errorf("failed to list nodepool: %w", err)
		}
		var np *gkecontainerv1beta1.NodePool
		for _, item := range nodepools.Items {
			if strings.HasPrefix(item.Name, "tidbcloud-dedicated-") || strings.HasPrefix(item.Name, "tidbcloud-prod-") {
				np = &item
				break
			}
		}
		if np == nil {
			return fmt.Errorf("admin nodepool not found")
		}
		if err := verifyNodePoolSpec(np.Name, np.Spec); err != nil {
			return err
		}

		cluster = np.Spec.ForProvider.Cluster
		clusterRef = np.Spec.ForProvider.ClusterRef.Name
		nodeSA = *np.Spec.ForProvider.Config.ServiceAccount
		providerConfigReference = np.Spec.ProviderConfigReference.Name
		imageType = *np.Spec.ForProvider.Config.ImageType
		version = *np.Spec.ForProvider.Version
		zone = np.Spec.ForProvider.Locations[0]
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeGKECluster, cluster)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeGKEClusterRef, clusterRef)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeSA, nodeSA)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeProviderConfigReference, providerConfigReference)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeImageType, imageType)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeVersion, version)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeZone, zone)
	}

	minSize, err := t.updateSizeAnnotation(consts.AnnotationKeyO11yNodeMinSize, defaultMinSize)
	if err != nil {
		return err
	}
	maxSize, err := t.updateSizeAnnotation(consts.AnnotationKeyO11yNodeMaxSize, defaultMaxSize)
	if err != nil {
		return err
	}
	desiredSize, err := t.updateSizeAnnotation(consts.AnnotationKeyO11yNodeDesiredSize, defaultDesiredSize)
	if err != nil {
		return err
	}
	diskSize, err := t.updateSizeAnnotation(consts.AnnotationKeyO11yNodeDiskSize, defaultDiskSize)
	if err != nil {
		return err
	}

	nodepool := &gkecontainerv1beta1.NodePool{
		ObjectMeta: metav1.ObjectMeta{Name: name},
	}
	_, err = controllerutil.CreateOrUpdate(ctx, t.Local, nodepool, func() error {
		if nodepool.Labels == nil {
			nodepool.Labels = make(map[string]string)
		}
		for k, v := range t.nodeMatchLabels {
			nodepool.Labels[k] = v
		}

		// Since the nodepool cannot be perceived by the caller, a separate dimension merge is required here.
		nodeConfigLbs := maps.Clone(t.nodeConfigLabels)
		nodeConfigLbs["Name"] = nodepool.Name
		nodeConfigLbs["tidbcloud.com/nodegroup-name"] = nodepool.Name

		nodepool.Spec = gkecontainerv1beta1.NodePoolSpec{
			ResourceSpec: xpv1.ResourceSpec{
				DeletionPolicy: xpv1.DeletionDelete,
				ProviderConfigReference: &xpv1.Reference{
					Name: providerConfigReference,
				},
			},
			ForProvider: gkecontainerv1beta1.NodePoolParameters{
				Autoscaling: &gkecontainerv1beta1.NodePoolAutoscaling{
					Enabled:      ptr.To(true),
					MaxNodeCount: ptr.To(int64(maxSize)),
					MinNodeCount: ptr.To(int64(minSize)),
				},
				Cluster:    cluster,
				ClusterRef: &xpv1.Reference{Name: clusterRef},
				Config: &gkecontainerv1beta1.NodeConfig{
					DiskSizeGb:  ptr.To(int64(diskSize)),
					DiskType:    ptr.To("pd-standard"),
					ImageType:   &imageType,
					Labels:      nodeConfigLbs,
					MachineType: &instanceType,
					Metadata: map[string]string{
						"disable-legacy-endpoints": "true",
					},
					OauthScopes:    []string{"https://www.googleapis.com/auth/cloud-platform"},
					ServiceAccount: &nodeSA,
					ShieldedInstanceConfig: &gkecontainerv1beta1.ShieldedInstanceConfig{
						EnableIntegrityMonitoring: ptr.To(true),
					},
					Taints: t.taints,
					WorkloadMetadataConfig: &gkecontainerv1beta1.WorkloadMetadataConfig{
						Mode: "GKE_METADATA",
					},
				},
				InitialNodeCount:  ptr.To(int64(desiredSize)),
				Locations:         []string{zone},
				Management:        &gkecontainerv1beta1.NodeManagementSpec{AutoRepair: ptr.To(false)},
				MaxPodsConstraint: &gkecontainerv1beta2.MaxPodsConstraint{MaxPodsPerNode: 32},
				UpgradeSettings:   &gkecontainerv1beta2.UpgradeSettings{MaxSurge: ptr.To(int64(1))},
				Version:           &version,
			},
		}
		return nil
	})
	return err
}

func verifyNodePoolSpec(name string, spec gkecontainerv1beta1.NodePoolSpec) error {
	if spec.ForProvider.ClusterRef == nil {
		return fmt.Errorf("clusterRef is nil in nodepool %s", name)
	}
	if spec.ForProvider.Config == nil {
		return fmt.Errorf("config is nil in nodepool %s", name)
	}
	if spec.ForProvider.Config.ServiceAccount == nil {
		return fmt.Errorf("serviceAccount is nil in nodepool %s", name)
	}
	if spec.ForProvider.Config.ImageType == nil {
		return fmt.Errorf("imageType is nil in nodepool %s", name)
	}
	if spec.ProviderConfigReference == nil {
		return fmt.Errorf("providerConfigReference is nil in nodepool %s", name)
	}
	if spec.ForProvider.Version == nil {
		return fmt.Errorf("version is nil in nodepool %s", name)
	}
	if len(spec.ForProvider.Locations) == 0 {
		return fmt.Errorf("locations is empty in nodepool %s", name)
	}
	return nil
}
