/*
Copyright 2023.

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 controller

import (
	"context"
	"fmt"
	"math"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/pingcap/errors"
	o11yv1 "github.com/tidbcloud/observability/api/regionalserver/resourceworker/v1"
	"github.com/tidbcloud/observability/pkg/common/constants"
	pkgconfig "github.com/tidbcloud/observability/vm-bench/pkg/config"
	"github.com/tidbcloud/observability/vm-bench/pkg/helm"
	"k8s.io/apimachinery/pkg/runtime"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	"sigs.k8s.io/controller-runtime/pkg/log"
	"sigs.k8s.io/controller-runtime/pkg/predicate"
	"sigs.k8s.io/controller-runtime/pkg/source"
)

const (
	benchFinalizer      = "bench.finalizers.observability.pingcap.com"
	proxiedVMInsertTmpl = "http://ingress-nginx-a-controller.kube-system.svc.cluster.local:80/%s/api/v1/write"
	// this will generate about 4K churn rate to vmstorage when a new cluster is created.
	targetMapTmpl  = "0.0.0.0:9102=5,basic-tidb-%d.basic-tidb-peer.tidb-admin:10080=1,basic-tikv-%d.basic-tikv-peer.tidb-admin:20180=1,basic-pd-%d.basic-pd-peer.tidb-admin:2379=3"
	largeTargetMap = "0.0.0.0:9102=100,basic-tidb-0.basic-tidb-peer.tidb-admin:10080=18,basic-tidb-1.basic-tidb-peer.tidb-admin:10080=18,basic-tikv-0.basic-tikv-peer.tidb-admin:20180=18,basic-tikv-1.basic-tikv-peer.tidb-admin:20180=18,basic-tikv-2.basic-tikv-peer.tidb-admin:20180=18,basic-pd-0.basic-pd-peer.tidb-admin:2379=3,basic-pd-1.basic-pd-peer.tidb-admin:2379=3,basic-pd-2.basic-pd-peer.tidb-admin:2379=4"
	largeCountBase = 100
	largeCountCeil = 400
	largeCountUnit = 20
	nodeAddr       = "0.0.0.0"
)

var (
	helmLogger = ctrl.Log.WithName("helm")
)

// ClusterReconciler reconciles a Cluster object
type ClusterReconciler struct {
	client.Client
	ClientGetter helm.ClientGetter
	Config       *pkgconfig.RootConfig
	Scheme       *runtime.Scheme
}

//+kubebuilder:rbac:groups=observability.pingcap.com,resources=clusters,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=observability.pingcap.com,resources=clusters/status,verbs=get;update;patch
//+kubebuilder:rbac:groups=observability.pingcap.com,resources=clusters/finalizers,verbs=update

// Reconcile is part of the main kubernetes reconciliation loop which aims to
// move the current state of the cluster closer to the desired state.
// TODO(user): Modify the Reconcile function to compare the state specified by
// the Cluster object against the actual cluster state, and then
// perform operations to make the cluster state reflect the state specified by
// the user.
//
// For more details, check Reconcile and its Result here:
// - https://pkg.go.dev/sigs.k8s.io/controller-runtime@v0.15.0/pkg/reconcile
func (r *ClusterReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	_ = log.FromContext(ctx)
	cluster := &o11yv1.Cluster{}

	if err := r.Get(ctx, req.NamespacedName, cluster); err != nil {
		if client.IgnoreNotFound(err) == nil {
			return ctrl.Result{}, nil
		}
		return ctrl.Result{}, fmt.Errorf("unable to get cluster %s: %v", req.NamespacedName, err)
	}

	// examine DeletionTimestamp to determine if object is under deletion
	if cluster.ObjectMeta.DeletionTimestamp.IsZero() {
		if !controllerutil.ContainsFinalizer(cluster, benchFinalizer) {
			controllerutil.AddFinalizer(cluster, benchFinalizer)
			if err := r.Update(ctx, cluster); err != nil {
				return ctrl.Result{}, err
			}
		}
	} else {
		if controllerutil.ContainsFinalizer(cluster, benchFinalizer) {
			if err := r.deleteExternalResources(ctx, cluster); err != nil {
				return ctrl.Result{}, err
			}

			// remove our finalizer from the list and update it.
			controllerutil.RemoveFinalizer(cluster, benchFinalizer)
			if err := r.Update(ctx, cluster); err != nil {
				return ctrl.Result{}, err
			}
		}
		return ctrl.Result{}, nil
	}
	if cluster.Status.Phase == "ACTIVE" {
		requeue, values, err := r.getValues(ctx, *cluster)
		if err != nil {
			return ctrl.Result{}, errors.Wrap(err, "failed to apply vm agent")
		}
		err = helm.ApplyV2(ctx, r.ClientGetter, &r.Config.VMAgent, releaseName(*cluster), cluster.Namespace, values)
		if err != nil {
			return ctrl.Result{}, errors.Wrap(err, "failed to apply vm agent")
		}
		registerStressQuerier(ctx, cluster)
		if requeue {
			return ctrl.Result{RequeueAfter: 60 * time.Second}, nil
		}
	} else {
		// wait to be ACTIVE
		return ctrl.Result{RequeueAfter: time.Second}, nil
	}

	return ctrl.Result{}, nil
}

func (r *ClusterReconciler) getValues(ctx context.Context, cluster o11yv1.Cluster) (bool, map[string]interface{}, error) {
	values, err := helm.GetReleaseValues(ctx, r.ClientGetter, releaseName(cluster))
	var targetAddrMap string
	var requeue bool
	if cluster.Labels[constants.LabelApplicationID] == "large" {
		if err != nil || values["targetAddrMap"] == nil {
			targetAddrMap = largeTargetMap
		} else {
			targetAddrMap = values["targetAddrMap"].(string)
		}
		return r.largeValues(ctx, targetAddrMap, cluster)
	}

	if err != nil || values["targetAddrMap"] == nil {
		rand := len(queriers) % 3
		tidbRand := len(queriers) % 2
		targetAddrMap = fmt.Sprintf(targetMapTmpl, tidbRand, rand, rand)
		requeue = true
	} else {
		// reuse existing targetAddrMap to avoid too much series churn.
		updated := updateAddrCount(values["targetAddrMap"].(string))
		targetAddrMap = updated.String()
	}
	return requeue, map[string]interface{}{
		"nginx": map[string]interface{}{
			"enabled": true,
		},
		"targetAddrMap": targetAddrMap,
		"remoteStorages": map[string]interface{}{
			"vm": map[string]interface{}{
				"writeURL": fmt.Sprintf(proxiedVMInsertTmpl, cluster.Name),
			},
		},
		"scrapeInterval": "15s",
		"externalLabels": []string{
			"cluster=stag-benchmark-gcp-us-west1-" + releaseName(cluster),
			"shoot_name=stag-benchmark-gcp-us-west1-" + releaseName(cluster),
			"project=" + cluster.Labels[constants.LabelProjectID],
			"seed_region=us-west1",
			"seed_provider=gcp",
			"shoot_infra=gke",
			"tenant=" + cluster.Labels[constants.LabelTenantID],
			"cluster_id=" + cluster.Labels[constants.LabelApplicationID],
			"tidb_cluster_id=large",
		},
	}, nil
}

func (r *ClusterReconciler) largeValues(ctx context.Context, targetAddrMap string, cluster o11yv1.Cluster) (bool, map[string]interface{}, error) {
	var requeue bool
	target := parseTargetPairs(targetAddrMap)
	cs := &o11yv1.ClusterList{}
	err := r.List(ctx, cs, client.InNamespace(cluster.Namespace))
	if err != nil {
		return true, nil, fmt.Errorf("failed to list clusters: %v", err)
	}
	clusterCount := len(cs.Items) - 1
	expected := clusterCount * largeCountUnit
	if expected > largeCountCeil {
		expected = largeCountCeil
	}
	if target.nodeCount < expected {
		requeue = true
		targetAddrMap = smoothIncrease(target)
	}
	if target.nodeCount > expected+largeCountUnit && target.nodeCount > largeCountBase {
		delta := target.nodeCount - expected
		if expected < largeCountBase {
			delta = target.nodeCount - largeCountBase
		}
		targetAddrMap = decrease(target, delta)
	}

	return requeue, map[string]interface{}{
		"nginx": map[string]interface{}{
			"enabled": true,
		},
		"nodeSelector": map[string]interface{}{
			"topology.kubernetes.io/zone": "us-west1-a",
			"node_type":                   "benchmark",
		},
		// for large Cluster with nodes to 200,
		// then around 400 * 0.25% = 1 targets gets updated labels on every scrape config update.
		// This generates around 10k*1=10k new time series every scrapeConfigUpdateInterval.
		"scrapeConfigUpdatePercent":  0.25,
		"scrapeConfigUpdateInterval": "30m",
		"targetAddrMap":              targetAddrMap,
		"remoteStorages": map[string]interface{}{
			"vm": map[string]interface{}{
				"writeURL": fmt.Sprintf(proxiedVMInsertTmpl, cluster.Name),
			},
		},
		"writeConcurrency": 8,
		"writeReplicaResource": map[string]interface{}{
			"limits": map[string]interface{}{
				"cpu":    "4",
				"memory": "2Gi",
			},
			"requests": map[string]interface{}{
				"cpu":    "1",
				"memory": "500Mi",
			},
		},
		"scrapeInterval": r.Config.VMAgent.ScrapeInterval,
		"externalLabels": []string{
			"cluster=stag-benchmark-gcp-us-west1-" + releaseName(cluster),
			"shoot_name=stag-benchmark-gcp-us-west1-" + releaseName(cluster),
			"project=" + cluster.Labels[constants.LabelProjectID],
			"seed_region=us-west1",
			"seed_provider=gcp",
			"shoot_infra=gke",
			"tenant=" + cluster.Labels[constants.LabelTenantID],
			"cluster_id=" + cluster.Labels[constants.LabelApplicationID],
			"tidb_cluster_id=large",
		},
	}, nil
}

func (r *ClusterReconciler) deleteExternalResources(ctx context.Context, cluster *o11yv1.Cluster) error {
	err := helm.Uninstall(ctx, r.ClientGetter, &r.Config.VMAgent, releaseName(*cluster))
	if err != nil {
		return errors.Annotate(err, "uninstall vmagent chart failed")
	}
	unregisterStressQuerier(*cluster)
	return nil
}

// SetupWithManager sets up the controller with the Manager.
func (r *ClusterReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).
		For(&o11yv1.Cluster{}).
		Watches(&source.Kind{Type: &o11yv1.Cluster{}}, handler.EnqueueRequestsFromMapFunc(r.addtionalEnqueueLargeCluster)).
		WithEventFilter(predicate.ResourceVersionChangedPredicate{}).
		Complete(r)
}

func (r *ClusterReconciler) addtionalEnqueueLargeCluster(cluster client.Object) []ctrl.Request {
	results := []ctrl.Request{}
	if cluster.GetLabels()[constants.LabelApplicationID] == "large" {
		return results
	}
	largeClusterList := &o11yv1.ClusterList{}
	err := r.List(context.TODO(), largeClusterList, client.MatchingLabels{constants.LabelApplicationID: "large"})
	if err != nil || len(largeClusterList.Items) == 0 {
		logger.Error(err, "failed to list large cluster when enqueueing")
	} else {
		results = append(results, ctrl.Request{
			NamespacedName: client.ObjectKeyFromObject(&largeClusterList.Items[0]),
		})
	}
	return results
}

func releaseName(cluster o11yv1.Cluster) string {
	if cluster.Labels[constants.LabelApplicationID] == "large" {
		return "large"
	}
	return cluster.Name[0:8]
}

func parseTargetPairs(s string) *TargetAddrPair {
	m := make(map[string]int)
	pairs := strings.Split(s, ",")
	total := 0
	for _, p := range pairs {
		kv := strings.Split(p, "=")
		if len(kv) == 2 {
			count, err := strconv.ParseInt(kv[1], 10, 0)
			if err != nil {
				panic(fmt.Errorf("failed to parse target addr map: %s", s))
			}
			m[strings.Trim(kv[0], " ")] = int(count)
			if !strings.Contains(kv[0], nodeAddr) {
				total += int(count)
			}
		}
	}
	return &TargetAddrPair{
		TargetAddrMap: m,
		nodeCount:     total,
	}
}

type TargetAddrPair struct {
	TargetAddrMap map[string]int
	nodeCount     int
}

func updateAddrCount(targetAddrMap string) *TargetAddrPair {
	pair := parseTargetPairs(targetAddrMap)
	for k := range pair.TargetAddrMap {
		switch {
		case strings.Contains(k, "tikv-peer"):
			pair.TargetAddrMap[k] = 7
		case strings.Contains(k, "pd-peer"):
			pair.TargetAddrMap[k] = 3
		case strings.Contains(k, "tidb-peer"):
			pair.TargetAddrMap[k] = 5
		default:
		}
	}
	pair.nodeCount = 15
	return pair
}

func smoothIncrease(pair *TargetAddrPair) string {
	for k := range pair.TargetAddrMap {
		switch {
		case strings.Contains(k, "tikv-peer"):
			pair.TargetAddrMap[k] += 1
			pair.nodeCount += 1
		case strings.Contains(k, "tidb-peer"):
			pair.TargetAddrMap[k] += 1
			pair.nodeCount += 1
		default:
		}
	}
	return pair.String()
}

func decrease(pair *TargetAddrPair, count int) string {
	d := int(math.Round(float64(count) / 5.0))
	if d == 0 {
		d = 1
	}
	for k := range pair.TargetAddrMap {
		switch {
		case strings.Contains(k, "tikv-peer"):
			pair.TargetAddrMap[k] -= d
			pair.nodeCount -= d
		case strings.Contains(k, "tidb-peer"):
			pair.TargetAddrMap[k] -= d
			pair.nodeCount -= d
		default:
		}
	}
	return pair.String()
}

func (t *TargetAddrPair) String() string {
	sb := strings.Builder{}
	keys := make([]string, 0, len(t.TargetAddrMap))
	for k := range t.TargetAddrMap {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	for _, k := range keys {
		if strings.Contains(k, nodeAddr) {
			sb.WriteString(fmt.Sprintf("%s=%d,", "0.0.0.0:9102", t.nodeCount))
			continue
		}
		sb.WriteString(fmt.Sprintf("%s=%d,", k, t.TargetAddrMap[k]))
	}
	return sb.String()[0 : len(sb.String())-1]
}