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

package e2e

import (
	"fmt"
	"os/exec"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pingcap/tidb-operator/pkg/scheduler/volume"
	"github.com/pingcap/tidb-operator/test/e2e/bootstrap"
	"github.com/pingcap/tidb-operator/test/e2e/framework"
	"github.com/pingcap/tidb-operator/test/e2e/util"
	yaml "gopkg.in/yaml.v2"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
)

func testHA(f *framework.Framework) {
	clusterName := "e2e-ha"
	var pdCount int32 = 1
	var tikvCount int32 = 6

	Expect(labelNode(f, "tikv")).To(Succeed())

	By("Given 1 node for TiDBVolume")
	Expect(generateTiDBVolumes(f, 1)).To(Succeed())

	By("When create a TiDB cluster")
	cmd := fmt.Sprintf("helm install /tmp/tidb-cluster --namespace=%s --name=%s --values=/tmp/data/e2e-ha-values.yaml", e2eNamespace, clusterName)
	data, err := exec.Command("/bin/sh", "-c", cmd).CombinedOutput()
	if err != nil {
		util.Logf(string(data))
	}
	Expect(err).NotTo(HaveOccurred())

	done := watchRestartCount(f, e2eNamespace, clusterName)

	By("Then all PDs are running")
	err = wait.Poll(5*time.Second, 4*time.Minute, func() (bool, error) {
		podList, err := f.GetAllPDPods(e2eNamespace, clusterName)
		if err != nil {
			return false, nil
		}

		// default: 1 pd
		if len(podList.Items) != int(pdCount) {
			return false, nil
		}

		for _, pod := range podList.Items {
			running, err := f.IsPodRunning(&pod)
			if err != nil || !running {
				return false, nil
			}
		}

		return true, err
	})
	Expect(err).NotTo(HaveOccurred())

	By("Then waiting 1 minute to make sure 1 tikv pod is scheduled and 2 tikv pods are unable to schedule")
	time.Sleep(1 * time.Minute)
	err = wait.Poll(2*time.Second, 2*time.Minute, func() (bool, error) {
		podList, err := f.GetAllTiKVPods(e2eNamespace, clusterName)
		if err != nil {
			return false, nil
		}
		if len(podList.Items) != 3 {
			return false, nil
		}
		var scheduled int
		for _, pod := range podList.Items {
			if pod.Spec.NodeName != "" {
				scheduled++
			}
		}
		if scheduled != 1 {
			return false, nil
		}
		return true, nil
	})
	Expect(err).NotTo(HaveOccurred())

	By("When given 2 nodes for TiDBVolume")
	Expect(generateTiDBVolumes(f, 2)).To(Succeed())

	By("Then waiting 1 minute to make sure 2 tikv pods are scheduled and 1 tikv pod is unable to schedule")
	time.Sleep(1 * time.Minute)
	err = wait.Poll(2*time.Second, 2*time.Minute, func() (bool, error) {
		podList, err := f.GetAllTiKVPods(e2eNamespace, clusterName)
		if err != nil {
			return false, nil
		}
		if len(podList.Items) != 3 {
			return false, nil
		}
		var scheduled int
		for _, pod := range podList.Items {
			if pod.Spec.NodeName != "" {
				scheduled++
			}
		}
		if scheduled != 2 {
			return false, nil
		}
		return true, nil
	})
	Expect(err).NotTo(HaveOccurred())

	By("When given 3 nodes for TiDBVolume")
	Expect(generateTiDBVolumes(f, 3)).To(Succeed())

	By("Then waiting 2 minutes to make sure all 3 tikv pods are scheduled")
	err = wait.Poll(1*time.Second, 2*time.Minute, func() (bool, error) {
		podList, err := f.GetAllTiKVPods(e2eNamespace, clusterName)
		if err != nil {
			return false, nil
		}
		if len(podList.Items) != 3 {
			return false, nil
		}
		nodes := make(map[string]int)
		for _, pod := range podList.Items {
			nodeName := pod.Spec.NodeName
			if nodeName == "" {
				return false, nil
			}
			nodes[nodeName]++
		}
		if len(nodes) != 3 {
			return false, errors.Errorf("tikv are not scheduled evenly: %v", nodes)
		}
		return true, nil
	})
	Expect(err).NotTo(HaveOccurred())

	By("When resize tikv's replicas to 5")
	tc, err := f.Cli.PingcapV1().TidbClusters(e2eNamespace).Get(clusterName, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())

	tc.Spec.TiKV.Size = 5

	_, err = f.Cli.PingcapV1().TidbClusters(e2eNamespace).Update(tc)
	Expect(err).NotTo(HaveOccurred())

	By("Then make sure no 3 tikv on the same node")
	err = wait.Poll(5*time.Second, 2*time.Minute, func() (bool, error) {
		podList, err := f.GetAllTiKVPods(e2eNamespace, clusterName)
		if err != nil {
			return false, nil
		}
		if len(podList.Items) != 5 {
			return false, nil
		}
		nodes := make(map[string]int)
		for _, pod := range podList.Items {
			nodeName := pod.Spec.NodeName
			if nodeName == "" {
				return false, nil
			}
			nodes[nodeName]++
		}
		for node, pods := range nodes {
			if pods >= 3 {
				return false, errors.Errorf("node %s has %d pods which is greater than 3", node, pods)
			}
		}
		return true, nil
	})
	Expect(err).NotTo(HaveOccurred())

	By(fmt.Sprintf("When resize tikv's replicas to %d", tikvCount))
	tc, err = f.Cli.PingcapV1().TidbClusters(e2eNamespace).Get(clusterName, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())

	tc.Spec.TiKV.Size = tikvCount

	_, err = f.Cli.PingcapV1().TidbClusters(e2eNamespace).Update(tc)
	Expect(err).NotTo(HaveOccurred())
	By("Then make sure each node has 2 tikv pods")
	err = wait.Poll(5*time.Second, 2*time.Minute, func() (bool, error) {
		podList, err := f.GetAllTiKVPods(e2eNamespace, clusterName)
		if err != nil {
			return false, nil
		}
		if len(podList.Items) != int(tikvCount) {
			return false, nil
		}
		nodes := make(map[string]int)
		for _, pod := range podList.Items {
			nodeName := pod.Spec.NodeName
			if nodeName == "" {
				return false, nil
			}
			nodes[nodeName]++
		}
		for node, pods := range nodes {
			if pods != 2 {
				return false, errors.Errorf("node %s has %d pods which is not 2", node, pods)
			}
		}
		return true, nil
	})
	Expect(err).NotTo(HaveOccurred())

	deleteCluster(f, e2eNamespace, clusterName, done, pdCount+tikvCount)
}

func labelNode(f *framework.Framework, kind string) error {
	nodeList, err := f.KubeCli.CoreV1().Nodes().List(metav1.ListOptions{})
	if err != nil {
		return err
	}
	for _, node := range nodeList.Items {
		schedulable := !node.Spec.Unschedulable
		for _, taint := range node.Spec.Taints {
			if taint.Effect == v1.TaintEffectNoSchedule {
				schedulable = false
				break
			}
		}
		if !schedulable {
			continue
		}
		err = wait.Poll(2*time.Second, 20*time.Second, func() (bool, error) {
			latestNode, err := f.KubeCli.CoreV1().Nodes().Get(node.Name, metav1.GetOptions{})
			if err != nil {
				return false, nil
			}
			latestNode.Labels["kind"] = kind
			_, err = f.KubeCli.CoreV1().Nodes().Update(latestNode)
			if err != nil {
				return false, nil
			}
			return true, nil
		})
		if err != nil {
			log.Errorf("failed to set node label kind=tikv: %v", err)
		}
	}
	return nil
}

func generateTiDBVolumes(f *framework.Framework, n int) error {
	nodeList, err := f.KubeCli.CoreV1().Nodes().List(metav1.ListOptions{})
	if err != nil {
		return err
	}
	nodes := []string{}
	for _, node := range nodeList.Items {
		schedulable := !node.Spec.Unschedulable
		for _, taint := range node.Spec.Taints {
			if taint.Effect == v1.TaintEffectNoSchedule {
				schedulable = false
				break
			}
		}
		if schedulable {
			nodes = append(nodes, node.Name)
		}
	}
	var tidbVolumes volume.TiDBVolume
	for _, node := range nodes {
		if len(tidbVolumes.Nodes) >= n {
			break
		}
		tidbVolumes.Nodes = append(tidbVolumes.Nodes, volume.Node{
			Name: node,
			Directories: []volume.Directory{
				{Path: "/tidb/pd-dir-1", Size: "1Gi"},
				{Path: "/tidb/pd-dir-2", Size: "1.5Gi"},
				{Path: "/tidb/pd-dir-3", Size: "2Gi"},
				{Path: "/tidb/pd-dir-4", Size: "3Gi"},
				{Path: "/tidb/pd-dir-5", Size: "2.5Gi"},

				{Path: "/tidb/tikv-dir-1", Size: "100Gi"},
				{Path: "/tidb/tikv-dir-2", Size: "100Gi"},
				{Path: "/tidb/tikv-dir-3", Size: "100Gi"},
				{Path: "/tidb/tikv-dir-4", Size: "100Gi"},
				{Path: "/tidb/tikv-dir-5", Size: "100Gi"},
			},
		})
	}
	data, err := yaml.Marshal(tidbVolumes)
	if err != nil {
		return err
	}
	cm, err := f.KubeCli.CoreV1().ConfigMaps(bootstrap.OperatorNamespace).Get(bootstrap.VolumeConfigMapName, metav1.GetOptions{})
	if err != nil {
		return err
	}
	cm.Data = map[string]string{"volume": string(data)}
	_, err = f.KubeCli.CoreV1().ConfigMaps(bootstrap.OperatorNamespace).Update(cm)
	return err
}