package eks

import (
	"context"
	"fmt"

	"github.com/aws/smithy-go/ptr"
	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"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	dmv1alpha1 "github.com/tidbcloud/eks-provider/api/v1alpha1"
	"github.com/tidbcloud/eks-provider/controllers/dm/task"
	"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"
	"github.com/tidbcloud/infra-provider/util/collections"
)

const (
	// this value is meaningless for now, nodes are created when pod triggers node pool scale out
	// the number of nodes will be the same as the number our pods(1 pod on 1 node)
	// keep this if we need to support share node-pool among jobs later.
	maxNodeCountPerZone = 100

	nodePoolNamePrefix = "dm"
)

// BaseNodePool manage lifecycle node pool for eks.
type BaseNodePool struct {
	*task.BaseTask
	dmSpec     *dmv1alpha1.DMSpec
	nodeLabels map[string]string
	taskName   string
}

func NewBaseNodePool(ctx *task.Context,
	taskName string,
	dmSpec *dmv1alpha1.DMSpec,
	nodeLabels map[string]string,
) *BaseNodePool {
	return &BaseNodePool{
		BaseTask:   &task.BaseTask{Context: ctx},
		dmSpec:     dmSpec,
		nodeLabels: nodeLabels,
		taskName:   taskName,
	}
}

func (t *BaseNodePool) Name() string {
	return t.taskName
}

func (t *BaseNodePool) Create(ctx context.Context) deploy.Result {
	// Gate: ensure AWS node role ARN is present before creating nodepool
	if res := ensureDMNodeRoleARNReady(t.buildNodeRoleARN()); !res.IsDoneState() {
		return res
	}

	zones := t.Zones()
	if len(zones) == 0 {
		return deploy.Error(fmt.Errorf("no zones found"))
	}

	// get the subnets list
	subnets := t.Subnets()
	if len(subnets) == 0 {
		return deploy.Error(fmt.Errorf("no subnets found"))
	}

	// use a specific subnet for the nodegroup, because DMCluster will create a PVC and attach to dm-worker
	// but when scale, the new node may ship to other subnets if we don't specify the subnet, which will cause PVC attach failed
	// and dm-worker can't be scheduled to the new node
	preferSubnetIndex := t.Config.DMConfig.AWS.PreferSubnetIndex
	if preferSubnetIndex >= len(subnets) {
		t.Logger.Info("preferSubnetIndex is out of range, use the default subnet", "preferSubnetIndex", preferSubnetIndex, "defaultSubnet", subnets[0])
		preferSubnetIndex = 0
	}
	selectedSubnets := []string{subnets[preferSubnetIndex]}
	t.Logger.Info("selected subnet for nodegroup", "subnet", selectedSubnets, "preferSubnetIndex", preferSubnetIndex)

	providerConfigName := t.GetProviderConfigName()
	taints := t.getSpecTaints()

	// Create launch template first
	lt := &ec2v1alpha1.LaunchTemplate{
		TypeMeta: metav1.TypeMeta{
			APIVersion: "ec2.aws.crossplane.io/v1alpha1",
			Kind:       "LaunchTemplate",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name: t.getLaunchTemplateName(),
		},
	}

	ltMutateFunc := func() error {
		lt.Labels = t.nodeLabels
		size, err := util.ParseDiskSizeAsGiB(t.dmSpec.SystemVolumeSize)
		if err != nil {
			t.Logger.Error(err, "fail to parse disk size")
			return err
		}

		t.Logger.Info("disk size", "size", size)

		lt.Spec = ec2v1alpha1.LaunchTemplateSpec{
			ForProvider: ec2v1alpha1.LaunchTemplateParameters{
				Region:             t.GetRegion(),
				LaunchTemplateName: ptr.String(t.getLaunchTemplateName()),
				LaunchTemplateData: &ec2v1alpha1.RequestLaunchTemplateData{
					BlockDeviceMappings: nodegroup.GetBlockDeviceMappingsForBottlerocket(size),
					TagSpecifications: []*ec2v1alpha1.LaunchTemplateTagSpecificationRequest{
						{
							ResourceType: ptr.String("instance"),
							Tags: []*ec2v1alpha1.Tag{
								{
									Key:   ptr.String("Name"),
									Value: ptr.String(fmt.Sprintf("%s-%s-%s", nodePoolNamePrefix, t.DMJob.Name, t.taskName)),
								},
							},
						},
					},
				},
			},
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: providerConfigName,
				},
			},
		}
		return nil
	}

	if _, err := controllerutil.CreateOrUpdate(ctx, t.Local, lt, ltMutateFunc); err != nil {
		return deploy.Error(err)
	}

	// Create nodegroup
	np := &manualv1alpha1.NodeGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name: t.getNodePoolName(),
		},
	}
	labels := t.nodeLabels
	mutateFunc := func() error {
		np.Labels = collections.MergeMaps(np.Labels, labels)
		instanceNumber := t.dmSpec.InstanceNumber
		if instanceNumber == 0 {
			instanceNumber = 1
		}

		np.Spec = manualv1alpha1.NodeGroupSpec{
			ForProvider: manualv1alpha1.NodeGroupParameters{
				Region:  t.GetRegion(),
				AMIType: &t.Config.Network.AdminNodesAMIType,
				ScalingConfig: &manualv1alpha1.NodeGroupScalingConfig{
					DesiredSize: ptr.Int32(int32(instanceNumber)),
					MaxSize:     ptr.Int32(int32(maxNodeCountPerZone * len(zones))),
					MinSize:     ptr.Int32(int32(instanceNumber)),
				},
				Subnets:     selectedSubnets,
				NodeRole:    t.buildNodeRoleARN(),
				Labels:      labels,
				Taints:      taints,
				ClusterName: t.Kubernetes.Name,
				UpdateConfig: &manualv1alpha1.NodeGroupUpdateConfig{
					MaxUnavailable: ptr.Int32(1),
				},
				CapacityType:  ptr.String("ON_DEMAND"),
				InstanceTypes: []string{t.dmSpec.InstanceType.Name},
				LaunchTemplate: &manualv1alpha1.LaunchTemplateSpecification{
					Name: ptr.String(t.getLaunchTemplateName()),
				},
			},
			ResourceSpec: xpv1.ResourceSpec{ProviderConfigReference: &xpv1.Reference{Name: providerConfigName}},
		}
		return nil
	}

	if _, err := controllerutil.CreateOrUpdate(ctx, t.Local, np, mutateFunc); err != nil {
		return deploy.Error(err)
	}
	return deploy.Done()
}

func (t *BaseNodePool) PostCreate(ctx context.Context) deploy.Result {
	if res := ensureDMNodeRoleARNReady(t.buildNodeRoleARN()); !res.IsDoneState() {
		return res
	}

	nodePoolName := t.getNodePoolName()
	ng := manualv1alpha1.NodeGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name: nodePoolName,
		},
	}
	if err := t.Local.Get(ctx, client.ObjectKeyFromObject(&ng), &ng); err != nil {
		if apierrors.IsNotFound(err) {
			return deploy.Wait(deploy.DefaultWaitRequeueAfter, "node group %s has not been created", nodePoolName)
		}
		t.Logger.Error(err, "fail to get node group "+nodePoolName)
		return deploy.Error(err)
	}

	if ng.Status.AtProvider.Status == "ACTIVE" {
		return deploy.Done()
	}
	// CREATE_FAILED/ERROR require manual recovery (reboot instance or recreate the nodegroup).
	// Keep retrying so the controller picks up automatically once the nodegroup is recovered.
	if ng.Status.AtProvider.Status == "CREATE_FAILED" || ng.Status.AtProvider.Status == "ERROR" {
		t.Logger.Error(fmt.Errorf("nodegroup requires manual recovery"),
			"nodegroup is in a failed state; please reboot instance or recreate the nodegroup to unblock the job",
			"nodegroup", nodePoolName,
			"status", ng.Status.AtProvider.Status,
			"health", ng.Status.AtProvider.Health)
	}
	return deploy.Wait(deploy.DefaultWaitRequeueAfter, "node group not ready: %s/%s", nodePoolName, ng.Status.AtProvider.Status)
}

func (t *BaseNodePool) Delete(ctx context.Context) deploy.Result {
	// Delete nodegroup first
	nodePool := &manualv1alpha1.NodeGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name: t.getNodePoolName(),
		},
	}
	if err := k8s.DeleteObject(ctx, t.Local, nodePool); err != nil {
		return deploy.Error(err)
	}

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

func (t *BaseNodePool) PostDelete(ctx context.Context) deploy.Result {
	nodePool := &manualv1alpha1.NodeGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name: t.getNodePoolName(),
		},
	}
	if err := k8s.DeleteObjectCompleted(ctx, t.Local, nodePool); err != nil {
		return deploy.Error(err)
	}

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

func (t *BaseNodePool) getNodePoolName() string {
	return getNodePoolName(t.BaseTask, t.dmSpec)
}

func (t *BaseNodePool) getSpecTaints() []manualv1alpha1.Taint {
	var taints []manualv1alpha1.Taint
	taints = append(taints, manualv1alpha1.Taint{
		Key:    ptr.String("dedicated"),
		Value:  ptr.String("tidb-" + task.ComponentName),
		Effect: "NO_SCHEDULE",
	})
	return taints
}

func (t *BaseNodePool) buildNodeRoleARN() string {
	return t.DMJob.Status.AWSNodeRoleARN
}

// ensureDMNodeRoleARNReady checks whether the AWS node role ARN is populated.
// If empty, it returns a retryable Result with a requeue after iamReadyRequeueInterval.
// If non-empty, it returns Done.
func ensureDMNodeRoleARNReady(arn string) deploy.Result {
	if arn == "" {
		return deploy.Wait(iamReadyRequeueInterval, "AWS node role ARN is not available in DMJob status")
	}
	return deploy.Done()
}

func (t *BaseNodePool) getLaunchTemplateName() string {
	return fmt.Sprintf("%s-lt-%s-%s", nodePoolNamePrefix, t.DMJob.Name, t.taskName)
}