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

import (
	"fmt"
	"strings"

	"github.com/juju/errors"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap.com/v1"
	"github.com/pingcap/tidb-operator/pkg/client/clientset/versioned"
	tsctl "github.com/pingcap/tidb-operator/pkg/tidbset/controller"
	"github.com/pingcap/tidb-operator/pkg/util/label"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

const (
	// OperatorName is the name of tidb-operator
	OperatorName = "tidb-operator"

	// SchedulerName is the  name of tidb-scheduler
	SchedulerName = "tidb-scheduler"

	// VolumeManagerName is the  name of tidb-volume-manager
	VolumeManagerName = "tidb-volume-manager"

	// LabelNodeRoleMaster specifies that a node is a master
	LabelNodeRoleMaster = "node-role.kubernetes.io/master"
)

// Framework supports common operations used by e2e tests
type Framework struct {
	BaseName string
	Cli      versioned.Interface
	KubeCli  kubernetes.Interface
}

// New returns a new framework and sets up a client for you
func New(baseName string) *Framework {
	f := &Framework{
		BaseName: baseName,
	}

	BeforeEach(f.BeforeEach)
	AfterEach(f.AfterEach)

	return f
}

// BeforeEach gets a client
func (f *Framework) BeforeEach() {
	By("Creating a client")

	cfg, err := rest.InClusterConfig()
	Expect(err).NotTo(HaveOccurred())

	cfg.QPS = 1000
	cfg.Burst = 1000

	cli, err := versioned.NewForConfig(cfg)
	Expect(err).NotTo(HaveOccurred())

	kubeCli, err := kubernetes.NewForConfig(cfg)
	Expect(err).NotTo(HaveOccurred())

	f.Cli = cli
	f.KubeCli = kubeCli
}

// AfterEach does nothing now
func (f *Framework) AfterEach() {
}

// GetAllClusterPods gets all pods of a TidbCluster, include: pd, tidb, tikv
func (f *Framework) GetAllClusterPods(ns, clusterName string) (*corev1.PodList, error) {
	return f.KubeCli.CoreV1().Pods(ns).List(label.ClusterListOptions(clusterName))
}

// GetAllClusterPVs gets all PVs of a TidbCluster
func (f *Framework) GetAllClusterPVs(clusterName string) (*corev1.PersistentVolumeList, error) {
	return f.KubeCli.CoreV1().PersistentVolumes().List(label.ClusterListOptions(clusterName))
}

// GetAllClusterPVCs gets all PVCs of a TidbCluster
func (f *Framework) GetAllClusterPVCs(ns, clusterName string) (*corev1.PersistentVolumeClaimList, error) {
	return f.KubeCli.CoreV1().PersistentVolumeClaims(ns).List(label.ClusterListOptions(clusterName))
}

// GetAllClusterServices gets all services of a TidbCluster
func (f *Framework) GetAllClusterServices(ns, clusterName string) (*corev1.ServiceList, error) {
	return f.KubeCli.CoreV1().Services(ns).List(label.ClusterListOptions(clusterName))
}

// GetAllClusterTidbSets gets all tidbSets of a TidbCluster
func (f *Framework) GetAllClusterTidbSets(ns, clusterName string) (*v1.TidbSetList, error) {
	return f.Cli.PingcapV1().TidbSets(ns).List(label.ClusterListOptions(clusterName))
}

// GetAllPDPods gets all PD pods of a TidbCluster
func (f *Framework) GetAllPDPods(ns, clusterName string) (*corev1.PodList, error) {
	return f.getAllPodsFor("pd", ns, clusterName)
}

// GetAllTiDBPods gets all TiDB pods of a TidbCluster
func (f *Framework) GetAllTiDBPods(ns, clusterName string) (*corev1.PodList, error) {
	return f.getAllPodsFor("tidb", ns, clusterName)
}

// GetAllTiKVPods gets all TiKV pods of a TidbCluster
func (f *Framework) GetAllTiKVPods(ns, clusterName string) (*corev1.PodList, error) {
	return f.getAllPodsFor("tikv", ns, clusterName)
}

func (f *Framework) getAllPodsFor(memberType, ns, clusterName string) (*corev1.PodList, error) {
	return f.KubeCli.CoreV1().Pods(ns).List(metav1.ListOptions{
		LabelSelector: labels.SelectorFromSet(
			label.New().Cluster(clusterName).App(memberType).Labels(),
		).String(),
	})
}

// IsPodRunning returns whether the pod is running and is ready in PD cluster
func (f *Framework) IsPodRunning(pod *corev1.Pod) (bool, error) {
	if pod.Status.Phase != corev1.PodRunning {
		return false, errors.Errorf("pod: %s is not running", f.PodName(pod))
	}

	var ownerRef *metav1.OwnerReference
	for _, or := range pod.OwnerReferences {
		if or.Kind == tsctl.ControllerKind.Kind {
			ownerRef = &or
		}
	}
	if ownerRef == nil {
		// other pods are not registered in pd, their status is indicated by pod.status.phase
		return true, nil
	}

	tidbset, err := f.Cli.PingcapV1().TidbSets(pod.GetNamespace()).Get(ownerRef.Name, metav1.GetOptions{})
	if err != nil {
		return false, err
	}

	if strings.HasSuffix(tidbset.Name, "-pdset") {
		return tidbset.IsReady(pod), nil
	}

	if strings.HasSuffix(tidbset.Name, "-tidbset") {
		return tidbset.IsReady(pod), nil
	}

	if strings.HasSuffix(tidbset.Name, "-tikvset") {
		return tidbset.IsReady(pod), nil
	}

	return true, nil
}

// PodName returns pod name with namespace
func (f *Framework) PodName(pod *corev1.Pod) string {
	return fmt.Sprintf("%s/%s", pod.GetNamespace(), pod.GetName())
}