// 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 scheduler

import (
	"fmt"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap.com/v1"
	"github.com/pingcap/tidb-operator/pkg/client/clientset/versioned"
	"github.com/pingcap/tidb-operator/pkg/scheduler/predicates"
	"github.com/pingcap/tidb-operator/pkg/scheduler/predicates/highavailability"
	"github.com/pingcap/tidb-operator/pkg/scheduler/predicates/volumebestfit"
	"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/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	schedulerapiv1 "k8s.io/kubernetes/pkg/scheduler/api/v1"
)

var errNoSuitableNode = errors.New("Could not find a suitable node")

type noSuitablePVErr struct {
	err string
}

func (e *noSuitablePVErr) Error() string {
	return e.err
}

// Scheduler is an interface for external processes to influence scheduling
// decisions made by kubernetes. This is typically needed for resources not directly
// managed by kubernetes.
type Scheduler interface {
	// Filter based on extender-implemented predicate functions. The filtered list is
	// expected to be a subset of the supplied list.
	Filter(*schedulerapiv1.ExtenderArgs) (*schedulerapiv1.ExtenderFilterResult, error)

	// Prioritize based on extender-implemented priority functions. The returned scores & weight
	// are used to compute the weighted score for an extender. The weighted scores are added to
	// the scores computed  by kubernetes scheduler. The total scores are used to do the host selection.
	Priority(*schedulerapiv1.ExtenderArgs) (schedulerapiv1.HostPriorityList, error)
}

var _ Scheduler = &tidbScheduler{}

type tidbScheduler struct {
	cli         versioned.Interface
	kubeCli     kubernetes.Interface
	predicates  []predicates.Predicate
	volumesPool volume.Pool
}

// NewTiDBScheduler returns a Scheduler
func NewTiDBScheduler(ns, cmTiDBVolume string, cli versioned.Interface, kubeCli kubernetes.Interface) Scheduler {
	return &tidbScheduler{
		cli:     cli,
		kubeCli: kubeCli,
		predicates: []predicates.Predicate{
			volumebestfit.NewVolumeBestFit(5),
			highavailability.NewHighAvailability(10, kubeCli),
		},
		volumesPool: volume.NewVolumePool(volumeConfigmap{
			ns: ns, name: cmTiDBVolume, kubeCli: kubeCli,
		}),
	}
}

// Filter select a node from *schedulerapiv1.ExtenderArgs.Nodes when this is a pd or tikv pod
// else return the original nodes.
func (ts *tidbScheduler) Filter(args *schedulerapiv1.ExtenderArgs) (*schedulerapiv1.ExtenderFilterResult, error) {
	pod := &args.Pod
	podName := podName(pod)
	ns := pod.GetNamespace()
	kubeNodes := args.Nodes.Items

	// node names from kubernetes's scheduler
	log.Infof("Start scheduling pod: [%s], tidb-scheduler received nodes: %+v", podName, getNodeNamesFrom(kubeNodes))

	// only for pd and tikv
	if l, ok := args.Pod.Labels[label.AppLabelKey]; !ok || (l != label.TiKVLabelVal && l != label.PDLabelVal) {
		return &schedulerapiv1.ExtenderFilterResult{
			Nodes: args.Nodes,
		}, nil
	}

	// PVC must exists
	pvc, pvcErr := ts.getPVCFrom(pod)
	if pvcErr != nil {
		log.Errorf("failed to get PVC from pod[%s]: %v", podName, pvcErr)
		return nil, errors.Trace(pvcErr)
	}
	pvcName := pvc.GetName()
	if pvc.Annotations == nil {
		errMsg := fmt.Sprintf("PVC: [%s/%s] doesn't have annotations", ns, pvcName)
		return &schedulerapiv1.ExtenderFilterResult{Error: errMsg}, errors.New(errMsg)
	}

	node := pvc.Annotations[volumeprovisioner.AnnProvisionerNode]
	hp := pvc.Annotations[volumeprovisioner.AnnProvisionerHostPath]
	if len(node) != 0 && len(hp) != 0 {
		if pvc.Status.Phase == apiv1.ClaimBound {
			log.Infof("Pod: [%s] will be scheduled on node: [%s] with a hostPath: [%s] associated with it", podName, node, hp)
			return findANode(kubeNodes, node)
		}

		return &schedulerapiv1.ExtenderFilterResult{Error: "waiting for the PVC bound with a PV"}, nil
	}

	clusterName := pod.Labels[label.ClusterLabelKey]
	if len(clusterName) == 0 {
		return nil, errors.Errorf("Pod: [%s/%s]'s labels: [%s] is empty", ns, pod.GetName(), label.ClusterLabelKey)
	}

	tc, clusterErr := ts.cli.PingcapV1().TidbClusters(ns).Get(clusterName, metav1.GetOptions{})
	if clusterErr != nil {
		return nil, errors.Trace(clusterErr)
	}

	// try to reuse cluster's PV
	err := ts.tryToReuse(pod, pvc, clusterName)
	// If we found a suitable PV
	if err == nil {
		return &schedulerapiv1.ExtenderFilterResult{Error: "waiting for the PVC bound with a PV"}, nil
	}
	// Else if we are restoring
	if tc.ReallyRestoring() {
		log.Errorf("Pod: [%s] restored failed: %v", podName, err)
		return nil, errors.Trace(err)
	}
	// If err is a error other than noSuitablePVErr
	if _, ok := err.(*noSuitablePVErr); !ok {
		return nil, errors.Trace(err)
	}
	// Else we can't find a suitable PV

	// storage size must exists
	pvcQuantity, ok := pvc.Spec.Resources.Requests[apiv1.ResourceStorage]
	if !ok {
		errMsg := fmt.Sprintf("PVC: [%s/%s] must have a request storage size", ns, pvcName)
		log.Error(errMsg)
		return nil, errors.New(errMsg)
	}

	tidbVol, poolErr := ts.volumesPool.Volumes()
	if poolErr != nil {
		log.Errorf("Error try to get tidb volumes when scheduling pod: [%s], %v", podName, poolErr)
		return nil, errors.Trace(poolErr)
	}

	// already provisioned volumes
	alreadyProvVolumes, provErr := ts.alreadyProvVolumes()
	if provErr != nil {
		log.Errorf("Error try to get already provisioned volumes when scheduling pod: [%s], %v", podName, provErr)
		return nil, errors.Trace(provErr)
	}
	log.Debugf("already provisioned volumes when scheduling pod: [%s], %v", podName, alreadyProvVolumes)

	// available volumes
	tv, volErr := tidbVol.From(kubeNodes)
	if volErr != nil {
		log.Errorf("Error fetching available volumes when scheduling pod: [%s], %v", podName, volErr)
		return nil, errors.Trace(volErr)
	}
	tv.Exclude(alreadyProvVolumes).WithQuantity(pvcQuantity)
	log.Debugf("available volumes when scheduling pod: [%s], %v", podName, tidbVol.Volumes)

	for _, predicate := range ts.predicates {
		if len(tidbVol.Volumes) == 0 {
			log.Error(errNoSuitableNode)
			return nil, errNoSuitableNode
		}

		fuse, err := predicate.Filter(tidbVol, pod, pvc)
		if err != nil {
			log.Errorf("Error when filter with predicate: %s when scheduling pod: [%s], %v", predicate.Name(), podName, err)
			return nil, errors.Trace(err)
		}

		if fuse {
			if len(tidbVol.Volumes) == 0 {
				errMsg := fmt.Sprintf("Predicate: %s is fusing, but don't return volume when scheduling pod: [%s]", predicate.Name(), podName)
				log.Error(errMsg)
				return nil, errors.New(errMsg)
			}

			nodeName := tidbVol.Volumes[0].NodeName
			log.Infof("Predicate: %s select a node: %s when scheduling pod: [%s]", predicate.Name(), nodeName, podName)

			break
		}

		log.Debugf("After Predicate: %s, volumes: %v when scheduling pod: [%s]", predicate.Name(), tidbVol.Volumes, podName)
	}

	vol := tidbVol.HighestPriorityVolume()
	if vol == nil {
		log.Error(errNoSuitableNode)
		return nil, errNoSuitableNode
	}

	if err := ts.markPVCProvisionedWith(pvc, vol.NodeName, vol.Path, nil); err != nil {
		return nil, errors.Trace(err)
	}

	return &schedulerapiv1.ExtenderFilterResult{Error: "waiting for the PVC bound with a PV"}, nil
}

// TODO add more e2e specs to ensure the logic is right
func (ts *tidbScheduler) tryToReuse(pod *apiv1.Pod, pvc *apiv1.PersistentVolumeClaim, clusterName string) error {
	ns := pvc.GetNamespace()
	app := pod.Labels[label.AppLabelKey]
	podName := pod.GetName()

	pvLabel := label.New().Cluster(pod.Labels[label.ClusterLabelKey]).App(app).Namespace(ns)
	pvListOptions := metav1.ListOptions{
		LabelSelector: labels.SelectorFromSet(pvLabel.Labels()).String(),
	}
	pvList, err := ts.kubeCli.CoreV1().PersistentVolumes().List(pvListOptions)
	if err != nil {
		return errors.Trace(err)
	}

	pvcLabel := label.New().Cluster(pod.Labels[label.ClusterLabelKey]).App(app)
	pvcListOptions := metav1.ListOptions{
		LabelSelector: labels.SelectorFromSet(pvcLabel.Labels()).String(),
	}
	pvcList, err := ts.kubeCli.CoreV1().PersistentVolumeClaims(ns).List(pvcListOptions)
	if err != nil {
		return errors.Trace(err)
	}

outerLabel:
	for _, pv := range pvList.Items {
		pvNode := pv.Annotations[volumeprovisioner.AnnProvisionerNode]
		pvHostPath := pv.Annotations[volumeprovisioner.AnnProvisionerHostPath]
		pvTombstoned := pv.Annotations[volumeprovisioner.AnnPVTombstoned]
		pvName := pv.GetName()

		if len(pvTombstoned) != 0 {
			log.Debug("PV: %s is marked tombstone at: %s, skipping", pvName, pvTombstoned)
			continue
		}

		for _, pc := range pvcList.Items {
			pvcNode := pc.Annotations[volumeprovisioner.AnnProvisionerNode]
			pvcHostPath := pc.Annotations[volumeprovisioner.AnnProvisionerHostPath]
			if len(pvNode) != 0 && len(pvHostPath) != 0 && pvNode == pvcNode && pvHostPath == pvcHostPath {
				continue outerLabel
			}
		}

		switch app {
		case label.PDLabelVal:
			pn := pv.Annotations[label.AnnPodNameKey]
			if pn != podName {
				continue outerLabel
			}
			return ts.markPVCProvisionedWith(
				pvc,
				pv.Annotations[volumeprovisioner.AnnProvisionerNode],
				pv.Annotations[volumeprovisioner.AnnProvisionerHostPath],
				&pvName,
			)
		case label.TiKVLabelVal:
			return ts.markPVCProvisionedWith(
				pvc,
				pv.Annotations[volumeprovisioner.AnnProvisionerNode],
				pv.Annotations[volumeprovisioner.AnnProvisionerHostPath],
				&pvName,
			)
		default:
		}
	}

	errMsg := fmt.Sprintf("can't find a PV for PVC: [%s/%s], Pod: [%s/%s]", ns, pvc.GetName(), ns, podName)
	log.Warn(errMsg)
	return &noSuitablePVErr{errMsg}
}

// We didn't pass `prioritizeVerb` to kubernetes scheduler extender's config file, this method will not be called.
func (ts *tidbScheduler) Priority(args *schedulerapiv1.ExtenderArgs) (schedulerapiv1.HostPriorityList, error) {
	result := schedulerapiv1.HostPriorityList{}

	// avoid index out of range panic
	if len(args.Nodes.Items) > 0 {
		result = append(result, schedulerapiv1.HostPriority{
			Host:  args.Nodes.Items[0].Name,
			Score: 0,
		})
	}

	return result, nil
}

func getNodeNamesFrom(nodes []apiv1.Node) []string {
	nodeNames := make([]string, 0)
	for _, n := range nodes {
		nodeNames = append(nodeNames, n.GetName())
	}

	return nodeNames
}

func (ts *tidbScheduler) markPVCProvisionedWith(pvc *apiv1.PersistentVolumeClaim, nodeName, hostPath string, shouldDeletedPVName *string) error {
	return wait.Poll(2*time.Second, 15*time.Second, func() (bool, error) {
		pvcNew, err := ts.kubeCli.CoreV1().PersistentVolumeClaims(pvc.Namespace).Get(pvc.Name, metav1.GetOptions{})
		if err != nil {
			log.Errorf("PVC[%s] get failed: %v", pvcName(pvc), err)
			return false, nil
		}

		if pvcNew.Annotations == nil {
			pvcNew.Annotations = make(map[string]string)
		}

		pvcNew.Annotations[volumeprovisioner.AnnProvisionerNode] = nodeName
		pvcNew.Annotations[volumeprovisioner.AnnProvisionerHostPath] = hostPath
		if shouldDeletedPVName != nil {
			pvcNew.Annotations[volumeprovisioner.AnnShouldDeletedPVName] = *shouldDeletedPVName
			log.Infof("Set PVC[%s] Annotations[%s] = %s", pvcName(pvcNew), volumeprovisioner.AnnShouldDeletedPVName, *shouldDeletedPVName)
		}
		log.Infof(
			"Set PVC[%s] Annotations[%s] = %s and Annotations[%s] = %s",
			pvcName(pvcNew),
			volumeprovisioner.AnnProvisionerNode,
			nodeName,
			volumeprovisioner.AnnProvisionerHostPath,
			hostPath,
		)

		_, err = ts.kubeCli.CoreV1().PersistentVolumeClaims(pvc.Namespace).Update(pvcNew)
		if err != nil {
			log.Errorf("PVC[%s] annotation updated failed: %v", pvcName(pvc), err)
			return false, nil
		}
		log.Infof("PVC[%s] annotation updated successfully", pvcName(pvc))

		return true, err
	})
}

func (ts *tidbScheduler) getPVCFrom(pod *apiv1.Pod) (*apiv1.PersistentVolumeClaim, error) {
	for _, vol := range pod.Spec.Volumes {
		if vol.Name == v1.TiDBVolumeName && vol.PersistentVolumeClaim != nil {
			return ts.kubeCli.CoreV1().PersistentVolumeClaims(pod.Namespace).Get(vol.PersistentVolumeClaim.ClaimName, metav1.GetOptions{})
		}
	}

	return nil, errors.Errorf("pod: %s doesn't have an associated PVC", podName(pod))
}

// merge PV and PVC's already provisioned volumes
func (ts *tidbScheduler) alreadyProvVolumes() ([]volume.Volume, error) {
	pvcList, err := ts.kubeCli.CoreV1().PersistentVolumeClaims(metav1.NamespaceAll).List(metav1.ListOptions{
		LabelSelector: labels.SelectorFromSet(
			label.New().Labels(),
		).String(),
	})
	if err != nil {
		log.Errorf("failed to list PVC from all namespaces: %v", err)
		return nil, errors.Trace(err)
	}

	pvList, err := ts.kubeCli.CoreV1().PersistentVolumes().List(metav1.ListOptions{
		LabelSelector: labels.SelectorFromSet(
			label.New().Labels(),
		).String(),
	})
	if err != nil {
		log.Errorf("failed to list PV: %v", err)
		return nil, errors.Trace(err)
	}

	volumes := make([]volume.Volume, 0)

	for _, pvc := range pvcList.Items {
		node, nodeOk := pvc.Annotations[volumeprovisioner.AnnProvisionerNode]
		hp, hpOk := pvc.Annotations[volumeprovisioner.AnnProvisionerHostPath]
		if !nodeOk || len(node) == 0 || !hpOk || len(hp) == 0 {
			continue
		}

		volumes = append(volumes, volume.Volume{NodeName: node, Path: hp})
	}

outer:
	for _, pv := range pvList.Items {
		node, nodeOk := pv.Annotations[volumeprovisioner.AnnProvisionerNode]
		hp, hpOk := pv.Annotations[volumeprovisioner.AnnProvisionerHostPath]
		if !nodeOk || len(node) == 0 || !hpOk || len(hp) == 0 {
			continue
		}

		for i := 0; i < len(volumes); i++ {
			vol := volumes[i]
			if vol.NodeName == node && vol.Path == hp {
				continue outer
			}
		}

		volumes = append(volumes, volume.Volume{NodeName: node, Path: hp})
	}

	return volumes, nil
}

func findANode(kubeNodes []apiv1.Node, nodeName string) (*schedulerapiv1.ExtenderFilterResult, error) {
	for _, n := range kubeNodes {
		if n.GetName() == nodeName {
			return &schedulerapiv1.ExtenderFilterResult{
				Nodes: &apiv1.NodeList{
					Items: []apiv1.Node{n},
				},
			}, nil
		}
	}

	return nil, errNoSuitableNode
}

func podName(pod *apiv1.Pod) string {
	return objName(&pod.ObjectMeta)
}

func pvcName(pvc *apiv1.PersistentVolumeClaim) string {
	return objName(&pvc.ObjectMeta)
}

func objName(om *metav1.ObjectMeta) string {
	return fmt.Sprintf("%s/%s", om.GetNamespace(), om.GetName())
}

// volumeConfigmap used as configuration for the volume pool.
type volumeConfigmap struct {
	ns      string
	name    string
	kubeCli kubernetes.Interface
}

func (cm volumeConfigmap) Name() string { return fmt.Sprintf("%s/%s", cm.ns, cm.name) }

func (cm volumeConfigmap) Config() (map[string]string, error) {
	cfg, err := cm.kubeCli.CoreV1().ConfigMaps(cm.ns).Get(cm.name, metav1.GetOptions{})
	if err != nil {
		return nil, errors.Trace(err)
	}
	return cfg.Data, err
}