package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/bitfield/script"
	"github.com/bramvdbogaerde/go-scp"
	"golang.org/x/crypto/ssh"
)

type envConfig struct {
	Host     string `json:"Host"`
	Port     int    `json:"Port"`
	Username string `json:"Username"`
	Password string `json:"Password"`
}

// 在本地启动一个 iperf3  服务端
func iperfServer() {
	cmd := exec.Command("benchmark/tools/iperf3", "-s")
	cmd.Run()
}

func randomString(length int) string {
	const charset = "abcdefghijklmnopqrstuvwxyz" +
		"ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

	var seededRand *rand.Rand = rand.New(
		rand.NewSource(time.Now().UnixNano()))

	b := make([]byte, length)
	for i := range b {
		b[i] = charset[seededRand.Intn(len(charset))]
	}
	return string(b)
}

// 封装 远程执行 shell 命令
func remoteShellExec(rClient ssh.Client, commands []string, logOut bool) error {

	outputStr := ""
	for _, command := range commands {
		session, err := rClient.NewSession()
		if err != nil {
			log.Fatal("Failed to create session: ", err)
		}
		defer session.Close()
		output, err := session.CombinedOutput(command)
		if err != nil {
			log.Fatal("Failed to run: ", err)
		}
		outputStr = string(output)
		if logOut {
			log.Println(outputStr)
		}
	}
	return nil
}

// 测试 SSH 目标主机是否可联通
func testSSH(rClient ssh.Client) bool {

	session, err := rClient.NewSession()
	if err != nil {
		log.Fatal("Failed to dial: ", err)
		return false
	}
	defer session.Close()
	return true
}

func main() {
	file, err := os.Open("config.json")
	if err != nil {
		fmt.Println("Error opening file:", err)
		return
	}
	defer file.Close()
	envConfig := &envConfig{}
	decoder := json.NewDecoder(file)
	err = decoder.Decode(&envConfig)
	if err != nil {
		fmt.Println("Error decoding file:", err)
		return
	}

	srcDir := "benchmark"
	dstDir := "/tmp/" + randomString(10) + "/"

	// 读取本地 config.json 文件, 注入成变量

	username := envConfig.Username
	password := envConfig.Password
	port := envConfig.Port
	HostExecOut, _ := script.Exec("hostname -I").String()
	hostIPS := strings.Split(HostExecOut, " ")

	config := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.AuthMethod{
			ssh.Password(password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	// 创建 iperfServer 协程
	go iperfServer()

	// 自动关闭 iperfServer 协程
	defer func() {
		cmd := exec.Command("killall", "iperf3")
		_, err := cmd.Output()
		if err != nil {
			log.Fatal(err)
		}
	}()

	// 使用 testSSH 函数测试目标主机是否可联通
	for _, host := range strings.Split(envConfig.Host, ",") {
		serverURL := fmt.Sprintf("%s:%d", host, port)
		client, err := ssh.Dial("tcp", serverURL, config)
		if err != nil {
			log.Fatal("Failed to dial: ", err)
		}
		if testSSH(*client) {
			log.Printf("Host: %s , SSH 连接成功", host)
		} else {
			log.Printf("Host: %s , SSH 连接失败", host)
		}
	}

	// 按 , 切 config.Host 中的每一个 host
	for _, host := range strings.Split(envConfig.Host, ",") {
		serverURL := fmt.Sprintf("%s:%d", host, port)

		client, err := ssh.Dial("tcp", serverURL, config)
		if err != nil {
			log.Fatal("Failed to dial: ", err)
		}

		_ = remoteShellExec(*client, []string{"mkdir -p " + dstDir}, false)

		filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if info.IsDir() {
				remoteShellExec(*client, []string{"mkdir -p " + dstDir + path}, false)
			} else {
				scpClient, err := scp.NewClientBySSH(client)
				scpClient.Connect()
				file, err := os.Open(path)
				defer file.Close()
				if err != nil {
					log.Fatal("Failed to open file: ", err)
				}
				err = scpClient.CopyFile(context.Background(), file, dstDir+path, "0655")
				if err != nil {
					log.Fatal("Failed to copy file: ", err)
				}
				scpClient.Close()

			}
			return nil
		})

		// 执行 远端 脚本
		log.Printf("Host: %s , 正在执行远端脚本...", host)

		_ = remoteShellExec(*client, []string{
			"cd " + dstDir + srcDir + " && " + "./run.sh " + hostIPS[0],
		}, true)

		// 删除远端文件
		remoteShellExec(*client, []string{"rm -rf " + dstDir}, false)
	}
}