package feature

import (
	"golang.org/x/mod/semver"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/cmpver"
	"github.com/tidbcloud/infra-provider/pkg/consts"
)

var minAddIndexAccelerateVersion = cmpver.MustNewConstraint(cmpver.GreaterOrEqual, "v6.3.0")

const (
	// https://docs.pingcap.com/tidb/stable/tidb-distributed-execution-framework#usage
	minDistTaskSupportedVersion = "v8.1.0"

	// The minimum TiDB version that support `tidb_token_auth`
	minTokenBasedAuthVersion = "v6.4.0"
)

const (
	UseClusterAutoscaler                = "UseClusterAutoscaler"
	AddK8sVersionToNodeSelector         = "AddK8sVersionToNodeSelector"
	Add3LayerNodeSelectorAndTolerations = "Add3LayerNodeSelectorAndTolerations"
	AddNodePoolDataPlaneLabel           = "AddNodePoolDataPlaneLabel"
	UseCMEK                             = "UseCMEK"
	UseVolumeAttributesClass            = "UseVolumeAttributesClass"
	IndexAccelerate                     = "IndexAccelerate"
	MountPlanReplayerV1                 = "MountPlanReplayerV1"
	MountPlanReplayerV2                 = "MountPlanReplayerV2"
	MountOOMPathV1                      = "MountOOMPathV1"
	MountOOMPathV2                      = "MountOOMPathV2"
	UseLegacyAuditPlugin                = "UseLegacyAuditPlugin"
	// SetupGCPCredentials should be enabled only when UseLegacyAuditPlugin is enabled and provider is GKE.
	// Because we need to push audit logs to GCS bucket.
	SetupGCPCredentials = "SetupGCPCredentials"

	TokenBasedAuth = "TokenBasedAuth"
	// https://docs.pingcap.com/tidb/stable/tidb-distributed-execution-framework#usage
	// https://github.com/tidbcloud/infra-provider/pull/4083
	DistTaskSupported = "DistTaskSupported"

	// for 6.1.3 <= TiDB-version < 6.5, we need to set `GOMEMLIMIT` env to avoid OOM.
	// as this env has no side effect for version >= 6.5 (TiDB kernel can handle it correctly, @XuHuaiyu),
	// and all newly created clusters should >= 6.1.3 now,
	// we can set it for all new clusters.
	// NOTE: do not set this env for existing clusters, otherwise the rolling update will be triggered.
	//
	// This feature is enabled for all tidbs managed by v2 operator now
	// It cannot be enabled for tidbs < v6.1.3
	// TODO(liubo02): remove it later if all tidbs are migrated to v2 and >= v6.1.3
	SetGoMemLimitEnv = "SetGoMemLimitEnv"

	// This feature was only enabled in gke before but not really worked
	// Set it to always false
	// TODO(liubo02): remove it later if all gke clusters don't enable kms with "tikv.cloud.pingcap.com/enable-kms=true" in annotation
	AddKmsCredentialsEnv = "AddKmsCredentialsEnv"

	// This feature is only for tidb operator v2
	// Current readiness probe produces too many warn logs and is not eanbled in v1
	// NOTE: If all clusters in prod have enabled this feature, it can be removed
	DisableV2PDReadinessProbe = "DisableV2PDReadinessProbe"

	// On GCP, we will set LB's connection draining timeout to 15s by default,
	// so we need to set preStop hook to sleep at least 30 seconds to achieve graceful shutdown.
	SetPreStopHook = "SetPreStopHook"

	// This feature is only for tidb operator v2.
	UsePDReadyAPI = "UsePDReadyAPI"
)

type gates map[string]struct{}

func (g gates) Enabled(feat string) bool {
	_, ok := g[feat]
	return ok
}

func (g gates) Set(feat string, val bool) {
	enabled := g.Enabled(feat)
	if val && !enabled {
		g[feat] = struct{}{}
		return
	}
	if !val && enabled {
		delete(g, feat)
	}
}

func NewGates(feats ...string) Gates {
	g := gates{}
	for _, feat := range feats {
		g[feat] = struct{}{}
	}

	return g
}

type Gates interface {
	Enabled(string) bool
}

func Init(
	c *cloudv1.Cluster,
	provider string,
	k8sVersion string,
	tidbVersion string,
) Gates {
	fg := gates{}
	fg.Set(SetGoMemLimitEnv, true)
	fg.Set(AddKmsCredentialsEnv, false)

	fg.Set(UseClusterAutoscaler,
		provider == consts.ProviderTypeGKE || provider == consts.ProviderTypeAKS)
	fg.Set(Add3LayerNodeSelectorAndTolerations,
		provider == consts.ProviderTypeGKE && c.FeatureEnabled(consts.ClusterFeatureGCP3LayerNodePool))
	fg.Set(UseCMEK,
		c.FeatureEnabled(consts.ClusterFeatureUseCMEK))
	fg.Set(UseVolumeAttributesClass,
		c.FeatureEnabled(consts.ClusterFeatureUseVolumeAttributesClass))
	fg.Set(IndexAccelerate,
		addIndexAccelerateSupported(tidbVersion))
	fg.Set(MountPlanReplayerV1,
		c.FeatureEnabled(consts.ClusterFeatureTiDBExposeReplayerPath) && checkMountPlanReplayerSupportVersionV1(tidbVersion))
	fg.Set(MountPlanReplayerV2,
		c.FeatureEnabled(consts.ClusterFeatureTiDBExposeReplayerPath) && checkMountPlanReplayerSupportVersionV2(tidbVersion))
	fg.Set(MountOOMPathV1,
		c.FeatureEnabled(consts.ClusterFeatureTiDBExposeOOMPath) && checkMountOOMPathSupportVersionV1(tidbVersion))
	fg.Set(MountOOMPathV2,
		c.FeatureEnabled(consts.ClusterFeatureTiDBExposeOOMPath) && checkMountOOMPathSupportVersionV2(tidbVersion))
	fg.Set(TokenBasedAuth, tokenBasedAuthEnabled(c, tidbVersion))
	fg.Set(DistTaskSupported,
		distTaskSupported(tidbVersion))
	fg.Set(AddNodePoolDataPlaneLabel,
		fg.Enabled(UseClusterAutoscaler) && !c.FeatureEnabled(consts.ClusterFeatureDisableNodePoolDataPlaneLabel))
	fg.Set(AddK8sVersionToNodeSelector,
		fg.Enabled(UseClusterAutoscaler) && k8sVersion != "")
	// FIXME(Huaxi): modify the feature gate when the new audit log is ready.
	fg.Set(UseLegacyAuditPlugin, true)
	fg.Set(SetupGCPCredentials, fg.Enabled(UseLegacyAuditPlugin) && provider == consts.ProviderTypeGKE)
	fg.Set(DisableV2PDReadinessProbe, c.FeatureEnabled(consts.ClusterFeatureDisableV2PDReadinessProbe))
	fg.Set(SetPreStopHook, c.FeatureEnabled(consts.ClusterFeatureEnableGracefulShutdown))
	fg.Set(UsePDReadyAPI, c.FeatureEnabled(consts.ClusterFeatureUsePDReadyAPIV2))
	return fg
}

func tokenBasedAuthEnabled(cluster *cloudv1.Cluster, version string) bool {
	if cluster.Labels[cloudv1.ClusterLabelWantTokenBasedAuth] != "true" {
		return false
	}

	if !tokenBasedAuthSupported(version) {
		return false
	}
	return true
}

func tokenBasedAuthSupported(version string) bool {
	return semver.Compare(version, minTokenBasedAuthVersion) >= 0
}

func addIndexAccelerateSupported(version string) bool {
	ok, err := minAddIndexAccelerateVersion.Check(version)
	if err != nil {
		return false
	}
	return ok
}

// TODO(liubo02): change to use cmpver.MustCompare ?
func checkMountPlanReplayerSupportVersionV1(version string) bool {
	return semver.Compare(version, "v5.3.0") >= 0 && semver.Compare(version, "v6.4.0") < 0
}

// [v6.4.0, +inf)
func checkMountPlanReplayerSupportVersionV2(version string) bool {
	return semver.Compare(version, "v6.4.0") >= 0
}

func checkMountOOMPathSupportVersionV1(version string) bool {
	return semver.Compare(version, "v6.1.0") >= 0 && semver.Compare(version, "v6.4.0") < 0
}

// [v6.4.0, +inf)
func checkMountOOMPathSupportVersionV2(version string) bool {
	return semver.Compare(version, "v6.4.0") >= 0
}

func distTaskSupported(version string) bool {
	return semver.Compare(version, minDistTaskSupportedVersion) >= 0
}