// Copyright 2017 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package highavailability

import (
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap.com/v1"
	"github.com/pingcap/tidb-operator/pkg/scheduler/predicates"
	"github.com/pingcap/tidb-operator/pkg/scheduler/volume"
	"github.com/pingcap/tidb-operator/pkg/util/label"
	"github.com/pingcap/tidb-operator/pkg/volumeprovisioner"
	apiv1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/kubernetes"
)

type highAvailability struct {
	kubeCli kubernetes.Interface
	weight  int32
}

// NewHighAvailability returns a predicates.Predicate
func NewHighAvailability(weight int32, kubeCli kubernetes.Interface) predicates.Predicate {
	return &highAvailability{weight: weight, kubeCli: kubeCli}
}

// Name returns predicate name
func (ha *highAvailability) Name() string {
	return "HighAvailability"
}

func (ha *highAvailability) Filter(tidbVol *volume.TiDBVolume, pod *apiv1.Pod, pvc *apiv1.PersistentVolumeClaim) (bool, error) {
	ns := pod.Namespace
	podLabel := label.New().Cluster(pod.Labels[label.ClusterLabelKey]).App(pod.Labels[label.AppLabelKey])
	podList, err := ha.kubeCli.CoreV1().Pods(ns).List(metav1.ListOptions{LabelSelector: labels.SelectorFromSet(podLabel.Labels()).String()})
	if err != nil {
		return true, errors.Trace(err)
	}

	var scheduledPods []string
	nodes := make(map[string]int)
	for _, pod := range podList.Items {
		var claimName string
		for _, vol := range pod.Spec.Volumes {
			if vol.Name == v1.TiDBVolumeName && vol.VolumeSource.PersistentVolumeClaim != nil {
				claimName = vol.VolumeSource.PersistentVolumeClaim.ClaimName
				break
			}
		}

		if len(claimName) == 0 {
			return true, errors.Errorf("Pod: [%s/%s]'s claimName is empty", pod.GetNamespace(), pod.GetName())
		}

		pvc, err := ha.kubeCli.CoreV1().PersistentVolumeClaims(ns).Get(claimName, metav1.GetOptions{})
		if err != nil {
			return true, errors.Trace(err)
		}

		nodeName, ok := pvc.Annotations[volumeprovisioner.AnnProvisionerNode]
		if !ok || nodeName == "" {
			continue
		}

		scheduledPods = append(scheduledPods, pod.Name)
		nodes[nodeName]++
	}
	log.Debugf("HA predicate, scheduled pods: %v", scheduledPods)
	idxs := make([]int, 0)
	for idx := 0; idx < len(tidbVol.Volumes); idx++ {
		vol := &tidbVol.Volumes[idx]
		// TODO: This only works for 3 region replicas
		// There should be at least 3 nodes:
		// we can't schedule the 2nd pod on the same node as the 1st pod
		// nor schedule the 3rd pod on the same node as the first two pod
		if (len(scheduledPods) == 1 || len(scheduledPods) == 2) && nodes[vol.NodeName] != 0 {
			idxs = append(idxs, idx)
			continue
		}
		vol.Priority -= int32(nodes[vol.NodeName]) * ha.weight
	}
	tidbVol.DeleteWithIdxs(idxs)
	log.Debugf("after HA filter tidb volumes are: %v", tidbVol.Volumes)
	return false, nil
}