package command

import (
	"math"
	"strings"

	"github.com/juju/errors"
	"github.com/pingcap/schrodinger/pkg/types"
	"github.com/spf13/cobra"
)

var (
	boxFormat         string
	boxDataDir        string
	boxFilePrefix     string
	boxOffline        bool
	storageClassRules string
)

// NewBoxCommand returns the cobra command for "box"
func NewBoxCommand() *cobra.Command {
	boxCmd := &cobra.Command{
		Use:   "box <subcommand>",
		Short: "Box related commands",
	}

	boxCmd.AddCommand(NewBoxDumpCommand())
	boxCmd.AddCommand(NewBoxLoadCommand())

	return boxCmd
}

// NewBoxDumpCommand returns the cobra command for "box dump",
// "box dump" dumps boxes from schrodinger api
// files formats supported: json
func NewBoxDumpCommand() *cobra.Command {
	bc := &cobra.Command{
		Use:   "dump [options]",
		Short: "dump boxes from schrodinger, file formats supported: json",

		Run: boxDumpCommandFunc,
	}

	bc.Flags().StringVar(&boxFormat, "format", "json", "file format for data")
	bc.Flags().StringVar(&boxDataDir, "dir", "./", "the storage directory of data")
	bc.Flags().StringVar(&boxFilePrefix, "file-prefix", "boxes", "the prefix of dump files")

	return bc
}

// NewBoxLoadCommand returns the cobra command for "box load",
// "box load" loads boxes from specific files,
// files formats supported: json
func NewBoxLoadCommand() *cobra.Command {
	bc := &cobra.Command{
		Use:   "load [options]",
		Short: "load boxes from files, file formats supported: json",

		Run: boxLoadCommandFunc,
	}

	bc.Flags().StringVar(&boxFormat, "format", "json", "file format for data")
	bc.Flags().StringVar(&boxDataDir, "dir", "./", "the storage directory of data")
	bc.Flags().StringVar(&boxFilePrefix, "file-prefix", "boxes", "the prefix of dump files")
	bc.Flags().BoolVar(&boxOffline, "offline", true, "set all box offline state")
	bc.Flags().StringVar(&storageClassRules, "storage-class-rules", "", "set storage class for box, example: \"ssd=nvme-disks,share-ssd=shared-nvme-disks\"")

	return bc
}

func boxDumpCommandFunc(cmd *cobra.Command, args []string) {
	checkDumpAndLoadArgs(args, boxFormat, boxDataDir, boxFilePrefix)

	schCli := mustClientFromCmd(cmd)
	boxes, err := schCli.ListBoxes(math.MaxInt32, 0)
	if err != nil {
		ExitWithError(ExitError, err)
	}

	saveDumpFile(boxes, boxDataDir, boxFilePrefix, boxFormat)
}

func boxLoadCommandFunc(cmd *cobra.Command, args []string) {
	checkDumpAndLoadArgs(args, boxFormat, boxDataDir, boxFilePrefix)

	boxes, err := readBoxesFromJSONFile(boxDataDir, boxFilePrefix)
	if err != nil {
		ExitWithError(ExitError, err)
	}

	schCli := mustClientFromCmd(cmd)
	for _, box := range boxes {
		if err := box.Verify(); err != nil {
			ExitWithError(ExitError, err)
		}

		if boxOffline {
			box.State = types.BoxStateOffline
		}

		if len(strings.TrimSpace(storageClassRules)) > 0 {
			setBoxStorageClass(box, storageClassRules)
		}

		if _, err = schCli.SetBox(box); err != nil {
			ExitWithError(ExitError, err)
		}
	}
}

func setBoxStorageClass(box *types.Box, ruleStr string) {
	rules := parseStorageClassRules(ruleStr)
	if len(rules) == 0 {
		return
	}

	if box.Cat == nil {
		return
	}

	box.Cat.PDConfig.ResourceConfig = setBoxStorageClassAction(box.Cat.PDConfig.ResourceConfig, rules)
	box.Cat.TiDBConfig.ResourceConfig = setBoxStorageClassAction(box.Cat.TiDBConfig.ResourceConfig, rules)
	box.Cat.TiKVConfig.ResourceConfig = setBoxStorageClassAction(box.Cat.TiKVConfig.ResourceConfig, rules)
}

func setBoxStorageClassAction(resource types.ResourceRequirements, rules map[string]string) types.ResourceRequirements {
	if resource.Limits != nil {
		v, ok := rules[resource.Limits.StorageClass]
		if ok {
			resource.Limits.StorageClass = v
		}
	}
	if resource.Requests != nil {
		v, ok := rules[resource.Requests.StorageClass]
		if ok {
			resource.Requests.StorageClass = v
		}
	}

	return resource
}

func parseStorageClassRules(rules string) map[string]string {
	rulesMap := make(map[string]string)

	strs := strings.Split(rules, ",")
	if len(strs) == 0 {
		return rulesMap
	}

	for _, ruleStr := range strs {
		ruleStr = strings.TrimSpace(ruleStr)
		if len(ruleStr) == 0 {
			continue
		}

		values := strings.Split(ruleStr, "=")
		if len(values) != 2 {
			ExitWithError(ExitBadArgs, errors.Errorf("storage rules %s", rules))
		}

		rulesMap[values[0]] = values[1]
	}

	return rulesMap
}