package qhad

import (
	"context"
	"time"

	"github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/endless/testcase/oltp/scenario/params"
)

func InitializeScale001Params(ctx context.Context, tidbClient *tidb.Client) (*params.Scale, error) {
	current, err := tidbClient.GetAllTypeInstanceNum(ctx)
	if err != nil {
		return nil, err
	}

	params := params.Scale{
		RunTime:    1 * time.Hour,
		WarmUpTime: 7 * time.Minute,

		BalanceTimeout:   3 * time.Hour,
		BalanceThreshold: 0.2,

		StoreLimit:        30,
		DefaultStoreLimit: 15,
	}

	// Calculate the scale out number.
	var target map[string]uint32
	for name, count := range current {
		if name == "tidb" {
			target[name] = count + 1
		} else if name == "pd" {
			target[name] = count + 1
		} else if name == "tikv" {
			target[name] = count + 1
		}
	}
	params.ScaleOutTargetNumber = target

	return &params, nil
}

func InitializeScale002Params(ctx context.Context, tidbClient *tidb.Client) (*params.Scale, error) {
	current, err := tidbClient.GetAllTypeInstanceNum(ctx)
	if err != nil {
		return nil, err
	}

	params := params.Scale{
		RunTime:    1 * time.Hour,
		WarmUpTime: 7 * time.Minute,

		BalanceTimeout:   3 * time.Hour,
		BalanceThreshold: 0.2,

		StoreLimit:        30,
		DefaultStoreLimit: 15,
	}

	// Calculate the scale out number.
	var target map[string]uint32
	for name, count := range current {
		if name == "tidb" {
			target[name] = count - 1
		} else if name == "pd" {
			target[name] = count - 1
		} else if name == "tikv" {
			target[name] = count - 1
		}
	}
	params.ScaleOutTargetNumber = target

	return &params, nil
}