package eks

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

	ec2v1alpha1 "github.com/crossplane-contrib/provider-aws/apis/ec2/v1alpha1"
	"github.com/crossplane-contrib/provider-aws/apis/eks/manualv1alpha1"
	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	"github.com/tidbcloud/eks-provider/controllers/replay/task"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/controllers/common"
	"github.com/tidbcloud/infra-provider/pkg/deploy"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/pkg/nodegroup"
	"github.com/tidbcloud/infra-provider/pkg/util"
)

type NodePool struct {
	*task.BaseTask
}

func NewNodePool(ctx *task.Context) *NodePool {
	return &NodePool{
		BaseTask: &task.BaseTask{Context: ctx},
	}
}

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

func (t *NodePool) Create(ctx context.Context) deploy.Result {
	// create launch template
	lt := &ec2v1alpha1.LaunchTemplate{
		TypeMeta: metav1.TypeMeta{
			APIVersion: "ec2.aws.crossplane.io/v1alpha1",
			Kind:       "LaunchTemplate",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name: t.getLaunchTemplateName(),
		},
	}

	labels := t.getNodeLabels()
	ltMutateFunc := func() error {
		lt.Labels = labels
		size := int64(20)
		region := t.GetRegion()

		// Encode UserData to base64
		userDataBase64 := base64.StdEncoding.EncodeToString([]byte(task.LaunchTemplateUserData))

		lt.Spec = ec2v1alpha1.LaunchTemplateSpec{
			ForProvider: ec2v1alpha1.LaunchTemplateParameters{
				Region:             region,
				LaunchTemplateName: ptr.To(t.getLaunchTemplateName()),
				LaunchTemplateData: &ec2v1alpha1.RequestLaunchTemplateData{
					BlockDeviceMappings: nodegroup.GetBlockDeviceMappingsForBottlerocket(size),
					TagSpecifications: []*ec2v1alpha1.LaunchTemplateTagSpecificationRequest{
						{
							ResourceType: ptr.To("instance"),
							Tags: []*ec2v1alpha1.Tag{
								{
									Key:   ptr.To("Name"),
									Value: ptr.To(fmt.Sprintf("replay-%s", t.Context.Replay.Name)),
								},
							},
						},
					},
					UserData: ptr.To(userDataBase64),
				},
			},
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: t.GetEKSProviderConfigName(),
				},
			},
		}
		return nil
	}

	if _, err := controllerutil.CreateOrUpdate(ctx, t.Context.Local, lt, ltMutateFunc); err != nil {
		t.Logger.Error(err, "fail to create launch template "+t.getLaunchTemplateName())
		return deploy.Error(err)
	}

	// Get all unique instance types from tasks
	instanceTypes := t.getUniqueInstanceTypes()
	t.Logger.Info("Creating node groups for instance types", "instanceTypes", instanceTypes)

	// Create node group for each instance type
	for _, instanceType := range instanceTypes {
		nodePoolName := t.generateNodePoolNameForInstanceType(instanceType)
		nodePool, err := t.desiredNodePoolForInstanceType(instanceType)
		if err != nil {
			t.Logger.Error(err, "fail to get node group "+nodePoolName, "instanceType", instanceType)
			return deploy.Error(err)
		}

		// Check if NodeGroup already exists
		existingNG := &manualv1alpha1.NodeGroup{}
		err = t.Context.Local.Get(ctx, client.ObjectKey{Name: nodePoolName}, existingNG)
		if err == nil {
			t.Logger.Info("NodeGroup already exists, skipping update", "nodegroup", nodePoolName, "instanceType", instanceType)
			continue
		}
		if !apierrors.IsNotFound(err) {
			t.Logger.Error(err, "fail to check node group existence "+nodePoolName, "instanceType", instanceType)
			return deploy.Error(err)
		}

		// NodeGroup doesn't exist, create it
		if err := k8s.Apply(ctx, t.Context.Local, nodePool); err != nil {
			t.Logger.Error(err, "fail to create node group "+nodePoolName, "instanceType", instanceType)
			return deploy.Error(err)
		}
		t.Logger.Info("Created node group", "nodegroup", nodePoolName, "instanceType", instanceType)
	}

	return deploy.Done()
}

func (t *NodePool) PostCreate(ctx context.Context) deploy.Result {
	// Get all unique instance types from tasks
	instanceTypes := t.getUniqueInstanceTypes()

	allReady := true
	var notReadyNodeGroups []string

	// Check status for each node group
	for _, instanceType := range instanceTypes {
		nodePoolName := t.generateNodePoolNameForInstanceType(instanceType)
		ng := manualv1alpha1.NodeGroup{
			ObjectMeta: metav1.ObjectMeta{
				Name: nodePoolName,
			},
		}

		if err := t.Context.Local.Get(ctx, client.ObjectKeyFromObject(&ng), &ng); err == nil {
			// Check if nodegroup is in a failed state
			if ng.Status.AtProvider.Status == "CREATE_FAILED" || ng.Status.AtProvider.Status == "ERROR" {
				t.Context.Logger.Error(fmt.Errorf("nodegroup creation failed"), "nodegroup creation failed",
					"nodegroup", nodePoolName,
					"instanceType", instanceType,
					"status", ng.Status.AtProvider.Status,
					"health", ng.Status.AtProvider.Health)
			}
			if ng.Status.AtProvider.Status != "ACTIVE" {
				allReady = false
				notReadyNodeGroups = append(notReadyNodeGroups, fmt.Sprintf("%s(%s)", nodePoolName, ng.Status.AtProvider.Status))
			}
		} else {
			t.Context.Logger.Error(err, "fail to get node group "+nodePoolName, "instanceType", instanceType)
			allReady = false
			notReadyNodeGroups = append(notReadyNodeGroups, nodePoolName)
		}
	}

	if allReady {
		return deploy.Done()
	} else {
		return deploy.Error(fmt.Errorf("node groups not ready: %s", strings.Join(notReadyNodeGroups, ", "))).WithRequeuAfter(time.Second * 20)
	}
}

func (t *NodePool) Delete(ctx context.Context) deploy.Result {
	// Get all unique instance types from tasks
	instanceTypes := t.getUniqueInstanceTypes()

	// Delete node groups for each instance type
	for _, instanceType := range instanceTypes {
		nodePoolName := t.generateNodePoolNameForInstanceType(instanceType)
		nodePool := &manualv1alpha1.NodeGroup{
			TypeMeta: metav1.TypeMeta{
				APIVersion: "eks.aws.crossplane.io/v1alpha1",
				Kind:       "NodeGroup",
			},
			ObjectMeta: metav1.ObjectMeta{
				Name: nodePoolName,
			},
		}
		if err := k8s.DeleteObject(ctx, t.Context.Local, nodePool); err != nil {
			t.Logger.Error(err, "fail to delete node group "+nodePoolName, "instanceType", instanceType)
			return deploy.Error(err)
		}
		t.Logger.Info("Deleted node group", "nodegroup", nodePoolName, "instanceType", instanceType)
	}

	// delete launch template
	lt := &ec2v1alpha1.LaunchTemplate{
		ObjectMeta: metav1.ObjectMeta{
			Name: t.getLaunchTemplateName(),
		},
	}
	return deploy.ErrorOrDone(k8s.DeleteObject(ctx, t.Context.Local, lt))
}

func (t *NodePool) PostDelete(ctx context.Context) deploy.Result {
	// Get all unique instance types from tasks
	instanceTypes := t.getUniqueInstanceTypes()

	// Check deletion status for each node group
	for _, instanceType := range instanceTypes {
		nodePoolName := t.generateNodePoolNameForInstanceType(instanceType)
		nodePool := &manualv1alpha1.NodeGroup{
			TypeMeta: metav1.TypeMeta{
				APIVersion: "eks.aws.crossplane.io/v1alpha1",
				Kind:       "NodeGroup",
			},
			ObjectMeta: metav1.ObjectMeta{
				Name: nodePoolName,
			},
		}
		if err := k8s.DeleteObjectCompleted(ctx, t.Context.Local, nodePool); err != nil {
			t.Logger.Info("Waiting for node group deletion", "nodegroup", nodePoolName, "instanceType", instanceType)
			return deploy.Error(err).WithRequeuAfter(time.Second * 20)
		}
	}

	lt := &ec2v1alpha1.LaunchTemplate{
		ObjectMeta: metav1.ObjectMeta{
			Name: t.getLaunchTemplateName(),
		},
	}
	return deploy.ErrorOrDone(k8s.DeleteObjectCompleted(ctx, t.Context.Local, lt))
}

func (t *NodePool) generateNodePoolNameForInstanceType(instanceType string) string {
	// Generate a unique NodeGroup name based on region, replay name, and instance type
	region := t.GetRegion()

	// Sanitize instance type for use in name (replace dots with dashes)
	sanitizedInstanceType := strings.ReplaceAll(instanceType, ".", "-")

	identifier := fmt.Sprintf("%s-%s-%s", region, t.Replay.Name, sanitizedInstanceType)
	hash := md5.Sum([]byte(identifier))
	hashString := hex.EncodeToString(hash[:])

	// Final format: replay-ng-{instance-type}-{hash} (e.g., "replay-ng-c6g-8xlarge-a28409aa123")
	// Keep it under 63 chars
	baseName := fmt.Sprintf("replay-ng-%s-%s", sanitizedInstanceType, hashString[:8])
	if len(baseName) > 63 {
		// If still too long, use shorter hash
		baseName = fmt.Sprintf("replay-ng-%s", hashString[:11])
	}
	return baseName
}

// getUniqueInstanceTypes returns a list of unique instance types from all tasks
func (t *NodePool) getUniqueInstanceTypes() []string {
	instanceTypeSet := make(map[string]struct{})

	for _, replayTask := range t.Replay.Spec.Tasks {
		instanceType := task.DefaultInstanceType
		if replayTask.Resources != nil && replayTask.Resources.InstanceType != "" {
			instanceType = replayTask.Resources.InstanceType
		}
		instanceTypeSet[instanceType] = struct{}{}
	}

	// Convert set to slice
	instanceTypes := make([]string, 0, len(instanceTypeSet))
	for instanceType := range instanceTypeSet {
		instanceTypes = append(instanceTypes, instanceType)
	}

	// If no tasks or no instance types found, use default
	if len(instanceTypes) == 0 {
		instanceTypes = append(instanceTypes, task.DefaultInstanceType)
	}

	return instanceTypes
}

func (t *NodePool) getLaunchTemplateName() string {
	return fmt.Sprintf("replay-lt-%s", t.Context.Replay.Name)
}

func (t *NodePool) desiredNodePoolForInstanceType(instanceType string) (*manualv1alpha1.NodeGroup, error) {
	labels := t.getNodeLabelsForInstanceType(instanceType)
	region := t.GetRegion()

	// Use the Node Role ARN created by iam_role task
	nodeRole := t.Replay.Status.AWSNodeRoleARN
	if nodeRole == "" {
		return nil, fmt.Errorf("AWSNodeRoleARN is not set in Replay status, iam_role task may not be ready")
	}

	nodePool := &manualv1alpha1.NodeGroup{
		TypeMeta: metav1.TypeMeta{
			APIVersion: "eks.aws.crossplane.io/v1alpha1",
			Kind:       "NodeGroup",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:   t.generateNodePoolNameForInstanceType(instanceType),
			Labels: labels,
		},
		Spec: manualv1alpha1.NodeGroupSpec{
			ForProvider: manualv1alpha1.NodeGroupParameters{
				Region: region,
				// Use ARM AMI for replay workloads (BOTTLEROCKET_ARM_64)
				// This is required for ARM-based instances like c6g.8xlarge
				AMIType:     ptr.To(nodegroup.ARMAMIType),
				ClusterName: t.Context.Replay.Spec.Kubernetes.Name,
				NodeRole:    nodeRole,
				ScalingConfig: &manualv1alpha1.NodeGroupScalingConfig{
					MaxSize: ptr.To[int32](100),
					MinSize: ptr.To[int32](0),
				},
				Subnets:       t.Subnets,
				InstanceTypes: []string{instanceType},
				Labels:        labels,
				Tags:          t.getNodeTags(),
				Taints: []manualv1alpha1.Taint{
					{
						Key:    ptr.To("dedicated"),
						Value:  ptr.To(fmt.Sprintf("tidb-%s", task.ReplayComponent)),
						Effect: util.ToCrossplaneTaintEffect(corev1.TaintEffectNoSchedule),
					},
				},
				CapacityType: ptr.To("ON_DEMAND"),
				LaunchTemplate: &manualv1alpha1.LaunchTemplateSpecification{
					Name: ptr.To(t.getLaunchTemplateName()),
				},
				UpdateConfig: &manualv1alpha1.NodeGroupUpdateConfig{
					MaxUnavailable: ptr.To[int32](1),
				},
			},
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: t.GetEKSProviderConfigName(),
				},
				// Set managementPolicies to avoid unnecessary updates
				// "ObserveOnly" would prevent all updates, but we still want to allow scaling
				//ManagementPolicies: []xpv1.ManagementAction{
				//	xpv1.ManagementActionCreate,
				//	xpv1.ManagementActionObserve,
				//	xpv1.ManagementActionUpdate,
				//	xpv1.ManagementActionDelete,
				//},
				DeletionPolicy: xpv1.DeletionDelete,
			},
		},
	}

	t.NodeLabels = labels

	return nodePool, nil
}

func (t *NodePool) getNodeLabels() map[string]string {
	tags := cloudv1.TagsFromLabels(t.Replay.Labels)
	// tags := map[string]string{}
	tags = common.ConcatenateTags(tags, map[string]string{
		// common.UsedBy:        common.CustomerTiDB,
		common.Component:     task.ReplayComponent,
		task.NodeSelectorKey: t.Replay.Name, // Required for pod node affinity matching
	})
	return tags
}

// getNodeLabelsForInstanceType returns node labels with instanceType information
func (t *NodePool) getNodeLabelsForInstanceType(instanceType string) map[string]string {
	tags := t.getNodeLabels()
	// Add instance type label for pod scheduling
	tags[task.NodeInstanceTypeKey] = instanceType
	return tags
}

func (t *NodePool) getNodeTags() map[string]string {
	tags := cloudv1.TagsFromLabels(t.Replay.Labels)
	tags = common.ConcatenateTags(tags, map[string]string{
		// common.UsedBy:        common.CustomerTiDB,
		common.Component:   task.ReplayComponent,
		common.Tenant:      t.Replay.GetTenantID(),
		common.Project:     t.Replay.GetProjectID(),
		common.ServiceType: common.Dedicated,
	})
	return tags
}