package main

import (
	"context"
	"flag"
	"log"
	"slices"
	"strings"
	"time"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

type args struct {
	ExcludedNS *string
	Memory     *string
	_          struct{}
}

func main() {
	args := parseArgs()

	ctx := context.Background()
	config, err := rest.InClusterConfig()
	if err != nil {
		panic(err.Error())
	}
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		panic(err.Error())
	}

	memory, err := resource.ParseQuantity(*args.Memory)
	if err != nil {
		panic(err.Error())
	}
	log.Printf("Memory: %v", memory.String())

	limits := []v1.LimitRangeItem{
		{
			DefaultRequest: v1.ResourceList{"memory": memory},
			Type:           "Container",
		},
	}
	limitRange := v1.LimitRange{
		Spec: v1.LimitRangeSpec{Limits: limits},
		ObjectMeta: metav1.ObjectMeta{
			Name: "extreme-request-defaults",
		},
	}

	excludedNS := strings.Split(*args.ExcludedNS, ",")

	for {
		namespaces, err := clientset.CoreV1().Namespaces().List(ctx, metav1.ListOptions{})
		if err != nil {
			panic(err)
		}

		for _, ns := range namespaces.Items {
			log.Printf("Checking for LimitRange named extreme-request-defaults in namespace '%v'\n", ns.Name)
			limitRanges, err := clientset.CoreV1().LimitRanges(ns.Name).List(ctx, metav1.ListOptions{FieldSelector: "metadata.name=extreme-request-defaults"})
			if err != nil {
				log.Printf("Unable to list LimitRanges in namespace '%v': Error: %v\n", ns.Name, err)
				continue
			}
			if !nsExcluded(ns.Name, excludedNS) {
				if len(limitRanges.Items) == 0 {
					log.Printf("Trying to create LimitRange\n")
					if _, err := clientset.CoreV1().LimitRanges(ns.Name).Create(ctx, &limitRange, metav1.CreateOptions{}); err != nil {
						log.Printf("Unable to create LimitRange in namespace '%v': Error: %v\n", ns.Name, err)
					} else {
						log.Printf("LimitRange extreme-request-defaults created in namespace '%v'\n", ns.Name)
					}
				}
			} else {
				if len(limitRanges.Items) > 0 {
					log.Printf("Trying to delete LimitRange\n")
					if err := clientset.CoreV1().LimitRanges(ns.Name).Delete(ctx, "extreme-request-defaults", metav1.DeleteOptions{}); err != nil {
						log.Printf("Unable to delete LimitRange in namespace '%v': Error: %v\n", ns.Name, err)
					}
				}
			}
		}

		time.Sleep(10 * time.Second)
	}
}

func nsExcluded(name string, excludedNS []string) bool {
	return slices.Contains(excludedNS, name) || slices.Contains(excludedNS, "*")
}

func parseArgs() args {
	args := args{
		ExcludedNS: flag.String("excluded-ns", "kube-system", "Comma-separated list of namespaces to be excluded"),
		Memory:     flag.String("memory", "1Ti", "The default memory requests to set in the LimitRange. Default: 1Ti"),
	}
	flag.Parse()

	return args
}