package tikvrel

import (
	"context"
	"errors"
	"io"
	"strconv"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
	. "github.com/onsi/gomega"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"

	"github.com/pingcap/endless/pkg/chaos"
	. "github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/pkg/matcher/metrics"
	. "github.com/pingcap/endless/pkg/util"
	"github.com/pingcap/endless/testcase/oltp/oltppkg"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s"
)

func (tikvRelOpts *TiKVRelOptions) NetworkPartitionTest(
	ctx context.Context,
	tc resource.TiDBCluster,
	pr *io.PipeReader,
	caseCmdParms oltppkg.CaseCmdParms,
) {
	oltppkg.LogAndBy("Network Partition Testing")
	tikvIdList := oltppkg.GetInstanceIdList(tc, resource.TiKV)
	tidbIdList := oltppkg.GetInstanceIdList(tc, resource.TiDB)
	pdMembers := oltppkg.GetPdMembersName(tc)
	pdLeader := oltppkg.GetPdLeaderName(tc)
	pdLeaderIndex, _ := strconv.Atoi(strings.TrimPrefix(pdLeader, "tc-pd-"))
	var pdFollowersIdList []int
	for _, pdMember := range pdMembers {
		pdMemberIndex, _ := strconv.Atoi(strings.TrimPrefix(pdMember, "tc-pd-"))
		if pdMemberIndex != pdLeaderIndex {
			pdFollowersIdList = append(pdFollowersIdList, pdMemberIndex)
		}
	}
	log.Info("PeersList",
		zap.Ints("tikvIdList: ", tikvIdList),
		zap.Ints("tidbIdList: ", tidbIdList),
		zap.Strings("pdMembers: ", pdMembers),
		zap.String("pdLeader: ", pdLeader),
		zap.Ints("pdFollowersIdList: ", pdFollowersIdList))

	var c1 chaos.Chaos
	var c2 chaos.Chaos
	var experimentSpec resource.ExperimentSpec
	var chaos1Peers = make(map[string][]string)
	var chaos2Peers = make(map[string][]string)

	switch tikvRelOpts.FaultType {
	case "network-partition-between-1IDC-and-others-flash":
		experimentSpec = &resource.NetworkPartitionSpec{
			Duration:  oltppkg.GetRandIntStr(3, 5) + "s",
			Scheduler: &resource.ExperimentScheduler{Cron: "@every 10m"},
		}
		chaos1Peers = map[string][]string{
			"selectorPeers": {"tc-pd-0", "tc-tikv-0", "tc-tidb-0"},
			"targetPeers":   {"tc-pd-1", "tc-tikv-1", "tc-tidb-1"},
		}
		chaos2Peers = map[string][]string{
			"selectorPeers": {"tc-pd-0", "tc-tikv-0", "tc-tidb-0"},
			"targetPeers":   {"tc-pd-2", "tc-tikv-2"},
		}
	case "network-partition-between-2-tikv":
		if tikvRelOpts.FaultDuration < 10*time.Second {
			experimentSpec = &resource.NetworkPartitionSpec{
				Duration:  strings.Split((tikvRelOpts.FaultDuration).String(), "s")[0] + "s",
				Scheduler: &resource.ExperimentScheduler{Cron: "@every 60m"},
			}
			tikvRelOpts.FaultDuration = 10 * time.Minute
		} else {
			experimentSpec = &resource.NetworkPartitionSpec{}
		}

		chaos1Peers = map[string][]string{
			"selectorPeers": {"tc-tikv-0"},
			"targetPeers":   {"tc-tikv-1"},
		}
	case "network-partition-between-2IDC":
		experimentSpec = &resource.NetworkPartitionSpec{}
		chaos1Peers = map[string][]string{
			"selectorPeers": {"tc-pd-0", "tc-tikv-0", "tc-tidb-0"},
			"targetPeers":   {"tc-pd-1", "tc-tikv-1", "tc-tidb-1"},
		}
	case "network-partition-between-1IDC-and-others":
		experimentSpec = &resource.NetworkPartitionSpec{}
		chaos1Peers = map[string][]string{
			"selectorPeers": {"tc-pd-0", "tc-tikv-0", "tc-tidb-0"},
			"targetPeers":   {"tc-pd-1", "tc-tikv-1", "tc-tidb-1"},
		}
		chaos2Peers = map[string][]string{
			"selectorPeers": {"tc-pd-0", "tc-tikv-0", "tc-tidb-0"},
			"targetPeers":   {"tc-pd-2", "tc-tikv-2"},
		}
	default:
		err := errors.New("invalid faultType")
		Expect(err).Should(Succeed())
	}
	c1 = Try(InitChaosByPeerName(tc, tikvRelOpts.FaultType, experimentSpec, chaos1Peers)).(chaos.Chaos)
	if len(chaos2Peers["selectorPeers"]) != 0 {
		c2 = Try(InitChaosByPeerName(tc, tikvRelOpts.FaultType, experimentSpec, chaos2Peers)).(chaos.Chaos)
	}

	waitForRaftLogCatchUpTime := 2 * time.Hour
	waitForBalanceTime := 3 * time.Hour
	runTime, recoveryCheckOffsetStr := tikvRelOpts.GetRunTimeAndCheckOffset(ctx, tc, caseCmdParms)
	maxRunTime := runTime + waitForRaftLogCatchUpTime + waitForBalanceTime
	ctxAll, cancel := context.WithTimeout(ctx, maxRunTime)
	defer cancel()
	g, gctx := errgroup.WithContext(ctxAll)

	tikvRelOpts.CommonCheckDuringCaseRunning(gctx, tc, g, maxRunTime, pr, caseCmdParms)

	oltppkg.WaitFor(gctx, "Warm Up before inject fault", tikvRelOpts.WarmUpTime)

	g.Go(WithGinkgoRecover(func() {
		err := oltppkg.InjectFaultByChaos(&c1, gctx, tikvRelOpts.FaultDuration)
		Expect(err).Should(Succeed())
		log.Info("Recovery after inject fault", zap.Duration("recovery-time: ", tikvRelOpts.RecoveryTime))
	}))

	if c2.TargetSelectorOpts != nil {
		g.Go(WithGinkgoRecover(func() {
			err := oltppkg.InjectFaultByChaos(&c2, gctx, tikvRelOpts.FaultDuration)
			Expect(err).Should(Succeed())
		}))
	}

	g.Go(WithGinkgoRecover(func() {
		oltppkg.WaitFor(gctx, "Wait QPS recover after inject fault.", tikvRelOpts.RTO)
		// know issue https://github.com/tikv/tikv/issues/12259 and incall 106, so change QPS jitter duration from 30s to 300s
		ConsistentlyWithMetrics(tc, maxRunTime).
			ShouldNot(QPSDropMoreThanOffset(0.5, recoveryCheckOffsetStr).For(300 * time.Second).WithCtx(gctx))
	}))

	oltppkg.WaitFor(gctx, "wait for fault recover.", tikvRelOpts.FaultDuration)

	tikvRelOpts.WaitClusterStatusRecover(gctx, tc, waitForRaftLogCatchUpTime, waitForBalanceTime)
	oltppkg.WaitFor(gctx, "Recovery observation period", tikvRelOpts.RecoveryTime)
	cancel()

	Expect(g.Wait()).Should(BeNil())
}

func InitChaosByPeerName(tc resource.TiDBCluster, name string, experimentSpec resource.ExperimentSpec, chaosPeers map[string][]string) (chaos.Chaos, error) {
	var c chaos.Chaos
	c.Tc = tc
	c.TargetTc = tc
	c.Name = name
	c.ExperimentSpec = experimentSpec
	var selectorOpts resource.TiDBClusterSelectorOpts
	var targetOpts resource.TiDBClusterSelectorOpts
	var pdSelectorSpec resource.TiDBClusterPDSelectorSpec
	var tikvSelectorSpec resource.TiDBClusterTiKVSelectorSpec
	var tidbSelectorSpec resource.TiDBClusterTiDBSelectorSpec
	var pdTargetSpec resource.TiDBClusterPDSelectorSpec
	var tikvTargetSpec resource.TiDBClusterTiKVSelectorSpec
	var tidbTargetSpec resource.TiDBClusterTiDBSelectorSpec
	log.Info("Peers",
		zap.Strings("selectorPeers: ", chaosPeers["selectorPeers"]),
		zap.Strings("targetPeers: ", chaosPeers["targetPeers"]),
		zap.String("fault type: ", name))

	for _, selectorPeer := range chaosPeers["selectorPeers"] {
		if strings.Contains(selectorPeer, "pd") {
			index := Try(strconv.Atoi(strings.TrimPrefix(selectorPeer, "tc-pd-"))).(int)
			pdSelectorSpec.Peer = &index
			selectorOpts.TiDBClusterPDSelectorSpec = &pdSelectorSpec
		} else if strings.Contains(selectorPeer, "tikv") {
			index := Try(strconv.Atoi(strings.TrimPrefix(selectorPeer, "tc-tikv-"))).(int)
			tikvSelectorSpec.Peer = &index
			selectorOpts.TiDBClusterTiKVSelectorSpec = &tikvSelectorSpec
		} else if strings.Contains(selectorPeer, "tidb") {
			index := Try(strconv.Atoi(strings.TrimPrefix(selectorPeer, "tc-tidb-"))).(int)
			tidbSelectorSpec.Peer = &index
			selectorOpts.TiDBClusterTiDBSelectorSpec = &tidbSelectorSpec
		}
	}

	for _, targetPeer := range chaosPeers["targetPeers"] {
		if strings.Contains(targetPeer, "pd") {
			index := Try(strconv.Atoi(strings.TrimPrefix(targetPeer, "tc-pd-"))).(int)
			pdTargetSpec.Peer = &index
			targetOpts.TiDBClusterPDSelectorSpec = &pdTargetSpec
		} else if strings.Contains(targetPeer, "tikv") {
			index := Try(strconv.Atoi(strings.TrimPrefix(targetPeer, "tc-tikv-"))).(int)
			tikvTargetSpec.Peer = &index
			targetOpts.TiDBClusterTiKVSelectorSpec = &tikvTargetSpec
		} else if strings.Contains(targetPeer, "tidb") {
			index := Try(strconv.Atoi(strings.TrimPrefix(targetPeer, "tc-tidb-"))).(int)
			tidbTargetSpec.Peer = &index
			targetOpts.TiDBClusterTiDBSelectorSpec = &tidbTargetSpec
		} else {
			return c, errors.New("invalid peer")
		}
	}
	c.SelectorOpts = selectorOpts
	c.TargetSelectorOpts = &targetOpts
	return c, nil
}