package chaos

import (
	"github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/endless/pkg/util/schedule"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
)

var _ schedule.DoUndo = &ConverTiFlashMode{}

type ConverTiFlashMode struct {
	tidbCli *tidb.Client
	dbs     []string
}

func NewConverTiFlashMode(tc resource.TiDBCluster, dbs []string) (*ConverTiFlashMode, error) {
	conn, err := tc.GetDB("")
	if err != nil {
		return nil, err
	}
	return &ConverTiFlashMode{
		tidbCli: tidb.NewClient(conn),
		dbs:     dbs,
	}, nil
}

func (c *ConverTiFlashMode) Do() error {
	for _, dbName := range c.dbs {
		tbls, err := c.tidbCli.ShowTableNames(dbName)
		if err != nil {
			return err
		}
		log.Info("convert tiflash mode", zap.String("db name", dbName), zap.Strings("table names", tbls), zap.String("mode", "fast"))
		err = c.tidbCli.AlterBaseTablesSetTiFlashMode(dbName, tbls, "fast")
		if err != nil {
			return err
		}
	}
	return nil
}

func (c *ConverTiFlashMode) Undo() error {
	for _, dbName := range c.dbs {
		tbls, err := c.tidbCli.ShowTableNames(dbName)
		if err != nil {
			return err
		}
		log.Info("convert tiflash mode", zap.String("db name", dbName), zap.Strings("table names", tbls), zap.String("mode", "normal"))
		err = c.tidbCli.AlterBaseTablesSetTiFlashMode(dbName, tbls, "normal")
		if err != nil {
			return err
		}
	}
	return nil
}