// Copyright 2026 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//go:build e2e

// Package data constructs the Kubernetes custom resources used by e2e suites.
package data

import (
	tioperatorv1alpha1 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"

	autoscalingv1alpha1 "github.com/tidbcloud/cluster-service-ng/internal/capacity-planner/api/v1alpha1"
)

const (
	ClusterName          = "e2e"
	TiDBGroupName        = "tidb"
	TiDBGroupKind        = "TiDBGroup"
	SmallTierCPU         = "100m"
	MediumTierCPU        = "300m"
	LargeTierCPU         = "600m"
	TierLabelKey         = "e2e.autoscaling.pingcap.com/tier"
	PreserveLabelKey     = "e2e.autoscaling.pingcap.com/preserve"
	TargetCPUUtilization = int32(65)

	tiDBVersion   = "v8.5.2"
	imageRegistry = "gcr.io/pingcap-public/dbaas/"
)

// NewCluster constructs the Cluster shared by the TiDB suite.
func NewCluster(namespace string) *tioperatorv1alpha1.Cluster {
	return &tioperatorv1alpha1.Cluster{
		TypeMeta: metav1.TypeMeta{
			APIVersion: tioperatorv1alpha1.GroupVersion.String(),
			Kind:       "Cluster",
		},
		ObjectMeta: metav1.ObjectMeta{Name: ClusterName, Namespace: namespace},
		Spec: tioperatorv1alpha1.ClusterSpec{
			UpgradePolicy: tioperatorv1alpha1.UpgradePolicyDefault,
		},
	}
}

// NewPDGroup constructs the single-replica PDGroup used by the TiDB suite.
func NewPDGroup(namespace string) *tioperatorv1alpha1.PDGroup {
	return &tioperatorv1alpha1.PDGroup{
		TypeMeta: metav1.TypeMeta{
			APIVersion: tioperatorv1alpha1.GroupVersion.String(),
			Kind:       "PDGroup",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      "pd",
			Namespace: namespace,
			Labels:    groupLabels("pd", tioperatorv1alpha1.LabelValComponentPD),
		},
		Spec: tioperatorv1alpha1.PDGroupSpec{
			Cluster:  tioperatorv1alpha1.ClusterReference{Name: ClusterName},
			Replicas: ptr.To[int32](1),
			Template: tioperatorv1alpha1.PDTemplate{
				Spec: tioperatorv1alpha1.PDTemplateSpec{
					Version: tiDBVersion,
					Image:   ptr.To(imageRegistry + "pd"),
					Volumes: []tioperatorv1alpha1.Volume{{
						Name:    "data",
						Mounts:  []tioperatorv1alpha1.VolumeMount{{Type: tioperatorv1alpha1.VolumeMountTypePDData}},
						Storage: resource.MustParse("1Gi"),
					}},
				},
			},
		},
	}
}

// NewTiKVGroup constructs the single-replica TiKVGroup used by the TiDB suite.
func NewTiKVGroup(namespace string) *tioperatorv1alpha1.TiKVGroup {
	return &tioperatorv1alpha1.TiKVGroup{
		TypeMeta: metav1.TypeMeta{
			APIVersion: tioperatorv1alpha1.GroupVersion.String(),
			Kind:       "TiKVGroup",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      "tikv",
			Namespace: namespace,
			Labels:    groupLabels("tikv", tioperatorv1alpha1.LabelValComponentTiKV),
		},
		Spec: tioperatorv1alpha1.TiKVGroupSpec{
			Cluster:  tioperatorv1alpha1.ClusterReference{Name: ClusterName},
			Replicas: ptr.To[int32](1),
			Template: tioperatorv1alpha1.TiKVTemplate{
				Spec: tioperatorv1alpha1.TiKVTemplateSpec{
					Version: tiDBVersion,
					Image:   ptr.To(imageRegistry + "tikv"),
					Volumes: []tioperatorv1alpha1.Volume{{
						Name:    "data",
						Mounts:  []tioperatorv1alpha1.VolumeMount{{Type: tioperatorv1alpha1.VolumeMountTypeTiKVData}},
						Storage: resource.MustParse("1Gi"),
					}},
				},
			},
		},
	}
}

// NewTiDBGroup constructs the autoscaling target for the TiDB suite.
func NewTiDBGroup(namespace string) *tioperatorv1alpha1.TiDBGroup {
	labels := groupLabels(TiDBGroupName, tioperatorv1alpha1.LabelValComponentTiDB)
	labels[PreserveLabelKey] = "true"
	return &tioperatorv1alpha1.TiDBGroup{
		TypeMeta: metav1.TypeMeta{
			APIVersion: tioperatorv1alpha1.GroupVersion.String(),
			Kind:       TiDBGroupKind,
		},
		ObjectMeta: metav1.ObjectMeta{Name: TiDBGroupName, Namespace: namespace, Labels: labels},
		Spec: tioperatorv1alpha1.TiDBGroupSpec{
			Cluster:  tioperatorv1alpha1.ClusterReference{Name: ClusterName},
			Replicas: ptr.To[int32](1),
			Template: tioperatorv1alpha1.TiDBTemplate{
				Spec: tioperatorv1alpha1.TiDBTemplateSpec{
					Version: tiDBVersion,
					Image:   ptr.To(imageRegistry + "tidb"),
					Resources: tioperatorv1alpha1.ResourceRequirements{
						CPU: ptr.To(resource.MustParse(SmallTierCPU)),
					},
					Probes: tioperatorv1alpha1.TiDBProbes{
						Readiness: &tioperatorv1alpha1.TiDBProb{Type: ptr.To(tioperatorv1alpha1.CommandProbeType)},
					},
				},
			},
		},
	}
}

// NewAutoscaler constructs a disabled Autoscaler for the TiDB target.
func NewAutoscaler(namespace, targetName string) *autoscalingv1alpha1.Autoscaler {
	averageUtilization := TargetCPUUtilization
	return &autoscalingv1alpha1.Autoscaler{
		TypeMeta: metav1.TypeMeta{
			APIVersion: autoscalingv1alpha1.GroupVersion.String(),
			Kind:       "Autoscaler",
		},
		ObjectMeta: metav1.ObjectMeta{Name: TiDBGroupName, Namespace: namespace},
		Spec: autoscalingv1alpha1.AutoscalerSpec{
			TargetRef: autoscalingv1alpha1.CrossVersionObjectReference{
				APIVersion: tioperatorv1alpha1.GroupVersion.String(),
				Kind:       TiDBGroupKind,
				Name:       targetName,
			},
			Mode:     autoscalingv1alpha1.AutoscalerModeDisabled,
			Replicas: &autoscalingv1alpha1.ReplicasSpec{Min: 1, Max: 10},
			Tiers: autoscalingv1alpha1.Tiers{
				Template: autoscalingv1alpha1.TierTemplateRef{Name: "e2e-tidb-node"},
				Sets: []autoscalingv1alpha1.TierSet{{
					Dimensions: []autoscalingv1alpha1.TierDimension{{
						Resources: []autoscalingv1alpha1.TierDimensionResource{{
							Name: corev1.ResourceCPU,
							Variants: []resource.Quantity{
								resource.MustParse(SmallTierCPU),
								resource.MustParse(MediumTierCPU),
								resource.MustParse(LargeTierCPU),
							},
						}},
					}},
				}},
			},
			Policy: autoscalingv1alpha1.AutoscalingPolicyCostOptimized,
			Metrics: []autoscalingv1alpha1.MetricSpec{{
				Type: autoscalingv1alpha1.MetricSourceTypeResource,
				Resource: &autoscalingv1alpha1.ResourceMetricSource{
					Name: corev1.ResourceCPU,
					Target: autoscalingv1alpha1.MetricTarget{
						Type:               autoscalingv1alpha1.MetricTargetTypeUtilization,
						AverageUtilization: &averageUtilization,
					},
				},
			}},
		},
	}
}

func groupLabels(name, component string) map[string]string {
	return map[string]string{
		tioperatorv1alpha1.LabelKeyCluster:   ClusterName,
		tioperatorv1alpha1.LabelKeyComponent: component,
		tioperatorv1alpha1.LabelKeyGroup:     name,
	}
}