package chaos

import (
	"fmt"
	"math/rand"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/pingcap/log"
	"go.uber.org/zap"
)

type RoleType int

const (
	AnyOne RoleType = iota
	GivenOne
	AllMembers
)

type ActionType int

const (
	KillServer ActionType = iota
	DelayNetwork
	RecoverNetwork
	InterruptNetwork
)

type Chaos interface {
	KillOneTiKV() error
	KillAllTiKV() error
	KillOnePD() error
	KillAllPD() error
	ChaosAction(string, RoleType, ActionType, []string, time.Duration, time.Duration) error
	String() string
}

var (
	// TODO: chaos for k8s environment
	_ Chaos = (*AnsibleChaos)(nil)

	HostRegexp = regexp.MustCompile(`^hosts \(([0-9]+)\):$`)
	// Clusters for Apiv2 test
	ApiV2Clusters = []string{"tikv", "pd"}
	// Clusters for TiKVCDC test
	TiKVCDCClusters = []string{"tikv", "pd", "tikv_cdc", "dst_tikv", "dst_pd"}
	// Set network latency
	TCSetLatencyCmd = "nohup sh -c 'tc qdisc add dev eth0 root netem delay %dms; sleep %d; tc qdisc del dev eth0 root netem' > /dev/null 2>&1 &"
	// Interrupt network Cmd
	InterruptNetworkCmd = "nohup sh -c 'ifdown eth0; sleep %d; ifup eth0' > /dev/null 2>&1 &"
)

// AnsibleChaos generates chaos by "ansible"
// Usage:
// 1. Install ansible. E.g. `yum install ansible`
// 2. Put a hosts file for tikv & pd in current directory
//  [tikv]
//  192.168.0.1
//  192.168.0.2
//  [pd]
//  192.168.0.3
//  192.168.0.4
// 3. Make ansible be able to connect to these hosts. See doc of ansible
type AnsibleChaos struct {
	hosts string
	key   string
	user  string

	// Cluster name -> cluster hosts
	// Now clusters have: tikv, pd, tikv_cdc, dst_tikv, dst_pd
	clustersHosts map[string][]string
	clustersName  []string

	dryRun bool // be true to do NO real change to system.
}

func NewAnsibleChaos(hosts string, key string, user string, withCdc bool) (chaos *AnsibleChaos, err error) {
	clustersName := ApiV2Clusters
	if withCdc {
		clustersName = TiKVCDCClusters
	}
	ansible := &AnsibleChaos{
		hosts:         hosts,
		key:           key,
		user:          user,
		clustersHosts: make(map[string][]string),
		clustersName:  clustersName,
	}

	for _, clusterName := range clustersName {
		if err = ansible.Ping(clusterName); err != nil {
			return nil, err
		}
		if ansible.clustersHosts[clusterName], err = ansible.ListHosts(clusterName); err != nil {
			return nil, err
		}
	}

	return ansible, err
}

func (c *AnsibleChaos) SetDryRun(dryRun bool) {
	c.dryRun = dryRun
}

func (c AnsibleChaos) String() string {
	clustersHosts := make([]string, 0, len(c.clustersName))
	for _, clusterName := range c.clustersName {
		clusterHosts := fmt.Sprintf("%s: %v", clusterName, c.clustersHosts[clusterName])
		clustersHosts = append(clustersHosts, clusterHosts)
	}
	return "{" + strings.Join(clustersHosts, ",") + "}"
}

func (c *AnsibleChaos) getCommendArgs(clusterName string) []string {
	args := []string{clusterName}

	if len(c.hosts) > 0 {
		args = append(args, "-i", c.hosts)
	}
	if len(c.key) > 0 {
		args = append(args, "--key-file", c.key)
	}
	if len(c.user) > 0 {
		args = append(args, "--user", c.user)
	}
	return args
}

func (c *AnsibleChaos) Ping(clusterName string) error {
	args := c.getCommendArgs(clusterName)
	args = append(args, "-m", "ping")
	cmd := exec.Command("ansible", args...)
	out, err := cmd.CombinedOutput()
	log.Info("AnsibleChaos.Ping", zap.String("cmd", cmd.String()), zap.Any("out", string(out)), zap.Error(err))
	return err
}

func (c *AnsibleChaos) ListHosts(clusterName string) (hosts []string, err error) {
	args := c.getCommendArgs(clusterName)
	args = append(args, "--list-hosts")
	cmd := exec.Command("ansible", args...)
	out, err := cmd.CombinedOutput()
	log.Info("AnsibleChaos.ListHosts", zap.String("cmd", cmd.String()), zap.String("out", string(out)), zap.Error(err))
	if err != nil {
		return nil, err
	}

	// output clusterName:
	// hosts (3):
	//   192.168.0.1
	//   192.168.0.2
	//   192.168.0.3
	lines := strings.Split(string(out), "\n")
	hostsNum := 0
	offset := 0
	for i, line := range lines {
		line := strings.TrimSpace(line)
		if matchs := HostRegexp.FindStringSubmatch(line); len(matchs) == 2 {
			num, err := strconv.Atoi(matchs[1])
			if err != nil {
				return nil, err
			}
			offset = i + 1
			hostsNum = num
			break
		}
	}

	if hostsNum == 0 {
		return nil, fmt.Errorf("no %s hosts in config", clusterName)
	}

	for i := 0; i < hostsNum; i++ {
		line := lines[offset+i]
		host := strings.TrimSpace(line)
		if len(host) > 0 {
			hosts = append(hosts, host)
		}
	}
	log.Info("Get cluster hosts", zap.String("cluster", clusterName), zap.Any("hosts", hosts))
	return hosts, nil
}

func (c *AnsibleChaos) pickHosts(clusterName string, count int) []string {
	hosts, ok := c.clustersHosts[clusterName]
	if !ok {
		return nil
	}

	rand.Shuffle(len(hosts), func(i, j int) {
		hosts[i], hosts[j] = hosts[j], hosts[i]
	})

	return hosts[:count]
}

func (c *AnsibleChaos) getServiceName(clusterName string) string {
	switch clusterName {
	case "tikv", "dst_tikv":
		return "tikv-server"
	case "pd", "dst_pd":
		return "pd-server"
	case "tikv_cdc":
		return "cdc"
	default:
		log.Panic("unexpected clusterName", zap.String("clusterName", clusterName))
		panic(clusterName)
	}
}

func (c *AnsibleChaos) excuteCmd(clusterName, remoteCmd string, hosts []string) error {
	args := c.getCommendArgs(clusterName)
	args = append(args, "--limit", strings.Join(hosts, " "))
	args = append(args, "-m", "shell", "-a", remoteCmd)

	cmd := exec.Command("ansible", args...)
	out, err := cmd.CombinedOutput()

	log.Warn("INFO: AnsibleChaos", zap.String("clusterName", clusterName), zap.Any("hosts", hosts),
		zap.String("cmd", cmd.String()), zap.String("out", string(out)), zap.Error(err))

	return err
}

func (c *AnsibleChaos) ChaosAction(clusterName string, role RoleType,
	action ActionType, hosts []string, latency time.Duration, outage time.Duration) error {

	switch role {
	case AnyOne:
		hosts = c.pickHosts(clusterName, 1)
	case GivenOne:
	case AllMembers:
		hosts = c.clustersHosts[clusterName]
	}
	if len(hosts) == 0 {
		return fmt.Errorf("no hosts found for cluster %s", clusterName)
	}

	var remoteCmd string
	switch action {
	case KillServer:
		if c.dryRun {
			remoteCmd = fmt.Sprintf("pgrep --list-name %v", c.getServiceName(clusterName))
		} else {
			// "killall --wait -9 tikv-server"
			remoteCmd = fmt.Sprintf("killall --wait -9 %v", c.getServiceName(clusterName))
		}
	case DelayNetwork:
		remoteCmd = fmt.Sprintf(TCSetLatencyCmd, latency.Milliseconds(), int(outage.Seconds()))
	case InterruptNetwork:
		remoteCmd = fmt.Sprintf(InterruptNetworkCmd, int(outage.Seconds()))
	}

	return c.excuteCmd(clusterName, remoteCmd, hosts)
}

func (c *AnsibleChaos) KillOneTiKV() error {
	return c.ChaosAction("tikv", AnyOne, KillServer, nil, 0, 0)
}
func (c *AnsibleChaos) KillOnePD() error {
	return c.ChaosAction("pd", AnyOne, KillServer, nil, 0, 0)
}

func (c *AnsibleChaos) KillAllTiKV() error {
	return c.ChaosAction("tikv", AllMembers, KillServer, nil, 0, 0)
}
func (c *AnsibleChaos) KillAllPD() error {
	return c.ChaosAction("pd", AllMembers, KillServer, nil, 0, 0)
}