package k8s

import (
	"context"

	"github.com/pingcap/errors"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/pkg/pingcap/label"
)

const (
	// TODO: move all consts to pkg/consts.go
	DefaultStorageClassName = "cloud-ssd"

	IopsToVolumeSizeRatioMaximum = 500

	PDGp3StorageClassName = "pd-gp3"
	PDGp3IOPS             = "4000"
	PDGp3ThroughputMBps   = "400"

	TiKVGp3StorageClassName = "tikv-gp3"
	TiKVGp3IOPS             = "4000"
	TiKVGp3ThroughputMBps   = "400"

	// TODO: add validation for tiflash storage.
	// Iops/Storage maximum is 500. so currently tiflash storage size must be larger than 12Gi.
	TiFlashGp3StorageClassName = "tiflash-gp3"
	TiFlashGp3IOPS             = "6000"
	TiFlashGp3ThroughputMBps   = "625"

	TiCDCGp3StorageClassName = "ticdc-gp3"
	TiCDCGp3IOPS             = "4000"
	TiCDCGp3ThroughputMBps   = "400"

	ImportGp3StorageClassName = "import-gp3"
	ImportGp3IOPS             = "16000"
	ImportGp3ThroughputMBps   = "1000"

	// ref: https://kubernetes.io/docs/reference/labels-annotations-taints/#failure-domainbetakubernetesiozone
	ZoneLabel       = "topology.kubernetes.io/zone"
	LegacyZoneLabel = "failure-domain.beta.kubernetes.io/zone"

	TiDBAddIndexAccelerateGp3StorageClassName      = "tidb-add-index-acceleration-gp3"
	TiDBAddIndexAccelerateGp3StorageDefaultSize    = "100Gi"
	TiDBAddIndexAccelerateGp3StorageDefaultDirName = "tidb-add-index-acceleration"
	TiDBAddIndexAccelerateGp3StorageDefaultDirPath = "/data/tmp"
)

type ComponentTopology struct {
	Replicas    int
	Scheduled   PodsByZone
	UnScheduled PodNames
}

type PodNames []string

type PodsByZone map[string]PodNames

func GetComponentTopology(ctx context.Context, tcCli client.Client, tc types.NamespacedName, component cloudv1.ComponentType) (*ComponentTopology, error) {
	podLabels := map[string]string{
		label.ComponentLabelKey: string(component),
		label.InstanceLabelKey:  tc.Name,
	}
	pods, err := SelectPods(ctx, tcCli, tc.Namespace, podLabels)
	if err != nil {
		return nil, errors.Annotate(err, "ListComponentPods")
	}
	replicas := len(pods)
	var unscheduled PodNames
	scheduled := map[string]PodNames{}
	for _, pod := range pods {
		if pod.Spec.NodeName == "" {
			unscheduled = append(unscheduled, pod.Name)
			continue
		}
		node := &corev1.Node{}
		// read from cache, no iterative RPC issue
		err := tcCli.Get(ctx, types.NamespacedName{Name: pod.Spec.NodeName}, node)
		if err != nil {
			return nil, errors.Annotate(err, "GetNode")
		}
		var zone string
		if v, ok := node.Labels[ZoneLabel]; ok {
			zone = v
		} else if v, ok := node.Labels[LegacyZoneLabel]; ok {
			zone = v // fallback to legacy (k8s version < 1.17)
		}
		if zone == "" {
			// impossible in the cloud, mock envs are expected to mock these labels too
			return nil, errors.Errorf("no zone labels found in node %s (mock it if not running in real cloud)", node.Name)
		}
		scheduled[zone] = append(scheduled[zone], pod.Name)
	}
	return &ComponentTopology{
		Replicas:    replicas,
		Scheduled:   scheduled,
		UnScheduled: unscheduled,
	}, nil
}

func SelectNodes(ctx context.Context, cli client.Client, selectors map[string]string) ([]corev1.Node, error) {
	nodeList := &corev1.NodeList{}
	if err := cli.List(ctx, nodeList, client.MatchingLabels(selectors)); err != nil {
		return nil, err
	}
	return nodeList.Items, nil
}

func SelectPods(ctx context.Context, cli client.Client, namespace string, selectors map[string]string) ([]corev1.Pod, error) {
	podList := &corev1.PodList{}
	if err := cli.List(ctx, podList, client.InNamespace(namespace), client.MatchingLabels(selectors)); err != nil {
		return nil, err
	}
	return podList.Items, nil
}