package main

import (
	"flag"
	"fmt"
	"math"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/tidbcloud/aws-shared-provider/pkg/vpa"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/informers"
	v1 "k8s.io/client-go/informers/core/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"
)

var (
	addr                         string
	memoryCapacity               string
	memoryUseRatio               float64
	kubeConfig                   string
	nodeName                     string
	maxReserved                  string
	enableMonitorRebalanceSignal bool
	enableKillPod                bool
	measureRealtimeMemUsage      bool
)

func main() {
	flag.StringVar(&addr, "addr", ":4040", "listening address")
	flag.StringVar(&memoryCapacity, "memoryCapacity", "", "Node memory capacity, autoload if empty")
	flag.Float64Var(&memoryUseRatio, "memoryUseRatio", 0.8, "Maximum node memory ratio")
	flag.StringVar(&kubeConfig, "kubeConfig", "", "kubeconfig")
	flag.StringVar(&nodeName, "nodeName", "", "node name")
	flag.StringVar(&maxReserved, "maxReserved", "3Gi", "Max Reserved memory for OS")
	flag.BoolVar(&enableMonitorRebalanceSignal, "enableMonitorRebalanceSignal", true, "Enable monitor rebalance signal")
	flag.BoolVar(&enableKillPod, "enableKillPod", true, "Enable kill pod")
	flag.BoolVar(&measureRealtimeMemUsage, "measureRealtimeMemUsage", false, "Enable measure realtime memory usage instead of using reported allocated memory, default is false")

	flag.Parse()

	opt := zap.Options{}
	opt.BindFlags(flag.CommandLine)
	ctrl.SetLogger(zap.New(zap.UseFlagOptions(&opt)))

	config, err := clientcmd.BuildConfigFromFlags("", kubeConfig)
	if err != nil {
		ctrl.Log.Error(err, "build kubeConfig error", "kubeConfig", kubeConfig)
		os.Exit(1)
	}
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		ctrl.Log.Error(err, "new clientset error", "kubeConfig", kubeConfig)
		os.Exit(1)
	}

	nodeFactory := informers.NewSharedInformerFactoryWithOptions(
		clientset, 10*time.Minute, informers.WithTweakListOptions(func(opts *metav1.ListOptions) {
			opts.FieldSelector = fmt.Sprintf("metadata.name=%s", nodeName)
		}))
	nodeInformer := nodeFactory.Core().V1().Nodes()
	nodeInformer.Informer()
	stopCh := make(chan struct{})
	nodeFactory.Start(stopCh)
	nodeFactory.WaitForCacheSync(stopCh)

	var q resource.Quantity

	if memoryCapacity != "" {
		q, err = resource.ParseQuantity(memoryCapacity)
		if err != nil {
			ctrl.Log.Error(err, "failed to parse memory capacity", "memoryCapacity", memoryCapacity)
			os.Exit(1)
		}
	} else {
		memory, err := loadNodeMemory(nodeInformer)
		if err != nil {
			ctrl.Log.Error(err, "failed to load node memory")
			os.Exit(1)
		}
		q = *memory
	}

	configuredReserved, _ := parseGiToBytes(maxReserved)
	maxUse := evalMaxUsableMemory(q.Value(), configuredReserved, memoryUseRatio)

	mc, err := vpa.NewMemoryAllocator(maxUse, enableKillPod, measureRealtimeMemUsage, clientset, nodeInformer, nodeName)
	if err != nil {
		ctrl.Log.Error(err, "failed to create memory controller")
		os.Exit(1)
	}

	if enableMonitorRebalanceSignal {
		err = vpa.StartMonitorRebalanceSignal(clientset, nodeInformer, nodeName)
		if err != nil {
			ctrl.Log.Error(err, "failed to start monitor rebalance signal")
			os.Exit(1)
		}
	}

	ctrl.Log.Info("memory usage capacity",
		"memoryCapacity", q.String(),
		"memoryUseRatio", memoryUseRatio,
		"maxUse", maxUse)

	mux := http.NewServeMux()
	mux.HandleFunc("/report", mc.HandleReport)
	mux.HandleFunc("/scale", mc.HandleScale)
	mux.HandleFunc("/status", mc.HandleStatus)
	mux.Handle("/metrics", promhttp.Handler())
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {})
	mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) {
		if mc.Ready() {
			w.WriteHeader(http.StatusOK)
		} else {
			w.WriteHeader(http.StatusServiceUnavailable)
		}
	})
	http.ListenAndServe(addr, mux) //nolint:errcheck // ignore error
}

func loadNodeMemory(nodeInformer v1.NodeInformer) (*resource.Quantity, error) {
	node, err := nodeInformer.Lister().Get(nodeName)
	if err != nil {
		return nil, err
	}
	return node.Status.Capacity.Memory(), nil
}

func parseGiToBytes(giga string) (int64, error) {
	if strings.HasSuffix(giga, "Gi") {
		numberStr := strings.TrimSuffix(giga, "Gi")
		value, err := strconv.Atoi(numberStr)
		if err != nil || value <= 0 {
			return math.MaxInt64, fmt.Errorf("MaxReserved is missconfigured")
		}
		return int64(value) * 1024 * 1024 * 1024, nil //nolint:mnd // 1024 is good
	} else {
		return math.MaxInt64, fmt.Errorf("MaxReserved is missconfigured")
	}
}

func evalMaxUsableMemory(nodeCapacity, configuredReserved int64, maxMemRatio float64) int64 {
	if int64(float64(nodeCapacity)*(1-maxMemRatio)) > configuredReserved {
		return nodeCapacity - configuredReserved
	}
	return int64(float64(nodeCapacity) * maxMemRatio)
}