package cdccluster

import (
	"context"
	"fmt"
	"math"

	opv1alpha1 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"

	cdcv1 "github.com/tidbcloud/ticdc-service-ng/api/v1"
	"github.com/tidbcloud/ticdc-service-ng/pkg/consts"
	"github.com/tidbcloud/ticdc-service-ng/pkg/k8s"
)

const (
	// reservedMemory for system components and daemon pods
	reservedMemoryInMiB = 1536
)

var toleration = corev1.Toleration{
	Key:      "dedicated",
	Operator: corev1.TolerationOpEqual,
	Value:    "worker",
	Effect:   corev1.TaintEffectNoSchedule,
}

type v2Cluster struct {
	desired          *opv1alpha1.TiCDCGroup
	current          opv1alpha1.TiCDCGroup
	preStopImage     string
	sharedPoolID     string
	schedulePolicies []opv1alpha1.SchedulePolicy
}

func buildSchedulePolicies(zones []string) []opv1alpha1.SchedulePolicy {
	topologies := make([]opv1alpha1.ScheduleTopology, len(zones))
	for i, zone := range zones {
		topologies[i] = opv1alpha1.ScheduleTopology{Topology: map[string]string{consts.ZoneLabelKey: zone}}
	}
	return []opv1alpha1.SchedulePolicy{
		{
			Type: opv1alpha1.SchedulePolicyTypeEvenlySpread,
			EvenlySpread: &opv1alpha1.SchedulePolicyEvenlySpread{
				Topologies: topologies,
			},
		},
	}
}

func NewV2ClusterProvider(preStopImage string, zones []string, tags map[string]string) Provider {
	return &v2Cluster{
		preStopImage:     preStopImage,
		schedulePolicies: buildSchedulePolicies(zones),
		sharedPoolID:     tags[consts.SharedPoolLabelKey],
	}
}

func (p *v2Cluster) Init(namespace, name string) {
	p.desired = &opv1alpha1.TiCDCGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
		Spec: opv1alpha1.TiCDCGroupSpec{
			SchedulePolicies: p.schedulePolicies,
			Template: opv1alpha1.TiCDCTemplate{
				Spec: opv1alpha1.TiCDCTemplateSpec{
					Overlay: &opv1alpha1.Overlay{
						Pod: &opv1alpha1.PodOverlay{
							ObjectMeta: opv1alpha1.ObjectMeta{
								Labels: map[string]string{
									"app.kubernetes.io/component":  "ticdc",
									"app.kubernetes.io/instance":   name,
									"app.kubernetes.io/managed-by": "tidb-operator",
									"app.kubernetes.io/name":       "ticdc-cluster",
									// Mandatory labels:
									"tags.tidbcloud.com/resource-pool": p.sharedPoolID,
									"tags.tidbcloud.com/component":     "ticdc",
								},
								Annotations: map[string]string{
									"pingcap.com/ignore-diff": "spec.dnsPolicy",
								},
							},
							Spec: &corev1.PodSpec{
								Affinity:        affinity(cloudv1.ComponentTiCDC),
								HostNetwork:     true,
								DNSPolicy:       corev1.DNSClusterFirstWithHostNet,
								SecurityContext: &corev1.PodSecurityContext{},
								Tolerations:     []corev1.Toleration{toleration},
							},
						},
					},
					PreStop: &opv1alpha1.TiCDCPreStop{
						Image: ptr.To(p.preStopImage),
					},
				},
			},
		},
	}
}

func (p *v2Cluster) SetReplicas(replicas int32) {
	p.desired.Spec.Replicas = ptr.To(replicas)
}

func (p *v2Cluster) SetVersion(baseImage string, version string, tag string) {
	image := baseImage
	if tag != "" {
		image = fmt.Sprintf("%s:%s", baseImage, tag)
	}
	p.desired.Spec.Template.Spec.Image = ptr.To(image)
	p.desired.Spec.Template.Spec.Version = version
}

func (p *v2Cluster) SetUpstreamCluster(_ string, name string) {
	// NOTE: TiDB Operator v2 does not support cross namespace yet
	p.desired.Spec.Cluster.Name = name
}

func (p *v2Cluster) SetTomlConfig(conf string) error {
	p.desired.Spec.Template.Spec.Config = opv1alpha1.ConfigFile(conf)
	return nil
}

func (p *v2Cluster) SetConfigField(key string, value interface{}) error {
	cdcConf := NewCDCConfig()
	if err := cdcConf.UnmarshalTOML([]byte(p.desired.Spec.Template.Spec.Config)); err != nil {
		return err
	}
	cdcConf.Set(key, value)
	b, err := cdcConf.MarshalTOML()
	if err != nil {
		return err
	}
	p.desired.Spec.Template.Spec.Config = opv1alpha1.ConfigFile(b)
	return nil
}

func (p *v2Cluster) SetRequest(cpuCore, memoryGi uint32) {
	memory := resource.MustParse(fmt.Sprintf("%dMi", ComputeAvailableMemory(memoryGi)))
	cpu := resource.MustParse(fmt.Sprintf("%d", cpuCore-1))
	p.desired.Spec.Template.Spec.Resources = opv1alpha1.ResourceRequirements{
		CPU:    &cpu,
		Memory: &memory,
	}
}

func (p *v2Cluster) SetNodeSelector(labelSelector map[string]string) {
	// Set node selector for TiCDC
	p.desired.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector = labelSelector
}

func (p *v2Cluster) SetStorage(size, vac string) {
	var vacName *string
	if vac != "" {
		vacName = ptr.To(vac)
	}
	p.desired.Spec.Template.Spec.Overlay.Pod.Spec.Volumes = []corev1.Volume{
		{
			Name: TiCDCDefaultDirName,
			VolumeSource: corev1.VolumeSource{
				Ephemeral: &corev1.EphemeralVolumeSource{
					VolumeClaimTemplate: &corev1.PersistentVolumeClaimTemplate{
						ObjectMeta: metav1.ObjectMeta{
							Labels: map[string]string{
								// Mandatory labels:
								"tags.tidbcloud.com/resource-pool": p.sharedPoolID,
								"tags.tidbcloud.com/component":     "ticdc",
							},
						},
						Spec: corev1.PersistentVolumeClaimSpec{
							AccessModes:      []corev1.PersistentVolumeAccessMode{corev1.ReadWriteOnce},
							StorageClassName: ptr.To("cloud-ssd"), // for nextgen we always use cloud-ssd
							Resources: corev1.VolumeResourceRequirements{
								Requests: map[corev1.ResourceName]resource.Quantity{
									corev1.ResourceStorage: resource.MustParse(size),
								},
							},
							VolumeAttributesClassName: vacName,
						},
					},
				},
			},
		},
	}
	p.desired.Spec.Template.Spec.Overlay.Pod.Spec.Containers = []corev1.Container{
		{
			Name: opv1alpha1.ContainerNameTiCDC,
			VolumeMounts: []corev1.VolumeMount{
				{
					Name:      TiCDCDefaultDirName,
					MountPath: TiCDCDefaultDirPath,
				},
			},
		},
	}
	// priority classs is created by cluster-service-ng, hardcode it here,
	p.desired.Spec.Template.Spec.Overlay.Pod.Spec.PriorityClassName = "tidb-ng-component-critical"
}

func (p *v2Cluster) GetClusterStatus(ctx context.Context, cli client.Client) (*cdcv1.TiCDCStatus, error) {
	if err := cli.Get(ctx, client.ObjectKeyFromObject(p.desired), &p.current); err != nil {
		return nil, err
	}
	phase, reason := convertPhaseReasonV2(&p.current)
	nodeSpec := p.current.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector[consts.NodeSpecLabelKey]
	cpuCore, memoryGi, err := ParseNodeSpecLabelValue(nodeSpec)
	if err != nil {
		return nil, err
	}
	return &cdcv1.TiCDCStatus{
		Phase:  phase,
		Reason: reason,
		NodeInfo: &cdcv1.NodeInfo{
			Replicas: int(*p.current.Spec.Replicas),
			NodeSpec: cdcv1.NodeSpec{
				CpuCore:  cpuCore,
				MemoryGi: memoryGi,
			},
		},
	}, nil
}

func (p *v2Cluster) NeedReschedule(ctx context.Context, cli client.Client, desiredNode *Node) (bool, error) {
	err := cli.Get(ctx, client.ObjectKeyFromObject(p.desired), &p.current)
	if err != nil {
		if apierrors.IsNotFound(err) {
			return false, nil
		}
		return false, err
	}

	nodeSelector := p.current.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector
	instanceType, ok := nodeSelector[consts.InstanceTypeLabelKey]
	if !ok || instanceType != desiredNode.InstanceType {
		return true, nil
	}
	k8sVersion, ok := nodeSelector[consts.K8sVersionLabelKey]
	if !ok || k8sVersion != desiredNode.K8sVersion {
		return true, nil
	}
	tolerations := p.current.Spec.Template.Spec.Overlay.Pod.Spec.Tolerations
	if !containTaintValue(tolerations, desiredNode.TaintValue) {
		return true, nil
	}
	return false, nil
}

func (p *v2Cluster) Rescheduled(ctx context.Context, cli client.Client) (bool, error) {
	podSelector := client.MatchingLabels{
		"pingcap.com/group":     p.desired.Name,
		"pingcap.com/component": string(cloudv1.ComponentTiCDC),
	}
	pods := &corev1.PodList{}
	// list pods for the TiCDCGroup
	if err := cli.List(ctx, pods, client.InNamespace(p.desired.Namespace), podSelector); err != nil {
		return false, err
	}
	desiredSelector := p.desired.Spec.Template.Spec.Overlay.Pod.Spec.NodeSelector
	desiredTolerations := p.desired.Spec.Template.Spec.Overlay.Pod.Spec.Tolerations
	scheduled := 0
	for _, pod := range pods.Items {
		if pod.Spec.NodeName == "" { // pod not scheduled
			continue
		}
		selector := pod.Spec.NodeSelector
		if selector[consts.InstanceTypeLabelKey] != desiredSelector[consts.InstanceTypeLabelKey] {
			continue
		}
		if selector[consts.K8sVersionLabelKey] != desiredSelector[consts.K8sVersionLabelKey] {
			continue
		}
		tolerationsMatched := true
		for _, desiredToleration := range desiredTolerations {
			if !containTaintValue(pod.Spec.Tolerations, desiredToleration.Value) {
				tolerationsMatched = false
				break
			}
		}
		if !tolerationsMatched {
			continue
		}
		scheduled++
	}

	return scheduled == int(*p.desired.Spec.Replicas), nil
}

func (p *v2Cluster) ApplyCR(ctx context.Context, cli client.Client) error {
	if err := k8s.Apply(ctx, cli, p.desired); err != nil {
		return fmt.Errorf("failed to apply TiCDCGroup CR: %w", err)
	}

	return nil
}

func (p *v2Cluster) DeleteCR(ctx context.Context, cli client.Client, name types.NamespacedName) error {
	return k8s.DeleteObject(ctx, cli, &opv1alpha1.TiCDCGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name.Name,
			Namespace: name.Namespace,
		},
	})
}

func convertPhaseReasonV2(cdcGroup *opv1alpha1.TiCDCGroup) (cdcv1.TiCDCPhase, string) {
	suspended := meta.IsStatusConditionTrue(cdcGroup.Status.Conditions, opv1alpha1.CondSuspended)
	ready := meta.IsStatusConditionTrue(cdcGroup.Status.Conditions, opv1alpha1.CondReady)
	synced := cdcGroup.Generation == cdcGroup.Status.ObservedGeneration &&
		meta.IsStatusConditionTrue(cdcGroup.Status.Conditions, opv1alpha1.CondSynced)
	switch {
	case suspended:
		return cdcv1.TiCDCPhaseSuspend, "Suspend"
	case synced && ready:
		return cdcv1.TiCDCPhaseRunning, "Running"
	case !synced:
		if cdcGroup.Status.Replicas != *cdcGroup.Spec.Replicas {
			return cdcv1.TiCDCPhaseScaling, "Scaling"
		}
		return cdcv1.TiCDCPhaseUpgrading, "Upgrading"
	default:
		// if some instances are ready, should we mark it as running?
		return cdcv1.TiCDCPhaseUnavailable, "Unavailable"
	}
}

func (p *v2Cluster) SetRestart(restartedAt string) {
	if len(p.desired.Spec.Template.Annotations) == 0 {
		p.desired.Spec.Template.Annotations = make(map[string]string)
	}
	p.desired.Spec.Template.Annotations[consts.RestartTiCDCOperatorV2Annotation] = restartedAt
}

func containTaintValue(tolerations []corev1.Toleration, taintValue string) bool {
	for _, toleration := range tolerations {
		if toleration.Value == taintValue {
			return true
		}
	}
	return false
}

// affinity returns a PodAntiAffinity configuration for the given component type.
// It ensures that pods of the same component type are not scheduled on the same node.
func affinity(comp cloudv1.ComponentType) *corev1.Affinity {
	return &corev1.Affinity{
		PodAntiAffinity: &corev1.PodAntiAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{
				{
					LabelSelector: &metav1.LabelSelector{
						MatchLabels: map[string]string{
							consts.ComponentLabelKey: string(comp),
						},
					},
					TopologyKey: "kubernetes.io/hostname",
				},
			},
		},
	}
}

// GenNodeSpecLabelValue generates a formatted string label from the given
// CPU core count and memory size in GiB.
// For example, if cpuCore is 8 and memoryGi is 16, it returns "8C16G".
func GenNodeSpecLabelValue(cpuCore, memoryGi uint32) string {
	return fmt.Sprintf("%dC%dG", cpuCore, memoryGi)
}

// ParseNodeSpecLabelValue is the reverse of GenNodeSpecLabelValue.
// It parses a formatted string label and extracts the CPU core count and memory size.
// It returns the parsed cpuCore, memoryGi, and an error if the input format is invalid.
// For example, "8C16G" will be parsed into cpuCore=8, memoryGi=16, err=nil.
func ParseNodeSpecLabelValue(labelValue string) (cpuCore uint32, memoryGi uint32, err error) {
	_, err = fmt.Sscanf(labelValue, "%dC%dG", &cpuCore, &memoryGi)
	if err != nil {
		return 0, 0, fmt.Errorf("invalid node spec format: %q. expected format like '8C16G'", labelValue)
	}
	return cpuCore, memoryGi, nil
}

// ref: https://cloud.google.com/kubernetes-engine/docs/concepts/cluster-architecture#node_allocatable
// and looks like Azure and EKS share the same algorithm https://github.com/awslabs/amazon-eks-ami/issues/387#issuecomment-582419188
// https://docs.microsoft.com/en-us/azure/aks/concepts-clusters-workloads#resource-reservations
var memTaxTable = []struct {
	boundary       uint64
	rate           float64
	quickDeduction uint64
}{
	{
		boundary:       1 * 1024, // 1GB
		rate:           0.0,
		quickDeduction: 255,
	},
	{
		boundary:       4 * 1024, // 4GB
		rate:           0.25,
		quickDeduction: 0,
	},
	{
		boundary:       8 * 1024, // 8GB
		rate:           0.2,
		quickDeduction: 205, // 0.05 * 4GB rounded
	},
	{
		boundary:       16 * 1024,
		rate:           0.1,
		quickDeduction: 1, // 0.1 * 4GB + 0.15 * 4GB
	},
	{
		boundary:       128 * 1024, // 128GB
		rate:           0.06,
		quickDeduction: 1680, // 0.19 * 4GB + 0.14 * 4GB + 0.04 * 8GB rounded
	},
	{
		boundary:       math.MaxInt64,
		rate:           0.02,
		quickDeduction: 6922, // (0.25 - 0.02) * 4 + (0.2 - 0.02) * 4 + (0.1 - 0.02) * 8 + (0.06 - 0.02) * 112 rounded
	},
}

// ComputeAvailableMemory computes the actual memory available to pods after reserving memory for system components and daemon pods.
func ComputeAvailableMemory(memoryInGiB uint32) uint64 {
	memoryInMiB := uint64(memoryInGiB) * 1024
	for _, level := range memTaxTable {
		// ASC order of memTaxTable is guaranteed
		if memoryInMiB <= level.boundary {
			memoryInMiB = memoryInMiB - uint64(math.Ceil(float64(memoryInMiB)*level.rate)) - level.quickDeduction
			break
		}
	}
	// subtract the reserved memory for daemon pods
	memoryInMiB -= reservedMemoryInMiB
	if memoryInMiB <= 0 {
		memoryInMiB = 0
	}
	return memoryInMiB
}