// 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"
	"time"

	"github.com/juju/errors"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap.com/v1"
	"github.com/pingcap/tidb-operator/test/e2e/framework"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
)

func testUpgrade(f *framework.Framework) {
	clusterName := "e2e-upgrade"
	var pdCount int32 = 3
	var tidbCount int32 = 1
	var tikvCount int32 = 3

	done := deployCluster(f, e2eNamespace, clusterName, pdCount, tidbCount, tikvCount)

	tidbVersion := "v2.0.4"
	pdImage := fmt.Sprintf("uhub.ucloud.cn/pingcap/pd:%s", tidbVersion)
	tikvImage := fmt.Sprintf("uhub.ucloud.cn/pingcap/tikv:%s", tidbVersion)
	tidbImage := fmt.Sprintf("uhub.ucloud.cn/pingcap/tidb:%s", tidbVersion)

	By("When upgrade the cluster")
	tc, err := f.Cli.PingcapV1().TidbClusters(e2eNamespace).Get(clusterName, metav1.GetOptions{})
	Expect(err).NotTo(HaveOccurred())

	tc.Spec.PD.Image = pdImage
	tc.Spec.TiKV.Image = tikvImage
	tc.Spec.TiDB.Image = tidbImage

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

	By("Then wait until pd/tikv/tidb is upgrading")
	err = wait.Poll(5*time.Second, 5*time.Minute, func() (bool, error) {
		tc, err := f.Cli.PingcapV1().TidbClusters(e2eNamespace).Get(clusterName, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}

		if tc.PDIsUpgrading() || tc.TiKVIsUpgrading() || tc.TiDBIsUpgrading() {
			return true, nil
		}

		return false, nil
	})
	Expect(err).NotTo(HaveOccurred())

	By("Then ensure upgrade order is pd -> tikv -> tidb")
	err = wait.Poll(100*time.Millisecond, 30*time.Minute, func() (bool, error) {
		tc, err := f.Cli.PingcapV1().TidbClusters(e2eNamespace).Get(clusterName, metav1.GetOptions{})
		if err != nil {
			return false, nil
		}
		if tc.PDIsUpgrading() {
			if tc.TiKVIsUpgrading() {
				return true, errors.New("pd and tikv are upgrading at the same time")
			}

			if tc.TiDBIsUpgrading() {
				return true, errors.New("pd and tidb are upgrading at the same time")
			}
		}

		if tc.TiKVIsUpgrading() {
			if tc.TiDBIsUpgrading() {
				return true, errors.New("tikv and tidb are upgrading at the same time")
			}
		}

		pdSet, err := f.Cli.PingcapV1().TidbSets(e2eNamespace).Get(fmt.Sprintf("%s-pdset", clusterName), metav1.GetOptions{})
		if err != nil {
			return false, nil
		}
		tikvSet, err := f.Cli.PingcapV1().TidbSets(e2eNamespace).Get(fmt.Sprintf("%s-tikvset", clusterName), metav1.GetOptions{})
		if err != nil {
			return false, nil
		}
		tidbSet, err := f.Cli.PingcapV1().TidbSets(e2eNamespace).Get(fmt.Sprintf("%s-tidbset", clusterName), metav1.GetOptions{})
		if err != nil {
			return false, nil
		}

		pdShouldUpgrade := tc.Spec.PD.Image != pdSet.ImageFor(v1.PDContainerType)
		tikvShouldUpgrade := tc.Spec.TiKV.Image != tikvSet.ImageFor(v1.TiKVContainerType)
		tidbShouldUpgrade := tc.Spec.TiDB.Image != tidbSet.ImageFor(v1.TiDBContainerType)

		if !pdShouldUpgrade && !tikvShouldUpgrade && !tidbShouldUpgrade &&
			!tc.PDIsUpgrading() && !tc.TiKVIsUpgrading() && !tc.TiDBIsUpgrading() {
			return true, nil
		}

		return false, nil
	})
	Expect(err).NotTo(HaveOccurred())

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