package arsenal

import (
	"context"
	"strings"
	"time"

	"github.com/juju/errors"
	"github.com/pingcap/schrodinger/agent/daemon"
	"github.com/pingcap/schrodinger/pkg/test_client"
)

// InjectFailpointWithRecover inject a random failpoint and then recover
func InjectFailpointWithRecover(ctx context.Context, nemesisCli *client.Client, name string, downTime time.Duration) error {
	client.Infof("[nemesis] inject failpoint %s", name)

	_, err := nemesisCli.FailWithRecover(ctx, name, downTime)
	return errors.Trace(err)
}

// KillWithRecover kill nodes for downTime and then recover
func KillWithRecover(ctx context.Context, nemesisCli *client.Client, name string, downTime time.Duration) error {
	client.Infof("[nemesis] kill %s", name)
	if _, err := nemesisCli.KillWithRecover(ctx, name, downTime); err != nil {
		if strings.Contains(err.Error(), "from state [ERROR EXITED]") {
			client.Infof("dont kill %s from state [ERROR EXITED]", name)
		} else {
			return errors.Trace(err)
		}
	}
	return nil
}

// SetCPUWithRecover set cpu limit for injectTime and then recover
func SetCPUWithRecover(ctx context.Context, nemesisCli *client.Client, name string, action string, rate string, injectTime time.Duration) error {
	client.Infof("[nemesis] set %s cpu", name)
	if _, err := nemesisCli.SetCPUWithRecover(ctx, name, action, rate, injectTime); err != nil {
		if strings.Contains(err.Error(), "resources do not exist") {
			client.Infof("set %s cpu resources do not exist,error:%s", name, err.Error())
		} else {
			return errors.Trace(err)
		}
	}
	return nil
}

// SetNetworkWithRecover set network limit injectTime and then recover
func SetNetworkWithRecover(ctx context.Context, nemesisCli *client.Client, name string, action string, rate string, injectTime time.Duration) error {
	client.Infof("[nemesis] set %s network", name)
	if _, err := nemesisCli.SetNetworkWithRecover(ctx, name, action, rate, injectTime); err != nil {
		if strings.Contains(err.Error(), "resources do not exist") {
			client.Infof("set %s network resources do not exist,error:%s", name, err.Error())
		} else {
			return errors.Trace(err)
		}
	}
	return nil
}

// SetDiskWithRecover set disk for injectTime and then recover
func SetDiskWithRecover(ctx context.Context, nemesisCli *client.Client, name string, action string, rate string, injectTime time.Duration) error {
	client.Infof("[nemesis] set %s disk", name)
	if _, err := nemesisCli.SetDiskWithRecover(ctx, name, action, rate, injectTime); err != nil {
		if strings.Contains(err.Error(), "resources do not exist") {
			client.Infof("set %s disk resources do not exist,error:%s", name, err.Error())
		} else {
			return errors.Trace(err)
		}
	}
	return nil
}

// RecoverKilledNodes recovers specified nodes
func RecoverKilledNodes(ctx context.Context, nemesisCli *client.Client, name string) {
	nodes, err := nemesisCli.GetNodesWithStatus(name)
	if err != nil {
		client.Warnf("Get %s nodes failed err:%s", name, err.Error())
		return
	}

	for _, node := range nodes {
		if node.State != daemon.ProcStatRunning {
			client.Debugf("start %s node %s ", name, node.Host())
			err := nemesisCli.RecoverNode(ctx, node)
			if err != nil {
				client.Warnf("Starts %s nodes %s failed err:%s", name, node.Host(), err.Error())
			}
		}
	}
}

// RecoverNodes recovers all pd/tidb/tikv nodes
func RecoverNodes(ctx context.Context, nemesisCli *client.Client) {
	RecoverKilledNodes(ctx, nemesisCli, "tikv")
	RecoverKilledNodes(ctx, nemesisCli, "pd")
	RecoverKilledNodes(ctx, nemesisCli, "tidb")
}