// Copyright 2018 TiKV Project Authors.
//
// 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 tests

import (
	"context"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/coreos/go-semver/semver"
	"github.com/pingcap/errors"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/pingcap/log"
	"github.com/tikv/pd/pkg/dashboard"
	"github.com/tikv/pd/pkg/swaggerserver"
	"github.com/tikv/pd/server"
	"github.com/tikv/pd/server/api"
	"github.com/tikv/pd/server/cluster"
	"github.com/tikv/pd/server/config"
	"github.com/tikv/pd/server/core"
	"github.com/tikv/pd/server/id"
	"github.com/tikv/pd/server/join"
	"github.com/tikv/pd/server/member"
	"go.etcd.io/etcd/clientv3"
	"go.uber.org/zap"
)

// TestServer states.
const (
	Initial int32 = iota
	Running
	Stop
	Destroy
)

var (
	// WaitLeaderReturnDelay represents the time interval of WaitLeader sleep before returning.
	WaitLeaderReturnDelay = 20 * time.Millisecond
	// WaitLeaderCheckInterval represents the time interval of WaitLeader running check.
	WaitLeaderCheckInterval = 500 * time.Millisecond
)

// TestServer is only for test.
type TestServer struct {
	sync.RWMutex
	server *server.Server
	state  int32
}

var zapLogOnce sync.Once

// NewTestServer creates a new TestServer.
func NewTestServer(ctx context.Context, cfg *config.Config) (*TestServer, error) {
	err := cfg.SetupLogger()
	if err != nil {
		return nil, err
	}
	zapLogOnce.Do(func() {
		log.ReplaceGlobals(cfg.GetZapLogger(), cfg.GetZapLogProperties())
	})
	err = join.PrepareJoinCluster(cfg)
	if err != nil {
		return nil, err
	}
	serviceBuilders := []server.HandlerBuilder{api.NewHandler, swaggerserver.NewHandler}
	serviceBuilders = append(serviceBuilders, dashboard.GetServiceBuilders()...)
	svr, err := server.CreateServer(ctx, cfg, serviceBuilders...)
	if err != nil {
		return nil, err
	}
	return &TestServer{
		server: svr,
		state:  Initial,
	}, nil
}

// Run starts to run a TestServer.
func (s *TestServer) Run() error {
	s.Lock()
	defer s.Unlock()
	if s.state != Initial && s.state != Stop {
		return errors.Errorf("server(state%d) cannot run", s.state)
	}
	if err := s.server.Run(); err != nil {
		return err
	}
	s.state = Running
	return nil
}

// Stop is used to stop a TestServer.
func (s *TestServer) Stop() error {
	s.Lock()
	defer s.Unlock()
	if s.state != Running {
		return errors.Errorf("server(state%d) cannot stop", s.state)
	}
	s.server.Close()
	s.state = Stop
	return nil
}

// Destroy is used to destroy a TestServer.
func (s *TestServer) Destroy() error {
	s.Lock()
	defer s.Unlock()
	if s.state == Running {
		s.server.Close()
	}
	if err := os.RemoveAll(s.server.GetConfig().DataDir); err != nil {
		return err
	}
	s.state = Destroy
	return nil
}

// ResignLeader resigns the leader of the server.
func (s *TestServer) ResignLeader() error {
	s.Lock()
	defer s.Unlock()
	return s.server.GetMember().ResignLeader(s.server.Context(), s.server.Name(), "")
}

// State returns the current TestServer's state.
func (s *TestServer) State() int32 {
	s.RLock()
	defer s.RUnlock()
	return s.state
}

// GetConfig returns the current TestServer's configuration.
func (s *TestServer) GetConfig() *config.Config {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetConfig()
}

// GetPersistOptions returns the current TestServer's schedule option.
func (s *TestServer) GetPersistOptions() *config.PersistOptions {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetPersistOptions()
}

// GetAllocator returns the current TestServer's ID allocator.
func (s *TestServer) GetAllocator() *id.AllocatorImpl {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetAllocator()
}

// GetAddr returns the address of TestCluster.
func (s *TestServer) GetAddr() string {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetAddr()
}

// GetServer returns the real server of TestServer.
func (s *TestServer) GetServer() *server.Server {
	s.RLock()
	defer s.RUnlock()
	return s.server
}

// GetClusterID returns the cluster ID.
func (s *TestServer) GetClusterID() uint64 {
	s.RLock()
	defer s.RUnlock()
	return s.server.ClusterID()
}

// GetLeader returns current leader of PD cluster.
func (s *TestServer) GetLeader() *pdpb.Member {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetLeader()
}

// GetCluster returns PD cluster.
func (s *TestServer) GetCluster() *metapb.Cluster {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetCluster()
}

// GetClusterVersion returns PD cluster version.
func (s *TestServer) GetClusterVersion() semver.Version {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetClusterVersion()
}

// GetServerID returns the unique etcd ID for this server in etcd cluster.
func (s *TestServer) GetServerID() uint64 {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetMember().ID()
}

// IsLeader returns whether the server is leader or not.
func (s *TestServer) IsLeader() bool {
	s.RLock()
	defer s.RUnlock()
	return !s.server.IsClosed() && s.server.GetMember().IsLeader()
}

// GetEtcdLeader returns the builtin etcd leader.
func (s *TestServer) GetEtcdLeader() (string, error) {
	s.RLock()
	defer s.RUnlock()
	req := &pdpb.GetMembersRequest{Header: &pdpb.RequestHeader{ClusterId: s.server.ClusterID()}}
	members, err := s.server.GetMembers(context.TODO(), req)
	if err != nil {
		return "", errors.WithStack(err)
	}
	return members.GetEtcdLeader().GetName(), nil
}

// GetEtcdLeaderID returns the builtin etcd leader ID.
func (s *TestServer) GetEtcdLeaderID() (uint64, error) {
	s.RLock()
	defer s.RUnlock()
	req := &pdpb.GetMembersRequest{Header: &pdpb.RequestHeader{ClusterId: s.server.ClusterID()}}
	members, err := s.server.GetMembers(context.TODO(), req)
	if err != nil {
		return 0, errors.WithStack(err)
	}
	return members.GetEtcdLeader().GetMemberId(), nil
}

// MoveEtcdLeader moves etcd leader from old to new.
func (s *TestServer) MoveEtcdLeader(old, new uint64) error {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetMember().MoveEtcdLeader(context.Background(), old, new)
}

// GetEtcdClient returns the builtin etcd client.
func (s *TestServer) GetEtcdClient() *clientv3.Client {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetClient()
}

// GetHTTPClient returns the builtin http client.
func (s *TestServer) GetHTTPClient() *http.Client {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetHTTPClient()
}

// GetStores returns the stores of the cluster.
func (s *TestServer) GetStores() []*metapb.Store {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetRaftCluster().GetMetaStores()
}

// GetStore returns the store with a given store ID.
func (s *TestServer) GetStore(storeID uint64) *core.StoreInfo {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetRaftCluster().GetStore(storeID)
}

// GetRaftCluster returns Raft cluster.
// If cluster has not been bootstrapped, return nil.
func (s *TestServer) GetRaftCluster() *cluster.RaftCluster {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetRaftCluster()
}

// GetRegions returns all regions' information in detail.
func (s *TestServer) GetRegions() []*core.RegionInfo {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetRaftCluster().GetRegions()
}

// GetRegionInfoByID returns regionInfo by regionID from cluster.
func (s *TestServer) GetRegionInfoByID(regionID uint64) *core.RegionInfo {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetRaftCluster().GetRegion(regionID)
}

// GetAdjacentRegions returns regions' information that are adjacent with the specific region ID.
func (s *TestServer) GetAdjacentRegions(region *core.RegionInfo) []*core.RegionInfo {
	s.RLock()
	defer s.RUnlock()
	left, right := s.server.GetRaftCluster().GetAdjacentRegions(region)
	return []*core.RegionInfo{left, right}
}

// GetStoreRegions returns all regions' information with a given storeID.
func (s *TestServer) GetStoreRegions(storeID uint64) []*core.RegionInfo {
	s.RLock()
	defer s.RUnlock()
	return s.server.GetRaftCluster().GetStoreRegions(storeID)
}

// BootstrapCluster is used to bootstrap the cluster.
func (s *TestServer) BootstrapCluster() error {
	bootstrapReq := &pdpb.BootstrapRequest{
		Header: &pdpb.RequestHeader{ClusterId: s.GetClusterID()},
		Store:  &metapb.Store{Id: 1, Address: "mock://1"},
		Region: &metapb.Region{Id: 2, Peers: []*metapb.Peer{{Id: 3, StoreId: 1, IsLearner: false}}},
	}
	_, err := s.server.Bootstrap(context.Background(), bootstrapReq)
	if err != nil {
		return err
	}
	return nil
}

// WaitLease is used to get leader lease.
// If it exceeds the maximum number of loops, it will return nil.
func (s *TestServer) WaitLease() *member.LeaderLease {
	for i := 0; i < 100; i++ {
		lease := s.server.GetLease()
		if lease != nil {
			return lease
		}
		time.Sleep(WaitLeaderCheckInterval)
	}
	return nil
}

// TestCluster is only for test.
type TestCluster struct {
	config  *clusterConfig
	servers map[string]*TestServer
}

// ConfigOption is used to define customize settings in test.
type ConfigOption func(conf *config.Config)

// NewTestCluster creates a new TestCluster.
func NewTestCluster(ctx context.Context, initialServerCount int, opts ...ConfigOption) (*TestCluster, error) {
	config := newClusterConfig(initialServerCount)
	servers := make(map[string]*TestServer)
	for _, conf := range config.InitialServers {
		serverConf, err := conf.Generate(opts...)
		if err != nil {
			return nil, err
		}
		s, err := NewTestServer(ctx, serverConf)
		if err != nil {
			return nil, err
		}
		servers[conf.Name] = s
	}
	return &TestCluster{
		config:  config,
		servers: servers,
	}, nil
}

// RunServer starts to run TestServer.
func (c *TestCluster) RunServer(server *TestServer) <-chan error {
	resC := make(chan error)
	go func() { resC <- server.Run() }()
	return resC
}

// RunServers starts to run multiple TestServer.
func (c *TestCluster) RunServers(servers []*TestServer) error {
	res := make([]<-chan error, len(servers))
	for i, s := range servers {
		res[i] = c.RunServer(s)
	}
	for _, c := range res {
		if err := <-c; err != nil {
			return errors.WithStack(err)
		}
	}
	return nil
}

// RunInitialServers starts to run servers in InitialServers.
func (c *TestCluster) RunInitialServers() error {
	servers := make([]*TestServer, 0, len(c.config.InitialServers))
	for _, conf := range c.config.InitialServers {
		servers = append(servers, c.GetServer(conf.Name))
	}
	return c.RunServers(servers)
}

// StopAll is used to stop all servers.
func (c *TestCluster) StopAll() error {
	for _, s := range c.servers {
		if err := s.Stop(); err != nil {
			return err
		}
	}
	return nil
}

// GetServer returns a server with a given name.
func (c *TestCluster) GetServer(name string) *TestServer {
	return c.servers[name]
}

// GetServers returns all servers.
func (c *TestCluster) GetServers() map[string]*TestServer {
	return c.servers
}

// GetLeader returns the leader of all servers
func (c *TestCluster) GetLeader() string {
	for name, s := range c.servers {
		if s.IsLeader() {
			return name
		}
	}
	return ""
}

// GetFollower returns an follower of all servers
func (c *TestCluster) GetFollower() string {
	for name, s := range c.servers {
		if !s.server.IsClosed() && !s.server.GetMember().IsLeader() {
			return name
		}
	}
	return ""
}

// WaitLeader is used to get leader.
// If it exceeds the maximum number of loops, it will return an empty string.
func (c *TestCluster) WaitLeader() string {
	for i := 0; i < 100; i++ {
		counter := make(map[string]int)
		running := 0
		for _, s := range c.servers {
			if s.state == Running {
				running++
			}
			n := s.GetLeader().GetName()
			if n != "" {
				counter[n]++
			}
		}
		for name, num := range counter {
			if num == running && c.GetServer(name).IsLeader() {
				time.Sleep(WaitLeaderReturnDelay)
				return name
			}
		}
		time.Sleep(WaitLeaderCheckInterval)
	}
	return ""
}

// ResignLeader resigns the leader of the cluster.
func (c *TestCluster) ResignLeader() error {
	leader := c.GetLeader()
	if len(leader) != 0 {
		return c.servers[leader].ResignLeader()
	}
	return errors.New("no leader")
}

// GetCluster returns PD cluster.
func (c *TestCluster) GetCluster() *metapb.Cluster {
	leader := c.GetLeader()
	return c.servers[leader].GetCluster()
}

// GetEtcdClient returns the builtin etcd client.
func (c *TestCluster) GetEtcdClient() *clientv3.Client {
	leader := c.GetLeader()
	return c.servers[leader].GetEtcdClient()
}

// GetHTTPClient returns the builtin http client.
func (c *TestCluster) GetHTTPClient() *http.Client {
	leader := c.GetLeader()
	return c.servers[leader].GetHTTPClient()
}

// GetConfig returns the current TestCluster's configuration.
func (c *TestCluster) GetConfig() *clusterConfig {
	return c.config
}

// HandleRegionHeartbeat processes RegionInfo reports from the client.
func (c *TestCluster) HandleRegionHeartbeat(region *core.RegionInfo) error {
	leader := c.GetLeader()
	cluster := c.servers[leader].GetRaftCluster()
	return cluster.HandleRegionHeartbeat(region)
}

// Join is used to add a new TestServer into the cluster.
func (c *TestCluster) Join(ctx context.Context) (*TestServer, error) {
	conf, err := c.config.Join().Generate()
	if err != nil {
		return nil, err
	}
	s, err := NewTestServer(ctx, conf)
	if err != nil {
		return nil, err
	}
	c.servers[conf.Name] = s
	return s, nil
}

// Destroy is used to destroy a TestCluster.
func (c *TestCluster) Destroy() {
	for _, s := range c.servers {
		err := s.Destroy()
		if err != nil {
			log.Error("failed to destroy the cluster:", zap.Error(err))
		}
	}
}