/*
Copyright 2024 The Kubernetes Authors.

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.
*/

package compatibility

import (
	"context"
	"fmt"
	"sort"
	"strings"
	"sync"

	"github.com/spf13/pflag"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/version"
	cliflag "k8s.io/component-base/cli/flag"
	"k8s.io/component-base/featuregate"
	"k8s.io/component-base/metrics/prometheus/compatversion"
	"k8s.io/klog/v2"
)

const (
	// DefaultKubeComponent is the component name for k8s control plane components.
	DefaultKubeComponent = "kube"

	klogLevel = 2
)

type VersionMapping func(from *version.Version) *version.Version

// ComponentGlobals stores the global variables for a component for easy access, including feature gate and effective version.
type ComponentGlobals struct {
	effectiveVersion MutableEffectiveVersion
	featureGate      featuregate.MutableVersionedFeatureGate

	// componentVersionMapping contains the mapping from the version of this component
	// to the version of another component.
	componentVersionMapping map[string]VersionMapping
	// dependentComponentVersion stores whether or not this component's version is dependent through mapping on another component.
	// If true, the emulation/minCompatibility version cannot be set from the flag, or version mapping from another component.
	dependentComponentVersion bool
}

// ComponentGlobalsRegistry stores the global variables for different components for easy access, including feature gate and effective version of each component.
type ComponentGlobalsRegistry interface {
	// EffectiveVersionFor returns the EffectiveVersion registered under the component.
	// Returns nil if the component is not registered.
	EffectiveVersionFor(component string) EffectiveVersion
	// FeatureGateFor returns the FeatureGate registered under the component.
	// Returns nil if the component is not registered.
	FeatureGateFor(component string) featuregate.FeatureGate
	// Register registers the EffectiveVersion and FeatureGate for a component.
	// returns error if the component is already registered.
	Register(component string, effectiveVersion MutableEffectiveVersion, featureGate featuregate.MutableVersionedFeatureGate) error
	// ComponentGlobalsOrRegister would return the registered global variables for the component if it already exists in the registry.
	// Otherwise, the provided variables would be registered under the component, and the same variables would be returned.
	ComponentGlobalsOrRegister(component string, effectiveVersion MutableEffectiveVersion, featureGate featuregate.MutableVersionedFeatureGate) (MutableEffectiveVersion, featuregate.MutableVersionedFeatureGate)
	// AddFlags adds flags of "--emulated-version", "--min-compatibility-version" and "--feature-gates"
	AddFlags(fs *pflag.FlagSet)
	// Set sets the flags for all global variables for all components registered.
	// A component's feature gate and effective version would not be updated until Set() is called.
	Set() error
	// SetFallback calls Set() if it has never been called.
	SetFallback() error
	// Validate calls the Validate() function for all the global variables for all components registered.
	Validate() []error
	// Reset removes all stored ComponentGlobals, configurations, and version mappings.
	Reset()
	// SetVersionMapping sets the mapping from the emulation/minCompatibility version of one component
	// to the emulation/minCompatibility version of another component.
	// Once set, the emulation/minCompatibility version of the toComponent will be determined by the emulation/minCompatibility version of the fromComponent,
	// and cannot be set from cmd flags anymore.
	// For a given component, its emulation/minCompatibility version can only depend on one other component, no multiple dependency is allowed.
	SetVersionMapping(fromComponent, toComponent string, f VersionMapping) error
	// AddMetrics adds metrics for the emulation version of a component.
	AddMetrics()
}

type componentGlobalsRegistry struct {
	componentGlobals map[string]*ComponentGlobals
	mutex            sync.RWMutex
	// emulationVersionConfig stores the list of component name to emulation version set from the flag.
	// When the `--emulated-version` flag is parsed, it would not take effect until Set() is called,
	// because we have to enforce of order of setting the emulation version first, then min compatibility version, and last the feature gate.
	emulationVersionConfig []string
	// minCompatibilityVersionConfig stores the list of component name to min compatibility version set from the flag.
	// When the `--min-compatibility-version` flag is parsed, it would not take effect until Set() is called,
	// because we have to enforce of order of setting the emulation version first, then min compatibility version, and last the feature gate.
	minCompatibilityVersionConfig []string
	// featureGatesConfig stores the map of component name to the list of feature gates set from the flag.
	// When the `--feature-gates` flag is parsed, it would not take effect until Set() is called,
	// because we have to enforce of order of setting the emulation version first, then min compatibility version, and last the feature gate.
	featureGatesConfig map[string][]string
	// featureGatesConfigFlags stores a pointer to the flag value, allowing other commands
	// to append to the feature gates configuration rather than overwriting it
	featureGatesConfigFlags *cliflag.ColonSeparatedMultimapStringString
	// set stores if the Set() function for the registry is already called.
	set bool
}

func NewComponentGlobalsRegistry() *componentGlobalsRegistry {
	return &componentGlobalsRegistry{
		componentGlobals:              make(map[string]*ComponentGlobals),
		emulationVersionConfig:        nil,
		minCompatibilityVersionConfig: nil,
		featureGatesConfig:            nil,
	}
}

func (r *componentGlobalsRegistry) AddMetrics() {
	for name, globals := range r.componentGlobals {
		effectiveVersion := globals.effectiveVersion
		if effectiveVersion == nil {
			continue
		}
		compatversion.RecordCompatVersionInfo(context.Background(), name, effectiveVersion.BinaryVersion().String(), effectiveVersion.EmulationVersion().String(), effectiveVersion.MinCompatibilityVersion().String())
	}
}

func (r *componentGlobalsRegistry) Reset() {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	r.componentGlobals = make(map[string]*ComponentGlobals)
	r.emulationVersionConfig = nil
	r.minCompatibilityVersionConfig = nil
	r.featureGatesConfig = nil
	r.featureGatesConfigFlags = nil
	r.set = false
}

func (r *componentGlobalsRegistry) EffectiveVersionFor(component string) EffectiveVersion {
	r.mutex.RLock()
	defer r.mutex.RUnlock()
	globals, ok := r.componentGlobals[component]
	if !ok {
		return nil
	}
	return globals.effectiveVersion
}

func (r *componentGlobalsRegistry) FeatureGateFor(component string) featuregate.FeatureGate {
	r.mutex.RLock()
	defer r.mutex.RUnlock()
	globals, ok := r.componentGlobals[component]
	if !ok {
		return nil
	}
	return globals.featureGate
}

func (r *componentGlobalsRegistry) unsafeRegister(component string, effectiveVersion MutableEffectiveVersion, featureGate featuregate.MutableVersionedFeatureGate) error {
	if _, ok := r.componentGlobals[component]; ok {
		return fmt.Errorf("component globals of %s already registered", component)
	}
	if featureGate != nil {
		if err := featureGate.SetEmulationVersionAndMinCompatibilityVersion(
			effectiveVersion.EmulationVersion(), effectiveVersion.MinCompatibilityVersion()); err != nil {
			return err
		}
	}
	c := ComponentGlobals{
		effectiveVersion:        effectiveVersion,
		featureGate:             featureGate,
		componentVersionMapping: make(map[string]VersionMapping),
	}
	r.componentGlobals[component] = &c
	return nil
}

func (r *componentGlobalsRegistry) Register(component string, effectiveVersion MutableEffectiveVersion, featureGate featuregate.MutableVersionedFeatureGate) error {
	if effectiveVersion == nil {
		return fmt.Errorf("cannot register nil effectiveVersion")
	}
	r.mutex.Lock()
	defer r.mutex.Unlock()
	return r.unsafeRegister(component, effectiveVersion, featureGate)
}

func (r *componentGlobalsRegistry) ComponentGlobalsOrRegister(component string, effectiveVersion MutableEffectiveVersion, featureGate featuregate.MutableVersionedFeatureGate) (MutableEffectiveVersion, featuregate.MutableVersionedFeatureGate) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	globals, ok := r.componentGlobals[component]
	if ok {
		return globals.effectiveVersion, globals.featureGate
	}
	utilruntime.Must(r.unsafeRegister(component, effectiveVersion, featureGate))
	return effectiveVersion, featureGate
}

func (r *componentGlobalsRegistry) unsafeKnownFeatures() []string {
	var known []string
	for component, globals := range r.componentGlobals {
		if globals.featureGate == nil {
			continue
		}
		for _, f := range globals.featureGate.KnownFeatures() {
			known = append(known, component+":"+f)
		}
	}
	sort.Strings(known)
	return known
}

func (r *componentGlobalsRegistry) unsafeVersionFlagOptions(isEmulation bool) []string {
	var vs []string
	for component, globals := range r.componentGlobals {
		if globals.dependentComponentVersion {
			continue
		}
		if isEmulation {
			vs = append(vs, fmt.Sprintf("%s=%s", component, globals.effectiveVersion.AllowedEmulationVersionRange()))
		} else {
			vs = append(vs, fmt.Sprintf("%s=%s", component, globals.effectiveVersion.AllowedMinCompatibilityVersionRange()))
		}
	}
	sort.Strings(vs)
	return vs
}

func (r *componentGlobalsRegistry) AddFlags(fs *pflag.FlagSet) {
	if r == nil {
		return
	}
	r.mutex.Lock()
	defer r.mutex.Unlock()
	for _, globals := range r.componentGlobals {
		if globals.featureGate != nil {
			globals.featureGate.Close()
		}
	}

	fs.StringSliceVar(&r.emulationVersionConfig, "emulated-version", r.emulationVersionConfig, ""+
		"The versions different components emulate their capabilities (APIs, features, ...) of.\n"+
		"If set, the component will emulate the behavior of this version instead of the underlying binary version.\n"+
		"Version format could only be major.minor, for example: '--emulated-version=wardle=1.2,kube=1.31'.\nOptions are: "+strings.Join(r.unsafeVersionFlagOptions(true), ",")+
		"\nIf the component is not specified, defaults to \"kube\"")

	// min-compatibility-version is typically the minimal binary version of all control plane components the server expects to coexist with throughout the lifecycle of this server.
	// If set to smaller than the server emulated version, the newest CEL features/libraries/parameters introduced after the min compatibility version would not be turned on,
	// feature stages with a MinCompatibilityVersion higher than this value will be skipped, and the resource storage version will be set based on the min compatibility version.
	// It can be set to equal to the binary version after all control plane components are upgraded, after which features with compatibility implications could be enabled, the newest CEL features/libraries/parameters would turned on, but rollbacks are no longer supported.
	fs.StringSliceVar(&r.minCompatibilityVersionConfig, "min-compatibility-version", r.minCompatibilityVersionConfig, ""+
		"The min version of control plane components the server should be compatible with.\n"+
		"Must be less or equal to the emulated-version. Version format could only be major.minor, for example: '--min-compatibility-version=wardle=1.2,kube=1.31'.\n"+
		"Options are: "+strings.Join(r.unsafeVersionFlagOptions(false), ",")+
		"\nIf the component is not specified, defaults to \"kube\"")

	if r.featureGatesConfigFlags == nil {
		r.featureGatesConfigFlags = cliflag.NewColonSeparatedMultimapStringStringAllowDefaultEmptyKey(&r.featureGatesConfig)
	}
	fs.Var(r.featureGatesConfigFlags, "feature-gates", "Comma-separated list of component:key=value pairs that describe feature gates for alpha/experimental features of different components.\n"+
		"If the component is not specified, defaults to \"kube\". This flag can be repeatedly invoked. For example: --feature-gates 'wardle:featureA=true,wardle:featureB=false' --feature-gates 'kube:featureC=true'"+
		"Options are:\n"+strings.Join(r.unsafeKnownFeatures(), "\n"))
}

type componentVersion struct {
	component string
	ver       *version.Version
}

// getFullVersionConfig expands the given version config with registered version mapping,
// and returns the map of component to Version.
func (r *componentGlobalsRegistry) getFullVersionConfig(
	versionConfigMap map[string]*version.Version) (map[string]*version.Version, error) {
	result := map[string]*version.Version{}
	setQueue := []componentVersion{}
	for comp, ver := range versionConfigMap {
		if _, ok := r.componentGlobals[comp]; !ok {
			return result, fmt.Errorf("component not registered: %s", comp)
		}
		klog.V(klogLevel).Infof("setting version %s=%s", comp, ver.String())
		setQueue = append(setQueue, componentVersion{comp, ver})
	}
	for len(setQueue) > 0 {
		cv := setQueue[0]
		if _, visited := result[cv.component]; visited {
			return result, fmt.Errorf("setting version of %s more than once, probably version mapping loop", cv.component)
		}
		setQueue = setQueue[1:]
		result[cv.component] = cv.ver
		for toComp, f := range r.componentGlobals[cv.component].componentVersionMapping {
			toVer := f(cv.ver)
			if toVer == nil {
				return result, fmt.Errorf("got nil version from mapping of %s=%s to component:%s", cv.component, cv.ver.String(), toComp)
			}
			klog.V(klogLevel).Infof("setting version %s=%s from version mapping of %s=%s", toComp, toVer.String(), cv.component, cv.ver.String())
			setQueue = append(setQueue, componentVersion{toComp, toVer})
		}
	}
	return result, nil
}

func toVersionMap(versionConfig []string) (map[string]*version.Version, error) {
	m := map[string]*version.Version{}
	for _, compVer := range versionConfig {
		// default to "kube" of component is not specified
		k := "kube"
		v := compVer
		if strings.Contains(compVer, "=") {
			arr := strings.SplitN(compVer, "=", 2)
			if len(arr) != 2 {
				return m, fmt.Errorf("malformed pair, expect string=string")
			}
			k = strings.TrimSpace(arr[0])
			v = strings.TrimSpace(arr[1])
		}
		ver, err := version.Parse(v)
		if err != nil {
			return m, err
		}
		if ver.Patch() != 0 {
			return m, fmt.Errorf("patch version not allowed, got: %s=%s", k, ver.String())
		}
		if existingVer, ok := m[k]; ok {
			return m, fmt.Errorf("duplicate version flag, %s=%s and %s=%s", k, existingVer.String(), k, ver.String())
		}
		m[k] = ver
	}
	return m, nil
}

func (r *componentGlobalsRegistry) SetFallback() error {
	r.mutex.Lock()
	set := r.set
	r.mutex.Unlock()
	if set {
		return nil
	}
	klog.Warning("setting componentGlobalsRegistry in SetFallback. We recommend calling componentGlobalsRegistry.Set()" +
		" right after parsing flags to avoid using feature gates before their final values are set by the flags.")
	return r.Set()
}

func (r *componentGlobalsRegistry) Set() error {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	r.set = true
	emulationVersionConfigMap, err := toVersionMap(r.emulationVersionConfig)
	if err != nil {
		return err
	}
	for comp := range emulationVersionConfigMap {
		if _, ok := r.componentGlobals[comp]; !ok {
			return fmt.Errorf("component not registered: %s", comp)
		}
		// only components without any dependencies can be set from the flag.
		if r.componentGlobals[comp].dependentComponentVersion {
			return fmt.Errorf("EmulationVersion of %s is set by mapping, cannot set it by flag", comp)
		}
	}
	if emulationVersions, err := r.getFullVersionConfig(emulationVersionConfigMap); err != nil {
		return err
	} else {
		for comp, ver := range emulationVersions {
			r.componentGlobals[comp].effectiveVersion.SetEmulationVersion(ver)
		}
	}

	minCompatibilityVersionConfigMap, err := toVersionMap(r.minCompatibilityVersionConfig)
	if err != nil {
		return err
	}
	for comp := range minCompatibilityVersionConfigMap {
		if _, ok := r.componentGlobals[comp]; !ok {
			return fmt.Errorf("component not registered: %s", comp)
		}
		// only components without any dependencies can be set from the flag.
		if r.componentGlobals[comp].dependentComponentVersion {
			return fmt.Errorf("MinCompatibilityVersion of %s is set by mapping, cannot set it by flag", comp)
		}
	}
	if minCompatibilityVersions, err := r.getFullVersionConfig(minCompatibilityVersionConfigMap); err != nil {
		return err
	} else {
		for comp, ver := range minCompatibilityVersions {
			r.componentGlobals[comp].effectiveVersion.SetMinCompatibilityVersion(ver)
		}
	}
	// Set feature gate emulation version before setting feature gate flag values.
	for comp, globals := range r.componentGlobals {
		if globals.featureGate == nil {
			continue
		}
		klog.V(klogLevel).Infof("setting %s:feature gate emulation version to %s, min compatibility version to %s", comp, globals.effectiveVersion.EmulationVersion().String(), globals.effectiveVersion.MinCompatibilityVersion().String())
		if err := globals.featureGate.SetEmulationVersionAndMinCompatibilityVersion(globals.effectiveVersion.EmulationVersion(), globals.effectiveVersion.MinCompatibilityVersion()); err != nil {
			return err
		}
	}
	for comp, fg := range r.featureGatesConfig {
		if comp == "" {
			if _, ok := r.featureGatesConfig[DefaultKubeComponent]; ok {
				return fmt.Errorf("set kube feature gates with default empty prefix or kube: prefix consistently, do not mix use")
			}
			comp = DefaultKubeComponent
		}
		if _, ok := r.componentGlobals[comp]; !ok {
			return fmt.Errorf("component not registered: %s", comp)
		}
		featureGate := r.componentGlobals[comp].featureGate
		if featureGate == nil {
			return fmt.Errorf("component featureGate not registered: %s", comp)
		}
		flagVal := strings.Join(fg, ",")
		klog.V(klogLevel).Infof("setting %s:feature-gates=%s", comp, flagVal)
		if err := featureGate.Set(flagVal); err != nil {
			return err
		}
	}
	return nil
}

func (r *componentGlobalsRegistry) Validate() []error {
	var errs []error
	r.mutex.Lock()
	defer r.mutex.Unlock()
	for _, globals := range r.componentGlobals {
		errs = append(errs, globals.effectiveVersion.Validate()...)
		var features map[featuregate.Feature]featuregate.FeatureSpec
		if globals.featureGate != nil {
			errs = append(errs, globals.featureGate.Validate()...)
			features = globals.featureGate.GetAll()
		}
		binaryVersion := globals.effectiveVersion.BinaryVersion()
		emulatedVersion := globals.effectiveVersion.EmulationVersion()
		if binaryVersion.GreaterThan(emulatedVersion) {
			if enabled := enabledAlphaFeatures(features, globals); len(enabled) != 0 {
				klog.Warningf("component has alpha features enabled in emulated version, this is unsupported: features=%v", enabled)
			}
		}
	}
	return errs
}

func enabledAlphaFeatures(features map[featuregate.Feature]featuregate.FeatureSpec, globals *ComponentGlobals) []string {
	var enabled []string
	for feat, featSpec := range features {
		if featSpec.PreRelease == featuregate.Alpha && globals.featureGate.Enabled(feat) {
			enabled = append(enabled, string(feat))
		}
	}
	return enabled
}

func (r *componentGlobalsRegistry) SetVersionMapping(fromComponent, toComponent string, f VersionMapping) error {
	if f == nil {
		return nil
	}
	klog.V(klogLevel).Infof("setting EmulationVersion mapping from %s to %s", fromComponent, toComponent)
	r.mutex.Lock()
	defer r.mutex.Unlock()
	if _, ok := r.componentGlobals[fromComponent]; !ok {
		return fmt.Errorf("component not registered: %s", fromComponent)
	}
	if _, ok := r.componentGlobals[toComponent]; !ok {
		return fmt.Errorf("component not registered: %s", toComponent)
	}
	// check multiple dependency
	if r.componentGlobals[toComponent].dependentComponentVersion {
		return fmt.Errorf("mapping of %s already exists from another component", toComponent)
	}
	r.componentGlobals[toComponent].dependentComponentVersion = true

	versionMapping := r.componentGlobals[fromComponent].componentVersionMapping
	if _, ok := versionMapping[toComponent]; ok {
		return fmt.Errorf("EmulationVersion from %s to %s already exists", fromComponent, toComponent)
	}
	versionMapping[toComponent] = f
	klog.V(klogLevel).Infof("setting the default EmulationVersion of %s based on mapping from the default EmulationVersion of %s", toComponent, fromComponent)
	defaultFromVersion := r.componentGlobals[fromComponent].effectiveVersion.EmulationVersion()
	emulationVersions, err := r.getFullVersionConfig(map[string]*version.Version{fromComponent: defaultFromVersion})
	if err != nil {
		return err
	}
	for comp, ver := range emulationVersions {
		r.componentGlobals[comp].effectiveVersion.SetEmulationVersion(ver)
	}
	return nil
}
