package main

import (
	"path/filepath"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/tidbcloud/cluster-tier-config/script/model"
)

var (
	// unit: MB/s
	// get baseline from here: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ebs-optimized.html#current-general-purpose
	// Max Throughput*/Volume of gp3 is 1,000 MB/s, Refer to https://aws.amazon.com/ebs/general-purpose/, so we use 1000 as baseline if base of ec2 instance is more than 1000.
	instance2InstanceEBSThroughputBaseline = map[string]int{
		"t4g.medium":  43,
		"m6g.large":   79,
		"m6g.xlarge":  149,
		"m6g.2xlarge": 297,
		"m6g.4xlarge": 594,
		"m6g.8xlarge": 1000,
		"c6g.xlarge":  149,
		"c6g.2xlarge": 297,
		"c6g.4xlarge": 594,
		"c6g.8xlarge": 1000,
		"r6g.2xlarge": 297,
		"r6g.4xlarge": 594,
		"r6g.8xlarge": 1000,

		"t3a.medium":   44,
		"m5.large":     81,
		"m5.xlarge":    144,
		"m5.2xlarge":   288,
		"m5.4xlarge":   594,
		"r5.2xlarge":   288,
		"r5.4xlarge":   594,
		"r5.8xlarge":   850,
		"c5.xlarge":    144,
		"c5.2xlarge":   288,
		"c5.4xlarge":   594,
		"i3.2xlarge":   213,
		"i3.4xlarge":   438,
		"r6i.8xlarge":  1000, // in fact: 1250.0
		"m6i.8xlarge":  1000, // in fact: 1250.0
		"c6g.16xlarge": 125,  // in fact: 2375.0
		"m6g.16xlarge": 1000, // in fact: 2375.0
	}

	// whitelistTierUnitKeys defines which tier units should be ignored
	whitelistTierUnitKeys = map[string]bool{
		"aws:pd:m6g.large:gp3:3000:125":    true,
		"aws:pd:c6g.xlarge:gp3:3000:125":   true,
		"aws:tikv:m6g.large:gp3:3000:125":  true,
		"aws:tikv:m6g.xlarge:gp3:3000:144": true,

		// TODO(chenfei): remove these line after we migrate tier units for prod env.
		"aws:tikv:m6g.2xlarge:gp3:4000:288":     true,
		"aws:tikv:r6g.2xlarge:gp3:4000:288":     true,
		"aws:tiflash:r6g.2xlarge:gp3:12000:288": true,

		"aws:pd:m5.large:gp3:3000:125":    true,
		"aws:pd:c5.xlarge:gp3:3000:125":   true,
		"aws:pd:c5.2xlarge:gp3:3000:125":  true,
		"aws:pd:c5.4xlarge:gp3:3000:125":  true,
		"aws:pd:m5.2xlarge:gp3:3000:125":  true,
		"aws:pd:c6g.2xlarge:gp3:3000:125": true,
		"aws:pd:c6g.4xlarge:gp3:3000:125": true,
		"aws:pd:c6g.8xlarge:gp3:3000:125": true,
		"aws:pd:m6g.2xlarge:gp3:3000:125": true,
		"aws:tikv:m5.large:gp3:3000:125":  true,

		// dev test tier
		"aws:pd:t4g.medium:gp3:4000:400":       true,
		"aws:tikv:t4g.medium:gp3:4000:400":     true,
		"aws:tiflash:t4g.medium:gp3:6000:625":  true,
		"aws:pd:c6g.xlarge:gp3:4000:400":       true,
		"aws:pd:m6g.xlarge:gp3:4000:400":       true,
		"aws:tikv:m6g.xlarge:gp3:4000:400":     true,
		"aws:tiflash:m6g.xlarge:gp3:6000:625":  true,
		"aws:tikv:m6g.2xlarge:gp3:4000:400":    true,
		"aws:tikv:r6g.2xlarge:gp3:4000:400":    true,
		"aws:tiflash:r6g.2xlarge:gp3:6000:625": true,
		"aws:pd:t3a.medium:gp3:4000:400":       true,
		"aws:pd:c5.xlarge:gp3:4000:400":        true,
		"aws:tikv:t3a.medium:gp3:4000:400":     true,
		"aws:tikv:m5.xlarge:gp3:3000:128":      true,
		"aws:tikv:m5.2xlarge:gp3:4000:400":     true,
		"aws:tikv:r5.2xlarge:gp3:4000:400":     true,
		"aws:tiflash:t3a.medium:gp3:6000:625":  true,
		"aws:tiflash:r5.2xlarge:gp3:6000:625":  true,
	}
)

type storageThroughputCheckerParams struct {
	enableWhitelist bool
	skipPD          bool
}

func newStorageThroughputChecker() *cobra.Command {
	params := &storageThroughputCheckerParams{}
	cmd := &cobra.Command{
		Use:   "check-storage-throughput",
		Short: "check tierUnit.storage.throughput against AWS EC2 EBS throughput baseline",
		Run: func(cmd *cobra.Command, args []string) {
			checkStorageThroughput(params)
		},
	}
	cmd.Flags().BoolVar(&params.enableWhitelist, "enable-whitelist", true, "whether enable whitelist")
	cmd.Flags().BoolVar(&params.skipPD, "skip-pd", false, "whether skip check for pd component")

	return cmd
}

func checkStorageThroughput(params *storageThroughputCheckerParams) {
	fullPath, _ := filepath.Abs(yamlDir)
	logrus.Infof("run storage throughput in %s", fullPath)

	tierUnits := readTierUnits(yamlDir)

	var hasDiff bool
	for _, each := range tierUnits {
		if params.skipPD && each.ComponentType == "pd" {
			logrus.Debugf("skip check for pd component")
			continue
		}

		if each.CloudProvider != model.CloudProviderAWS {
			continue
		}

		// only gp3 EBS can customize throughput
		if each.Storage.InstanceType != "gp3" {
			continue
		}

		// ignore shared tier unit
		if each.DeployType == "shared" {
			continue
		}

		baselineThroughput, ok := instance2InstanceEBSThroughputBaseline[each.Compute.InstanceType]
		if !ok {
			logrus.Fatalf("can't find baseline for %s", each.Compute.InstanceType)
		}

		if each.Storage.ThroughputMi != baselineThroughput {
			if params.enableWhitelist && whitelistTierUnitKeys[each.Key] {
				logrus.Infof("ignore diff, because this tier unit is in whitelist. tierUnitKey=%s", each.Key)
				continue
			}

			hasDiff = true
			logrus.Warnf("tierUnitKey=%s, tier-throughput=%d, baseline-throughput=%d \n",
				each.Key, each.Storage.ThroughputMi, baselineThroughput)
		}
	}

	if hasDiff {
		logrus.Fatalf("has throughput diff between tier_unit.storage.throughput and ec2 ebs throughput baseline")
	} else {
		logrus.Infof("no diff found\n")
	}
}