package tikvrel

import (
	"context"
	"errors"
	"io"
	"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/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 OLTP_AZ_REL_BASELINE(
	ctx context.Context,
	tc resource.TiDBCluster,
	benchToolset resource.WorkloadNode,
	pr *io.PipeReader,
	caseCmdParms oltppkg.CaseCmdParms,
) {
	type testParams struct {
		faultTpye     string
		faultDuration time.Duration
	}

	testCasesParams := []testParams{
		testParams{faultTpye: "network-partition-between-2-az", faultDuration: 10 * time.Minute},
		testParams{faultTpye: "network-partition-between-1-to-others-az", faultDuration: 10 * time.Minute},
		testParams{faultTpye: "network-delay-between-2-az-latency:2ms", faultDuration: 10 * time.Minute},
		testParams{faultTpye: "network-delay-between-2-az-latency:50ms", faultDuration: 2 * time.Minute},
		testParams{faultTpye: "network-delay-between-1-to-others-az-latency:2ms", faultDuration: 10 * time.Minute},
		testParams{faultTpye: "network-delay-between-1-to-others-az-latency:50ms", faultDuration: 2 * time.Minute},
	}
	warmUpTime := 30 * time.Minute
	recoveryTime := 30 * time.Minute
	waitForRaftLogCatchUp := 5 * time.Hour
	waitForBalance := 5 * time.Hour

	oltppkg.LogAndBy("Waiting for raft log lag to be less than a certain value before test.")
	oltppkg.WaitForRaftLogLagLessThan(ctx, tc, waitForRaftLogCatchUp, 2*time.Minute, 1000)

	oltppkg.LogAndBy("Waiting for balance before test.")
	oltppkg.WaitForBalance(ctx, tc, waitForBalance, 5*time.Minute, 0.1)

	for _, testCase := range testCasesParams {
		var c1 chaos.Chaos
		var c2 chaos.Chaos

		if strings.Contains(testCase.faultTpye, "network-partition-between-2-az") {
			targetSelectOpts := chaos.SelectAllTiKVInOneAZ("zone2")
			c1 = chaos.Chaos{
				Name:               "Network partition between zone1 and zone2",
				Tc:                 tc,
				TargetTc:           tc,
				ExperimentSpec:     &resource.NetworkPartitionSpec{},
				SelectorOpts:       chaos.SelectAllTiKVInOneAZ("zone1"),
				TargetSelectorOpts: &targetSelectOpts,
			}
		} else if strings.Contains(testCase.faultTpye, "network-partition-between-1-to-others-az") {
			targetSelectOpts := chaos.SelectAllTiKVInOneAZ("zone3")
			c1 = chaos.Chaos{
				Name:               "Network partition between zone1 and zone3",
				Tc:                 tc,
				TargetTc:           tc,
				ExperimentSpec:     &resource.NetworkPartitionSpec{},
				SelectorOpts:       chaos.SelectAllTiKVInOneAZ("zone1"),
				TargetSelectorOpts: &targetSelectOpts,
			}
			c2 = chaos.Chaos{
				Name:               "Network partition between zone2 and zone3",
				Tc:                 tc,
				TargetTc:           tc,
				ExperimentSpec:     &resource.NetworkPartitionSpec{},
				SelectorOpts:       chaos.SelectAllTiKVInOneAZ("zone2"),
				TargetSelectorOpts: &targetSelectOpts,
			}

		} else if strings.Contains(testCase.faultTpye, "network-delay-between-2-az") {
			targetSelectOpts := chaos.SelectAllTiKVInOneAZ("zone2")
			c1 = chaos.Chaos{
				Name:     "Network partition between zone1 and zone2",
				Tc:       tc,
				TargetTc: tc,
				ExperimentSpec: &resource.NetworkLatencySpec{
					Latency:     strings.TrimPrefix(testCase.faultTpye, "network-delay-between-2-az-latency:"),
					Correlation: "25",
				},
				SelectorOpts:       chaos.SelectAllTiKVInOneAZ("zone1"),
				TargetSelectorOpts: &targetSelectOpts,
			}
		} else if strings.Contains(testCase.faultTpye, "network-delay-between-1-to-others-az") {
			targetSelectOpts := chaos.SelectAllTiKVInOneAZ("zone3")
			c1 = chaos.Chaos{
				Name:     "Network partition between zone1 and zone3",
				Tc:       tc,
				TargetTc: tc,
				ExperimentSpec: &resource.NetworkLatencySpec{
					Latency:     strings.TrimPrefix(testCase.faultTpye, "network-delay-between-1-to-others-az-latency:"),
					Correlation: "25",
				},
				SelectorOpts:       chaos.SelectAllTiKVInOneAZ("zone1"),
				TargetSelectorOpts: &targetSelectOpts,
			}
			c2 = chaos.Chaos{
				Name:     "Network partition between zone2 and zone3",
				Tc:       tc,
				TargetTc: tc,
				ExperimentSpec: &resource.NetworkLatencySpec{
					Latency:     strings.TrimPrefix(testCase.faultTpye, "network-delay-between-1-to-others-az-latency:"),
					Correlation: "25",
				},
				SelectorOpts:       chaos.SelectAllTiKVInOneAZ("zone2"),
				TargetSelectorOpts: &targetSelectOpts,
			}
		} else {
			Expect(errors.New("invalid faultType")).Should(Succeed())
		}

		maxRunTime := warmUpTime + testCase.faultDuration + recoveryTime + waitForRaftLogCatchUp + waitForBalance
		ctxAll, cancel := context.WithTimeout(ctx, maxRunTime)
		defer cancel()
		g, gctx := errgroup.WithContext(ctxAll)
		g.Go(WithGinkgoRecover(func() {
			err := oltppkg.WorkLoadErrLogCheck(gctx, pr, maxRunTime, false, []string{})
			Expect(err).Should(Succeed())
		}))

		oltppkg.WaitFor(gctx, "Warm Up before inject fault", warmUpTime)
		log.Info("faultInfo",
			zap.String("fault-type", testCase.faultTpye),
			zap.Duration("faultDuration", testCase.faultDuration))
		g.Go(WithGinkgoRecover(func() {
			err := oltppkg.InjectFaultByChaos(&c1, gctx, testCase.faultDuration)
			Expect(err).Should(Succeed())

			oltppkg.LogAndBy("Waiting for raft log lag to be less than a certain value after recover fault.")
			oltppkg.WaitForRaftLogLagLessThan(gctx, tc, waitForRaftLogCatchUp, 2*time.Minute, 1000)

			oltppkg.LogAndBy("Waiting for balance.")
			oltppkg.WaitForBalance(gctx, tc, waitForBalance, 5*time.Minute, 0.1)

			oltppkg.WaitFor(gctx, "Recovery observation period", recoveryTime)
			cancel()
		}))
		if c2.TargetSelectorOpts != nil {
			g.Go(WithGinkgoRecover(func() {
				err := oltppkg.InjectFaultByChaos(&c2, gctx, testCase.faultDuration)
				Expect(err).Should(Succeed())
			}))
		}
		Expect(g.Wait()).Should(BeNil())
	}
}