package builder

import (
	"fmt"
	"path/filepath"
	"strconv"
	"strings"

	opv2 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/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"

	"github.com/tidbcloud/eks-provider/controllers/cluster/task/common/audit"
	"github.com/tidbcloud/eks-provider/controllers/cluster/task/common/feature"
	"github.com/tidbcloud/eks-provider/controllers/utils"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/util"
)

const (
	AddIndexAccelerateVolumeName                        = "index-acceleration"
	AddIndexAccelerateStorageClassName                  = "index-acceleration"
	AddIndexAccelerateHyperdiskBalancedStorageClassName = "index-acceleration-hyperdisk-balanced"
	AddIndexAccelerateMountPath                         = "/data/tmp"
	// TODO: it should be allowed to be configured
	AddIndexAccelerateVolumeSize = "100Gi"
)

const (
	TiDBOOMRecordVolumeName = "tidb-oom-record"

	// Fixed oom record path for tidb of version in [6.1.0, 6.4.0)
	// It's encoded from "/tmp/{uid}_tidb/{{ base64({host}:{port}/{status_host}:{status_port}) }}/tmp-storage/record"
	// in which all the variables are fixed:
	//     uid = 0, host = 0.0.0.0, port = 4000, status_host = 0.0.0.0, status_port = 10080
	TiDBOOMRecordContainerPathV1 = "/tmp/0_tidb/MC4wLjAuMDo0MDAwLzAuMC4wLjA6MTAwODA=/tmp-storage/record"
	// Fixed oom record path for tidb of version in [6.4.0, +inf)
	// It's a sub path in log directory. Due to not setting the log directory, it's placed to the root directory.
	TiDBOOMRecordContainerPathV2 = "/oom_record"
	// "/var/tidb_oom_record/<tenant_id>/<project_id>/<cluster_id>"
	TiDBOOMRecordHostPathTmp      = "/var/tidb_oom_record/%s/%s/%s"
	TiDBOOMRecordMountSubPathExpr = "$(NAMESPACE)/$(POD_NAME)"

	TiDBPlanReplayerVolumeName = "tidb-plan-replayer"
	// Fixed plan replayer path for tidb of version in [5.3.0, 6.4.0)
	// It's encoded from "/tmp/replayer/{pid}" in which pid is always 1.
	TiDBPlanReplayerContainerPathV1 = "/tmp/replayer/1"
	// Fixed plan replayer path for tidb of version in [6.4.0, +inf)
	// It's a sub path in log directory. Due to not setting the log directory, it's placed to the root directory.
	TiDBPlanReplayerContainerPathV2 = "/replayer"
	// "/var/tidb_plan_replayer/<tenant_id>/<project_id>/<cluster_id>"
	TiDBPlanReplayerHostPathTmp      = "/var/tidb_plan_replayer/%s/%s/%s"
	TiDBPlanReplayerMountSubPathExpr = "$(NAMESPACE)/$(POD_NAME)"
)

const (
	// the value of `GOMEMLIMIT` env is set to 80% of the memory limit of tidb container. @XuHuaiyu
	GOMEMLIMIT_RATE = 0.8
)

var defaultAuditLogSidecarCommand = []string{
	"sh",
	"-c",
	"touch /var/log/auditlog/tidb-audit.log; tail -n0 -F /var/log/auditlog/tidb-audit.log;",
}

// MachineTypeToIndexAccelerateSC maps GCP machine type prefixes to the corresponding
// StorageClass to be used for audit log volumes. Currently only n4-series nodes use HyperdiskBalancedStorageClass.
var MachineTypeToIndexAccelerateSC = map[string]string{
	// GCP:
	// n4-standard-2
	"n4-":  AddIndexAccelerateHyperdiskBalancedStorageClassName,
	"c4a-": AddIndexAccelerateHyperdiskBalancedStorageClassName,
}

type TiDBGroupConfig struct {
	// from namer
	Name       string
	Namespace  string
	ClusterRef string
	// infra api cluster name
	ClusterName         string
	AuthTokenSecretName string

	// from gitops config
	ConfigRevisionName string
	ImageTag           string
	BaseImage          string
	OriginalConfig     map[string]any

	Resources        corev1.ResourceRequirements
	StorageClassName string
	StorageSize      *resource.Quantity

	// from cluster component
	InstanceType string
	Replicas     int32
	Arch         string // CPU architecture of the instance type (e.g. "amd64", "arm64")

	// from context
	Zones      []string
	K8sVersion string

	HelperImage            string
	AuditLogSidecarImage   string
	AuditLogSidecarCommand []string
	AuditLogSidecarEnvs    []corev1.EnvVar

	// from cluster
	TenantID  string
	ProjectID string
	ClusterID string
	Provider  string

	// only
	TiDBServiceScope string

	PreStopSleepSeconds int32
}

type TiDBGroup interface {
	Build(gates feature.Gates, cfg *TiDBGroupConfig) *opv2.TiDBGroup
}

func NewTiDBGroupBuilder() TiDBGroup {
	return dbgBuilder{}
}

type dbgBuilder struct{}

func (dbgBuilder) Build(gates feature.Gates, cfg *TiDBGroupConfig) *opv2.TiDBGroup {
	slowlogVolumeName := "tidb-" + audit.SlowlogVolumeName

	configFile := mustBuildTiDBConfigFile(gates, cfg.OriginalConfig, util.FnvHash(cfg.Namespace), cfg.TiDBServiceScope)

	dbg := &opv2.TiDBGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name:      cfg.Name,
			Namespace: cfg.Namespace,
			Annotations: map[string]string{
				cloudv1.ConfigRevisionIdKey: cfg.ConfigRevisionName,
			},
		},
		Spec: opv2.TiDBGroupSpec{
			Cluster: opv2.ClusterReference{
				Name: cfg.ClusterRef,
			},
			Replicas:         &cfg.Replicas,
			SchedulePolicies: buildSchedulePolicies(cfg.Zones),
			Template: opv2.TiDBTemplate{
				ObjectMeta: opv2.ObjectMeta{
					Annotations: map[string]string{
						consts.TcInstanceTypeAnnotationKey: cfg.InstanceType,
						// NOTE(liubo02): useless, just keep to avoid restarting
						audit.MetricsPortAnnotationKey: strconv.Itoa(audit.MetricsPort),
					},
				},
				Spec: opv2.TiDBTemplateSpec{
					PreStop: setPreStopHook(gates, cfg),
					Version: cfg.ImageTag,
					Image:   ptr.To(cfg.BaseImage),
					Config:  opv2.ConfigFile(configFile),
					UpdateStrategy: opv2.UpdateStrategy{
						Config: opv2.ConfigUpdateStrategyHotReload,
					},
					Security: &opv2.TiDBSecurity{
						TLS: &opv2.TiDBTLS{
							MySQL: &opv2.TLS{Enabled: true},
						},
					},
					Volumes: []opv2.Volume{
						{
							Name: slowlogVolumeName,
							Mounts: []opv2.VolumeMount{
								{
									MountPath: audit.MountPath,
									Type:      opv2.VolumeMountTypeTiDBSlowLog,
								},
							},
							Storage:          resource.MustParse(audit.DefaultStorageSize),
							StorageClassName: ptr.To(audit.GetStorageClassName(cfg.InstanceType)),
						},
					},
					SlowLog: &opv2.TiDBSlowLog{
						Image: &cfg.HelperImage,
					},
					Overlay: &opv2.Overlay{
						PersistentVolumeClaims: []opv2.NamedPersistentVolumeClaimOverlay{
							{
								Name: slowlogVolumeName,
								PersistentVolumeClaim: opv2.PersistentVolumeClaimOverlay{
									ObjectMeta: opv2.ObjectMeta{
										// All app.kubernetes.io labels are for compatibility with v1
										// It may be used in prometheus rules and svc selectors in v1
										Labels: map[string]string{
											"app.kubernetes.io/component":  "tidb",
											"app.kubernetes.io/instance":   cfg.Name,
											"app.kubernetes.io/managed-by": "tidb-operator",
											"app.kubernetes.io/name":       "tidb-cluster",
										},
									},
								},
							},
						},
						Pod: &opv2.PodOverlay{
							ObjectMeta: opv2.ObjectMeta{
								Annotations: map[string]string{
									consts.TcInstanceTypeAnnotationKey: cfg.InstanceType,
									audit.MetricsPortAnnotationKey:     strconv.Itoa(audit.MetricsPort),
								},
								// All app.kubernetes.io labels are for compatibility with v1
								// It may be used in prometheus rules and svc selectors in v1
								Labels: map[string]string{
									"app.kubernetes.io/component":  "tidb",
									"app.kubernetes.io/instance":   cfg.Name,
									"app.kubernetes.io/managed-by": "tidb-operator",
									"app.kubernetes.io/name":       "tidb-cluster",
								},
							},
							Spec: &corev1.PodSpec{
								InitContainers: []corev1.Container{
									{
										// override requests and limits
										Name: opv2.ContainerNameTiDBSlowLog,
										Resources: corev1.ResourceRequirements{
											Requests: corev1.ResourceList{
												corev1.ResourceCPU:    resource.MustParse("20m"),
												corev1.ResourceMemory: resource.MustParse("5Mi"),
											},
											Limits: corev1.ResourceList{
												corev1.ResourceCPU:    resource.MustParse("100m"),
												corev1.ResourceMemory: resource.MustParse("50Mi"),
											},
										},
									},
									// Because the `net.*` sysctls are not allowed with host networking enabled,
									// we have to set them through a privileged init container.
									{
										Name:            "init",
										Image:           cfg.HelperImage,
										Command:         []string{"sh", "-c", "sysctl", "-w", "net.ipv4.tcp_keepalive_intvl=75", "net.ipv4.tcp_keepalive_time=300"},
										SecurityContext: &corev1.SecurityContext{Privileged: ptr.To(true)},
										Resources: corev1.ResourceRequirements{
											// same resources as the main container
											Requests: corev1.ResourceList{
												corev1.ResourceCPU:    *cfg.Resources.Requests.Cpu(),
												corev1.ResourceMemory: *cfg.Resources.Requests.Memory(),
											},
											Limits: corev1.ResourceList{
												// unset cpu limit
												corev1.ResourceMemory: *cfg.Resources.Limits.Memory(),
											},
										},
									},
								},
								Containers:                buildContainersForTiDB(gates, cfg, slowlogVolumeName),
								HostNetwork:               true,
								DNSPolicy:                 corev1.DNSClusterFirstWithHostNet,
								PriorityClassName:         consts.PriorityClasName,
								Affinity:                  utils.Affinity(cloudv1.ComponentTiDB),
								Tolerations:               buildTolerations(gates, string(cloudv1.ComponentTiDB), cfg.Arch, cfg.Provider),
								NodeSelector:              buildNodeSelector(gates, cfg.Namespace, cfg.ClusterName, string(cloudv1.ComponentTiDB), cfg.InstanceType, cfg.K8sVersion),
								TopologySpreadConstraints: buildTopologySpreadConstraints(gates, cfg.Zones),
								Volumes:                   buildTiDBVolumes(gates, cfg.InstanceType, cfg.TenantID, cfg.ProjectID, cfg.ClusterID),
							},
						},
					},
				},
			},
		},
	}
	if gates.Enabled(feature.TokenBasedAuth) {
		dbg.Spec.Template.Spec.Security.AuthToken = &opv2.TiDBAuthToken{
			JWKs: corev1.LocalObjectReference{
				Name: cfg.AuthTokenSecretName,
			},
		}
	}
	return dbg
}

func buildContainersForTiDB(gates feature.Gates, cfg *TiDBGroupConfig, slowlogVolumeName string) []corev1.Container {
	slowlogVolumeMount := corev1.VolumeMount{
		Name:      "ti-vol-" + slowlogVolumeName,
		MountPath: "/var/log/auditlog",
	}

	containers := []corev1.Container{
		*buildTiDBContainer(gates, cfg.Resources),
		{
			Name:  "statementlog",
			Image: cfg.HelperImage,
			Command: []string{
				"sh",
				"-c",
				"touch /var/log/auditlog/tidb-statements.log; tail -n0 -F /var/log/auditlog/tidb-statements.log;",
			},
			Resources: corev1.ResourceRequirements{
				Limits: corev1.ResourceList{
					corev1.ResourceCPU:    resource.MustParse("100m"),
					corev1.ResourceMemory: resource.MustParse("50Mi"),
				},
				Requests: corev1.ResourceList{
					corev1.ResourceCPU:    resource.MustParse("0m"),
					corev1.ResourceMemory: resource.MustParse("0Mi"),
				},
			},
			VolumeMounts: []corev1.VolumeMount{
				// manually mount slowlog volume to this container
				slowlogVolumeMount,
			},
		},
	}

	if gates.Enabled(feature.UseLegacyAuditPlugin) {
		cmd := defaultAuditLogSidecarCommand
		if len(cfg.AuditLogSidecarCommand) > 0 {
			cmd = cfg.AuditLogSidecarCommand
		}

		auditLogContainer := corev1.Container{
			Name:    audit.AuditLogName,
			Image:   cfg.AuditLogSidecarImage,
			Command: cmd,
			Resources: corev1.ResourceRequirements{
				Limits: corev1.ResourceList{
					corev1.ResourceCPU:    resource.MustParse("100m"),
					corev1.ResourceMemory: resource.MustParse("256Mi"),
				},
				Requests: corev1.ResourceList{
					corev1.ResourceCPU:    resource.MustParse("0m"),
					corev1.ResourceMemory: resource.MustParse("0Mi"),
				},
			},
			VolumeMounts: []corev1.VolumeMount{
				{
					Name:      audit.ConfigVolumeName,
					MountPath: audit.ConfigMountPath(),
				},
				slowlogVolumeMount,
			},
			Env: []corev1.EnvVar{
				{
					Name: "POD_NAME",
					ValueFrom: &corev1.EnvVarSource{
						FieldRef: &corev1.ObjectFieldSelector{
							FieldPath: "metadata.name",
						},
					},
				},
			},
		}

		if len(cfg.AuditLogSidecarEnvs) > 0 {
			auditLogContainer.Env = append(auditLogContainer.Env, cfg.AuditLogSidecarEnvs...)
		}

		if gates.Enabled(feature.SetupGCPCredentials) {
			auditLogContainer.Env = append(auditLogContainer.Env, corev1.EnvVar{
				Name:  "GOOGLE_APPLICATION_CREDENTIALS",
				Value: audit.GCPCredentialPath,
			})

			auditLogContainer.VolumeMounts = append(auditLogContainer.VolumeMounts, corev1.VolumeMount{
				Name:      audit.GCPCredentialsVolumeName,
				MountPath: filepath.Dir(audit.GCPCredentialPath),
			})
		}
		containers = append(containers, auditLogContainer)
	}

	return containers
}

func buildTiDBContainer(gates feature.Gates, res corev1.ResourceRequirements) *corev1.Container {
	c := &corev1.Container{
		Name: opv2.ContainerNameTiDB,
		Resources: corev1.ResourceRequirements{
			Requests: corev1.ResourceList{
				corev1.ResourceCPU:    *res.Requests.Cpu(),
				corev1.ResourceMemory: *res.Requests.Memory(),
			},
			Limits: corev1.ResourceList{
				corev1.ResourceMemory: *res.Limits.Memory(),
			},
		},
		Env: []corev1.EnvVar{
			{
				Name:  "TZ",
				Value: "UTC",
			},
			{
				Name: "NAMESPACE",
				ValueFrom: &corev1.EnvVarSource{
					FieldRef: &corev1.ObjectFieldSelector{
						FieldPath: "metadata.namespace",
					},
				},
			},
			{
				Name: "POD_NAME",
				ValueFrom: &corev1.EnvVarSource{
					FieldRef: &corev1.ObjectFieldSelector{
						FieldPath: "metadata.name",
					},
				},
			},
		},
	}

	if gates.Enabled(feature.IndexAccelerate) {
		c.VolumeMounts = append(c.VolumeMounts, corev1.VolumeMount{
			Name:      AddIndexAccelerateVolumeName,
			MountPath: AddIndexAccelerateMountPath,
		})
	}

	if gates.Enabled(feature.MountPlanReplayerV1) {
		c.VolumeMounts = append(c.VolumeMounts, corev1.VolumeMount{
			Name:        TiDBPlanReplayerVolumeName,
			MountPath:   TiDBPlanReplayerContainerPathV1,
			SubPathExpr: TiDBPlanReplayerMountSubPathExpr,
		})
	}

	if gates.Enabled(feature.MountPlanReplayerV2) {
		c.VolumeMounts = append(c.VolumeMounts, corev1.VolumeMount{
			Name:        TiDBPlanReplayerVolumeName,
			MountPath:   TiDBPlanReplayerContainerPathV2,
			SubPathExpr: TiDBPlanReplayerMountSubPathExpr,
		})
	}

	if gates.Enabled(feature.MountOOMPathV1) {
		c.VolumeMounts = append(c.VolumeMounts, corev1.VolumeMount{
			Name:        TiDBOOMRecordVolumeName,
			MountPath:   TiDBOOMRecordContainerPathV1,
			SubPathExpr: TiDBOOMRecordMountSubPathExpr,
		})
	}
	if gates.Enabled(feature.MountOOMPathV2) {
		c.VolumeMounts = append(c.VolumeMounts, corev1.VolumeMount{
			Name:        TiDBOOMRecordVolumeName,
			MountPath:   TiDBOOMRecordContainerPathV2,
			SubPathExpr: TiDBOOMRecordMountSubPathExpr,
		})
	}

	if gates.Enabled(feature.SetGoMemLimitEnv) {
		if mem := res.Limits.Memory(); mem != nil && !mem.IsZero() {
			tidbGoLimit := strconv.FormatInt(int64(GOMEMLIMIT_RATE*float64(mem.Value())), 10)
			c.Env = append(c.Env, corev1.EnvVar{Name: "GOMEMLIMIT", Value: tidbGoLimit})
		}
	}

	return c
}

func buildTiDBVolumes(gates feature.Gates, instanceType, tenantID, projectID, clusterID string) []corev1.Volume {
	var volumes []corev1.Volume

	if gates.Enabled(feature.IndexAccelerate) {
		volumes = append(volumes, corev1.Volume{
			Name: AddIndexAccelerateVolumeName,
			VolumeSource: corev1.VolumeSource{
				Ephemeral: &corev1.EphemeralVolumeSource{
					VolumeClaimTemplate: &corev1.PersistentVolumeClaimTemplate{
						ObjectMeta: metav1.ObjectMeta{},
						Spec: corev1.PersistentVolumeClaimSpec{
							AccessModes: []corev1.PersistentVolumeAccessMode{
								corev1.ReadWriteOnce,
							},
							StorageClassName: ptr.To(GetAddIndexAccelerateStorageClassName(instanceType)),
							Resources: corev1.VolumeResourceRequirements{
								Requests: corev1.ResourceList{
									corev1.ResourceStorage: resource.MustParse(AddIndexAccelerateVolumeSize),
								},
							},
						},
					},
				},
			},
		})
	}

	if gates.Enabled(feature.MountPlanReplayerV1) || gates.Enabled(feature.MountPlanReplayerV2) {
		volumes = append(volumes, corev1.Volume{
			Name: TiDBPlanReplayerVolumeName,
			VolumeSource: corev1.VolumeSource{
				HostPath: &corev1.HostPathVolumeSource{
					Path: fmt.Sprintf(TiDBPlanReplayerHostPathTmp, tenantID, projectID, clusterID),
					Type: ptr.To(corev1.HostPathDirectoryOrCreate),
				},
			},
		})
	}

	if gates.Enabled(feature.MountOOMPathV1) || gates.Enabled(feature.MountOOMPathV2) {
		volumes = append(volumes, corev1.Volume{
			Name: TiDBOOMRecordVolumeName,
			VolumeSource: corev1.VolumeSource{
				HostPath: &corev1.HostPathVolumeSource{
					Path: fmt.Sprintf(TiDBOOMRecordHostPathTmp, tenantID, projectID, clusterID),
					Type: ptr.To(corev1.HostPathDirectoryOrCreate),
				},
			},
		})
	}

	if gates.Enabled(feature.UseLegacyAuditPlugin) {
		volumes = append(volumes, audit.BuildConfigMapVolume(audit.ConfigMapName(clusterID)))
	}

	if gates.Enabled(feature.SetupGCPCredentials) {
		volumes = append(volumes, corev1.Volume{
			Name: audit.GCPCredentialsVolumeName,
			VolumeSource: corev1.VolumeSource{
				Secret: &corev1.SecretVolumeSource{
					SecretName: audit.GCPSecretName,
				},
			},
		})
	}

	return volumes
}

func mustBuildTiDBConfigFile(gates feature.Gates, cfg map[string]any, allowedCN, serviceScope string) []byte {
	// TODO(liubo02): use another lib to generate toml cfg
	conf := opv1alpha1.NewTiDBConfig()
	conf.MP = cfg
	conf.Set("security.cluster-verify-cn", []string{allowedCN})
	conf.Set("performance.tcp-keep-alive", true)
	conf.Set("security.enable-sem", true)
	if gates.Enabled(feature.IndexAccelerate) {
		conf.Set("temp-dir", AddIndexAccelerateMountPath)
	}

	// always mount audit plugins
	conf.Set("instance.plugin_dir", "/plugins")
	conf.Set("instance.plugin_load", "audit-1")

	if gates.Enabled(feature.DistTaskSupported) && serviceScope != "" {
		// do not set service scope for default tidb group
		conf.Set("instance.tidb_service_scope", serviceScope)
	}

	configBytes, err := conf.MarshalTOML()
	if err != nil {
		panic("failed to marshal tidb config" + err.Error())
	}
	return configBytes
}

func setPreStopHook(gates feature.Gates, cfg *TiDBGroupConfig) *opv2.TiDBPreStop {
	if gates.Enabled(feature.SetPreStopHook) {
		return &opv2.TiDBPreStop{
			SleepSeconds: cfg.PreStopSleepSeconds,
		}
	}
	return nil
}

func GetAddIndexAccelerateStorageClassName(instanceType string) string {
	storageClass := AddIndexAccelerateStorageClassName // default storage class
	// Check if instanceType has a special storage class mapping
	for prefix, sc := range MachineTypeToIndexAccelerateSC {
		if strings.HasPrefix(instanceType, prefix) {
			storageClass = sc
			break
		}
	}
	return storageClass
}