package k8s

import (
	"context"

	"github.com/pingcap/errors"
	oplabel "github.com/pingcap/tidb-operator/pkg/apis/label"
	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"
	EBSZoneLabel    = "topology.ebs.csi.aws.com/zone"
	GKEZoneLabel    = "topology.gke.io/zone"

	// ref: https://github.com/pingcap/tidb-operator/blob/2e78d9c8ff86449360d5025a1c40dc175c04b195/pkg/apis/label/label.go#L83:2
	PVCDeferDeletingAnno = oplabel.AnnPVCDeferDeleting
)

type ComponentTopology struct {
	Replicas    int
	Scheduled   PodsByZone
	UnScheduled PodNames
}

type ComponentPVCTopology struct {
	// Bound is the PVCs that are bound to PVs
	// PVCs that are scaled in are not included.
	Bound PVCsByZone
}

type PodNames []string

type PodsByZone map[string]PodNames

type PVCNames []string

type PVCsByZone map[string]PVCNames

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 GetComponentPVCTopology(ctx context.Context, tcCli client.Client, tc types.NamespacedName, component cloudv1.ComponentType) (*ComponentPVCTopology, error) {
	pvcLabels := map[string]string{
		label.ComponentLabelKey: string(component),
		label.InstanceLabelKey:  tc.Name,
	}
	pvcs, err := SelectPVCs(ctx, tcCli, tc.Namespace, pvcLabels)
	if err != nil {
		return nil, errors.Annotate(err, "ListComponentPVCs")
	}

	bound := map[string]PVCNames{}
	for _, pvc := range pvcs {
		if pvc.Status.Phase != corev1.ClaimBound {
			continue
		}

		if _, exist := pvc.Annotations[PVCDeferDeletingAnno]; exist {
			continue
		}

		if pvc.Spec.VolumeName == "" {
			continue
		}

		pv := &corev1.PersistentVolume{}
		err := tcCli.Get(ctx, types.NamespacedName{Name: pvc.Spec.VolumeName}, pv)
		if err != nil {
			return nil, errors.Annotate(err, "GetPV")
		}

		zone := readZoneOfPV(pv)
		if zone != "" {
			bound[zone] = append(bound[zone], pvc.Name)
		}
	}
	return &ComponentPVCTopology{
		Bound: bound,
	}, 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
}

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

func readZoneOfPV(pv *corev1.PersistentVolume) string {
	supportedZoneLabels := map[string]struct{}{
		ZoneLabel:       {},
		LegacyZoneLabel: {},
		EBSZoneLabel:    {},
		GKEZoneLabel:    {},
	}

	// read from label
	for k, v := range pv.Labels {
		if _, ok := supportedZoneLabels[k]; ok {
			return v
		}
	}

	// read from node affinity
	if pv.Spec.NodeAffinity != nil && pv.Spec.NodeAffinity.Required != nil {
		for _, term := range pv.Spec.NodeAffinity.Required.NodeSelectorTerms {
			for _, exp := range term.MatchExpressions {
				if exp.Operator != corev1.NodeSelectorOpIn {
					continue
				}

				if _, exist := supportedZoneLabels[exp.Key]; exist {
					return exp.Values[0]
				}
			}
		}
	}

	return ""
}