package kubelet

import (
	"context"
	_ "embed"
	"encoding/json"
	"fmt"
	"net"
	"net/url"
	"os"
	"path"
	"strings"
	"time"

	"go.uber.org/zap"
	"golang.org/x/mod/semver"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	k8skubelet "k8s.io/kubelet/config/v1beta1"

	"github.com/avast/retry-go/v5"
	"github.com/aws/smithy-go/ptr"
	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/api"
	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/aws/imds"
	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/containerd"
	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/system"
	"github.com/awslabs/amazon-eks-ami/nodeadm/internal/util"
)

const (
	kubeletConfigRoot = "/etc/kubernetes/kubelet"
	kubeletConfigFile = "config.json"
	kubeletConfigDir  = "config.json.d"
	kubeletConfigPerm = 0644

	// systemdResolvedUplinkResolvConf is maintained by systemd-resolved and
	// always lists the real upstream DNS servers, never the 127.0.0.53 stub.
	// see: https://www.freedesktop.org/software/systemd/man/latest/systemd-resolved.service.html#/etc/resolv.conf
	systemdResolvedUplinkResolvConf = "/run/systemd/resolve/resolv.conf"
)

// kubeletConfig is an internal-only representation of the kubelet configuration
// that is generated using sane defaults for EKS. It is a subset of the upstream
// KubeletConfiguration types:
// https://pkg.go.dev/k8s.io/kubelet/config/v1beta1#KubeletConfiguration
type kubeletConfig struct {
	Address                         string                           `json:"address"`
	Authentication                  k8skubelet.KubeletAuthentication `json:"authentication"`
	Authorization                   k8skubelet.KubeletAuthorization  `json:"authorization"`
	CgroupDriver                    string                           `json:"cgroupDriver"`
	CgroupRoot                      string                           `json:"cgroupRoot"`
	ClusterDNS                      []string                         `json:"clusterDNS"`
	ClusterDomain                   string                           `json:"clusterDomain"`
	ContainerRuntimeEndpoint        string                           `json:"containerRuntimeEndpoint"`
	ImageServiceEndpoint            string                           `json:"imageServiceEndpoint,omitempty"`
	EvictionHard                    map[string]string                `json:"evictionHard,omitempty"`
	FeatureGates                    map[string]bool                  `json:"featureGates"`
	HairpinMode                     string                           `json:"hairpinMode"`
	KubeAPIBurst                    *int                             `json:"kubeAPIBurst,omitempty"`
	KubeAPIQPS                      *int                             `json:"kubeAPIQPS,omitempty"`
	KubeReserved                    map[string]string                `json:"kubeReserved,omitempty"`
	KubeReservedCgroup              *string                          `json:"kubeReservedCgroup,omitempty"`
	Logging                         loggingConfiguration             `json:"logging"`
	MaxPods                         int32                            `json:"maxPods,omitempty"`
	ProtectKernelDefaults           bool                             `json:"protectKernelDefaults"`
	ProviderID                      *string                          `json:"providerID,omitempty"`
	ReadOnlyPort                    int                              `json:"readOnlyPort"`
	RegisterWithTaints              []v1.Taint                       `json:"registerWithTaints,omitempty"`
	ResolverConfig                  *string                          `json:"resolvConf,omitempty"`
	SerializeImagePulls             bool                             `json:"serializeImagePulls"`
	ServerTLSBootstrap              bool                             `json:"serverTLSBootstrap"`
	ShutdownGracePeriod             *metav1.Duration                 `json:"shutdownGracePeriod,omitempty"`
	ShutdownGracePeriodCriticalPods *metav1.Duration                 `json:"shutdownGracePeriodCriticalPods,omitempty"`
	SystemReservedCgroup            *string                          `json:"systemReservedCgroup,omitempty"`
	TLSCipherSuites                 []string                         `json:"tlsCipherSuites"`
	metav1.TypeMeta                 `json:",inline"`
}

type loggingConfiguration struct {
	Verbosity int `json:"verbosity"`
}

// Creates an internal kubelet configuration from the public facing bootstrap
// kubelet configuration with additional sane defaults.
func defaultKubeletSubConfig() kubeletConfig {
	return kubeletConfig{
		TypeMeta: metav1.TypeMeta{
			Kind:       "KubeletConfiguration",
			APIVersion: "kubelet.config.k8s.io/v1beta1",
		},
		Address: "0.0.0.0",
		Authentication: k8skubelet.KubeletAuthentication{
			Anonymous: k8skubelet.KubeletAnonymousAuthentication{
				Enabled: ptr.Bool(false),
			},
			Webhook: k8skubelet.KubeletWebhookAuthentication{
				Enabled:  ptr.Bool(true),
				CacheTTL: metav1.Duration{Duration: time.Minute * 2},
			},
			X509: k8skubelet.KubeletX509Authentication{
				ClientCAFile: caCertificatePath,
			},
		},
		Authorization: k8skubelet.KubeletAuthorization{
			Mode: "Webhook",
			Webhook: k8skubelet.KubeletWebhookAuthorization{
				CacheAuthorizedTTL:   metav1.Duration{Duration: time.Minute * 5},
				CacheUnauthorizedTTL: metav1.Duration{Duration: time.Second * 30},
			},
		},
		CgroupDriver:             "systemd",
		CgroupRoot:               "/",
		ClusterDomain:            "cluster.local",
		ContainerRuntimeEndpoint: containerd.ContainerRuntimeEndpoint,
		EvictionHard: map[string]string{
			"memory.available":  "100Mi",
			"nodefs.available":  "10%",
			"nodefs.inodesFree": "5%",
		},
		FeatureGates: map[string]bool{
			"RotateKubeletServerCertificate": true,
		},
		HairpinMode:           "hairpin-veth",
		ProtectKernelDefaults: true,
		ReadOnlyPort:          0,
		Logging: loggingConfiguration{
			Verbosity: 2,
		},
		SerializeImagePulls: false,
		ServerTLSBootstrap:  true,
		TLSCipherSuites: []string{
			"TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256",
			"TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384",
			"TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305",
			"TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256",
			"TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384",
			"TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305",
			"TLS_RSA_WITH_AES_128_GCM_SHA256",
			"TLS_RSA_WITH_AES_256_GCM_SHA384",
		},
	}
}

// Update the ClusterDNS of the internal kubelet config using a heuristic based
// on the cluster service IP CIDR address.
func (ksc *kubeletConfig) withFallbackClusterDns(cluster *api.ClusterDetails) error {
	clusterDns, err := cluster.GetClusterDns()
	if err != nil {
		return err
	}
	ksc.ClusterDNS = []string{clusterDns}
	return nil
}

// On Amazon Linux 2027, point kubelet at systemd-resolved's uplink resolv.conf
// so pods never inherit the 127.0.0.53 stub address. Inside a pod's network
// namespace that address is the pod's own loopback, which makes CoreDNS forward
// to itself and exit. This is the fix upstream documents for systemd-resolved:
// https://kubernetes.io/docs/tasks/administer-cluster/dns-debugging-resolution/#known-issues
// A resolvConf in the user's kubelet config (the 40-nodeadm.conf drop-in) or a
// --resolv-conf flag still takes precedence.
func (ksc *kubeletConfig) withResolverConfig(osDistro api.OSDistro) {
	if osDistro == api.OSDistroAL2027 {
		ksc.ResolverConfig = ptr.String(systemdResolvedUplinkResolvConf)
	}
}

// To support worker nodes to continue to communicate and connect to local cluster even when the Outpost
// is disconnected from the parent AWS Region, the following specific setup are required:
//   - append entries to /etc/hosts with the mappings of control plane host IP address and API server
//     domain name. So that the domain name can be resolved to IP addresses locally.
//   - use aws-iam-authenticator as bootstrap auth for kubelet TLS bootstrapping which downloads client
//     X.509 certificate and generate kubelet kubeconfig file which uses the client cert. So that the
//     worker node can be authentiacated through X.509 certificate which works for both connected and
//     disconnected state.
func (ksc *kubeletConfig) withOutpostSetup(cfg *api.NodeConfig) error {
	if enabled := cfg.Spec.Cluster.EnableOutpost; enabled != nil && *enabled {
		zap.L().Info("Setting up outpost..")

		if cfg.Spec.Cluster.ID == "" {
			return fmt.Errorf("clusterId cannot be empty when outpost is enabled")
		}
		apiUrl, err := url.Parse(cfg.Spec.Cluster.APIServerEndpoint)
		if err != nil {
			return err
		}

		// TODO: cleanup
		var ipAddresses []string
		err = retry.New(
			retry.Attempts(6),
			retry.Delay(200*time.Millisecond),
			retry.OnRetry(func(n uint, err error) {
				zap.L().Info("Retrying DNS lookup after error", zap.Error(err))
			}),
		).Do(
			func() error {
				var err error
				ipAddresses, err = net.LookupHost(apiUrl.Host)
				return err
			},
		)
		if err != nil {
			return err
		}
		var ipHostMappings []string
		for _, ip := range ipAddresses {
			ipHostMappings = append(ipHostMappings, fmt.Sprintf("%s\t%s", ip, apiUrl.Host))
		}
		output := strings.Join(ipHostMappings, "\n") + "\n"
		// append to /etc/hosts file with shuffled mappings of "IP address to API server domain name"
		f, err := os.OpenFile("/etc/hosts", os.O_APPEND|os.O_WRONLY, kubeletConfigPerm)
		if err != nil {
			return err
		}
		defer f.Close()
		if _, err := f.WriteString(output); err != nil {
			return err
		}
	}
	return nil
}

func (ksc *kubeletConfig) withNodeLabels(flags map[string]string, nodeLabelFuncs map[string]LabelProvider) {
	var nodeLabels []string
	for nodeLabelKey, provider := range nodeLabelFuncs {
		nodeLabelValue, ok, err := provider.Get()
		if err != nil {
			zap.L().Error("Failed to get node label value", zap.String("key", nodeLabelKey), zap.Error(err))
			continue
		}
		if !ok {
			continue
		}
		nodeLabel := fmt.Sprintf("%s=%s", nodeLabelKey, nodeLabelValue)
		zap.L().Info("Adding node label", zap.String("label", nodeLabel))
		nodeLabels = append(nodeLabels, nodeLabel)
	}
	if len(nodeLabels) > 0 {
		flags["node-labels"] = strings.Join(nodeLabels, ",")
	}
}

func (ksc *kubeletConfig) withNodeIp(cfg *api.NodeConfig, flags map[string]string, imdsClient imds.IMDSClient) error {
	nodeIp, err := getNodeIp(context.TODO(), cfg, imdsClient)
	if err != nil {
		return err
	}
	flags["node-ip"] = nodeIp
	zap.L().Info("Setup IP for node", zap.String("ip", nodeIp))
	return nil
}

func (ksc *kubeletConfig) withVersionToggles(cfg *api.NodeConfig) {
	// EKS enables DRA on 1.33+
	if semver.Compare(cfg.Status.KubeletVersion, "v1.33.0") >= 0 {
		ksc.FeatureGates["DynamicResourceAllocation"] = true
	}
	// Enable MutableCSINodeAllocatableCount on 1.34+
	if semver.Compare(cfg.Status.KubeletVersion, "v1.34.0") >= 0 {
		ksc.FeatureGates["MutableCSINodeAllocatableCount"] = true
	}
	// Enable graceful node shutdown in 1.34+
	if semver.Compare(cfg.Status.KubeletVersion, "v1.34.0") >= 0 {
		ksc.ShutdownGracePeriod = &metav1.Duration{
			Duration: time.Second * 150, // 2m30s
		}
		ksc.ShutdownGracePeriodCriticalPods = &metav1.Duration{
			Duration: time.Second * 30, // allocated from the total above
		}
	}
}

func (ksc *kubeletConfig) withCloudProvider(cfg *api.NodeConfig, flags map[string]string) {
	// ref: https://github.com/kubernetes/kubernetes/pull/121367
	flags["cloud-provider"] = "external"
	// provider ID needs to be specified when the cloud provider is external
	ksc.ProviderID = ptr.String(getProviderId(cfg.Status.Instance.AvailabilityZone, cfg.Status.Instance.ID))
	var nodeName string
	if api.UseInstanceIdNodeName(cfg.Spec.FeatureGates, cfg.Status.OSDistro) {
		zap.L().Info("Using Instance Id naming strategy")
		nodeName = cfg.Status.Instance.ID
	} else {
		// the name of the Node object default to EC2 PrivateDnsName
		// see: https://github.com/awslabs/amazon-eks-ami/pull/1264
		nodeName = cfg.Status.Instance.PrivateDNSName
	}
	flags["hostname-override"] = nodeName
}

// Override the kubelet config with reserved cgroup values on behalf of the user
func (ksc *kubeletConfig) withDefaultReservedResources(cfg *api.NodeConfig, resources system.Resources) {
	ksc.SystemReservedCgroup = ptr.String("/system")
	ksc.KubeReservedCgroup = ptr.String("/runtime")
	if instanceInfo, err := GetInstanceInfo(context.TODO(), cfg.Status.Instance.Region, cfg.Status.Instance.Type); err != nil {
		zap.L().Warn("Failed to retrieve instance info, falling back to default", zap.Error(err))
		ksc.MaxPods = defaultMaxPods
	} else {
		ksc.MaxPods = CalcMaxPods(instanceInfo, cfg.Spec.Kubelet.MaxPodsExpression)
	}
	ksc.KubeReserved = map[string]string{
		"cpu":               fmt.Sprintf("%dm", getCPUMillicoresToReserve(resources)),
		"ephemeral-storage": "1Gi",
		"memory":            fmt.Sprintf("%dMi", getMemoryMebibytesToReserve(ksc.MaxPods)),
	}
}

func (ksc *kubeletConfig) withImageServiceEndpoint(cfg *api.NodeConfig, resources system.Resources) {
	if containerd.UseSOCISnapshotter(cfg, resources) {
		ksc.ImageServiceEndpoint = "unix:///run/soci-snapshotter-grpc/soci-snapshotter-grpc.sock"
	}
}

func (ksc *kubeletConfig) withRuntimeCgroups(flags map[string]string) {
	// Set runtime cgroups so that cadvisor metrics include container usage.
	// Reference https://github.com/awslabs/amazon-eks-ami/issues/1667
	flags["runtime-cgroups"] = "/runtime.slice/containerd.service"
}

func (k *kubelet) generateKubeletConfig(cfg *api.NodeConfig) (*kubeletConfig, error) {
	kubeletConfig := defaultKubeletSubConfig()

	if err := kubeletConfig.withFallbackClusterDns(&cfg.Spec.Cluster); err != nil {
		return nil, err
	}
	if err := kubeletConfig.withOutpostSetup(cfg); err != nil {
		return nil, err
	}
	if err := kubeletConfig.withNodeIp(cfg, k.flags, k.imdsClient); err != nil {
		return nil, err
	}

	kubeletConfig.withVersionToggles(cfg)
	kubeletConfig.withCloudProvider(cfg, k.flags)
	kubeletConfig.withDefaultReservedResources(cfg, k.resources)
	kubeletConfig.withImageServiceEndpoint(cfg, k.resources)
	kubeletConfig.withRuntimeCgroups(k.flags)
	kubeletConfig.withResolverConfig(cfg.Status.OSDistro)

	nodeLabelFuncs := map[string]LabelProvider{}
	if semver.Compare(cfg.Status.KubeletVersion, "v1.35.0") >= 0 {
		// see: https://github.com/NVIDIA/gpu-operator/commit/e25291b86cf4542ac62d8635cda4bd653c4face3
		nodeLabelFuncs["nvidia.com/gpu.present"] = NvidiaGPULabel{fs: system.RealFileSystem{}}
	}
	kubeletConfig.withNodeLabels(k.flags, nodeLabelFuncs)

	return &kubeletConfig, nil
}

// writeKubeletConfig writes nodeadm's generated kubelet config to the
// standard config file and writes the user's provided config to a directory for
// drop-in support.
func (k *kubelet) writeKubeletConfig(cfg *api.NodeConfig) error {
	kubeletConfig, err := k.generateKubeletConfig(cfg)
	if err != nil {
		return err
	}
	kubeletConfigBytes, err := json.MarshalIndent(kubeletConfig, "", strings.Repeat(" ", 4))
	if err != nil {
		return err
	}

	configPath := path.Join(kubeletConfigRoot, kubeletConfigFile)
	k.flags["config"] = configPath

	zap.L().Info("Writing kubelet config to file..", zap.String("path", configPath))
	if err := util.WriteFileWithDir(configPath, kubeletConfigBytes, kubeletConfigPerm); err != nil {
		return err
	}

	if len(cfg.Spec.Kubelet.Config) > 0 {
		dirPath := path.Join(kubeletConfigRoot, kubeletConfigDir)
		k.flags["config-dir"] = dirPath
		if semver.Compare(cfg.Status.KubeletVersion, "v1.30.0") < 0 {
			zap.L().Info("Enabling kubelet config drop-in dir..")
			k.environment["KUBELET_CONFIG_DROPIN_DIR_ALPHA"] = "on"
		}

		// merge in default type metadata like kind and apiVersion in case the
		// user has not specified this, as it is required to qualify a drop-in
		// config as a valid KubeletConfiguration
		userKubeletConfigMap, err := util.Merge(defaultKubeletSubConfig().TypeMeta, cfg.Spec.Kubelet.Config, json.Marshal, json.Unmarshal)
		if err != nil {
			return err
		}
		userKubeletConfigBytes, err := json.MarshalIndent(userKubeletConfigMap, "", strings.Repeat(" ", 4))
		if err != nil {
			return err
		}
		filePath := path.Join(dirPath, "40-nodeadm.conf")
		zap.L().Info("Writing user kubelet config to drop-in file..", zap.String("path", filePath))
		if err := util.WriteFileWithDir(filePath, userKubeletConfigBytes, kubeletConfigPerm); err != nil {
			return err
		}
	}

	return nil
}

func getProviderId(availabilityZone, instanceId string) string {
	return fmt.Sprintf("aws:///%s/%s", availabilityZone, instanceId)
}

// Get the IP of the node depending on the ipFamily configured for the cluster
func getNodeIp(ctx context.Context, cfg *api.NodeConfig, imdsClient imds.IMDSClient) (string, error) {
	ipFamily, err := api.GetCIDRIpFamily(cfg.Spec.Cluster.CIDR)
	if err != nil {
		return "", err
	}
	switch ipFamily {
	case api.IPFamilyIPv4:
		ipv4, err := imdsClient.GetProperty(ctx, imds.LocalIPv4)
		if err != nil {
			return "", err
		}
		return ipv4, nil
	case api.IPFamilyIPv6:
		prop := fmt.Sprintf("network/interfaces/macs/%s/ipv6s", cfg.Status.Instance.MAC)
		ipv6, err := imdsClient.GetProperty(ctx, imds.IMDSProperty(prop))
		if err != nil {
			return "", err
		}
		return ipv6, nil
	default:
		return "", fmt.Errorf("invalid ip-family. %s is not one of %v", ipFamily, []api.IPFamily{api.IPFamilyIPv4, api.IPFamilyIPv6})
	}
}

func getCPUMillicoresToReserve(resources system.Resources) int {
	totalCPUMillicores, err := resources.GetMilliNumCores()
	if err != nil {
		zap.L().Error("Error found when GetMilliNumCores", zap.Error(err))
		return 0
	}
	cpuRanges := []int{0, 1000, 2000, 4000, totalCPUMillicores}
	cpuPercentageReservedForRanges := []int{600, 100, 50, 25}
	cpuToReserve := 0

	for i, percentageToReserveForRange := range cpuPercentageReservedForRanges {
		startRange := cpuRanges[i]
		// #nosec G602 // cpuRanges is one item longer than cpuPercentageReservedForRanges
		endRange := cpuRanges[i+1]
		cpuToReserve += getResourceToReserveInRange(totalCPUMillicores, startRange, endRange, percentageToReserveForRange)
	}

	return cpuToReserve
}

// getResourceToReserveInRange calculates the CPU resources to reserve for a given range.
func getResourceToReserveInRange(totalCPU, startRange, endRange, percentage int) int {
	if totalCPU <= startRange {
		return 0
	}
	reserved := totalCPU
	if reserved > endRange {
		reserved = endRange
	}
	return (reserved - startRange) * percentage / 10000
}

func getMemoryMebibytesToReserve(maxPods int32) int32 {
	return 11*maxPods + 255
}