package main

import (
	"Walnut-HS/config"
	gConfig "Walnut-HS/genesis/config"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"slices"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
	"gopkg.in/yaml.v2"
)

var sshConnectionPool = make(map[string]*ssh.Client)
var sftpConnectionPool = make(map[string]*sftp.Client)
var ipList1 = []string{
	"10.18.1.4",
	"10.18.1.5",
	"10.18.1.6",
	"10.18.1.7",
	"10.18.1.8",

	"10.18.1.36",
	"10.18.1.37",
	"10.18.1.38",
	"10.18.1.39",
	"10.18.1.40",
}

func TestUploadFiles(t *testing.T) {
	genesisCfg, err := gConfig.NewConfiguration("./genesis_1000.yaml")
	if err != nil {
		t.Log(err.Error())
		return
	}
	sshCertsNameList := []string{}
	for i := 0; i < 10; i++ {
		sshCertsNameList = append(sshCertsNameList, fmt.Sprintf("vsm-scp%d.pem", i+1))
	}
	num := 1000
	committeeNum := 200
	var userName, password, prefixPath string
	tarFileName := "sc-s.tar.gz"
	localFilePath := "./" + tarFileName
	nodeProgramName := "scbft"
	genesisProgramName := "genesis"
	configFileName := "sc_bft.yaml"

	ipNodesMapping := map[string]int{}
	wg := sync.WaitGroup{}
	maxConcurrent := 10
	sem := make(chan struct{}, maxConcurrent)

	for i, ip := range genesisCfg.IpList[:num] {
		if slices.Contains(ipList1, ip) {
			userName = "gravity"
			password = "123456"
			prefixPath = "/home/gravity/app/scbft-s/"
		} else {
			userName = "guanyuan"
			password = "123456"
			prefixPath = "/home/guanyuan/app/scbft-s/"
		}
		index, ok := ipNodesMapping[ip]
		if !ok {
			ipNodesMapping[ip] = 1
		} else {
			ipNodesMapping[ip] = index + 1
		}
		index += 1

		filePath := fmt.Sprintf("%snode%d/", prefixPath, index)
		keyPath := fmt.Sprintf("%skeys/", filePath)
		publicKeyList := make([]string, len(genesisCfg.PublicKeys))
		for k, key := range genesisCfg.PublicKeys {
			publicKeyList[k] = keyPath + path.Base(key)
		}
		privateKeyList := make([]string, len(genesisCfg.PrivateKeys))
		for k, key := range genesisCfg.PrivateKeys {
			privateKeyList[k] = keyPath + path.Base(key)
		}
		vrfPrivateKeyList := make([]string, len(genesisCfg.VrfPrivateKeys))
		for k, key := range genesisCfg.VrfPrivateKeys {
			vrfPrivateKeyList[k] = keyPath + path.Base(key)
		}
		vrfPublicKeyList := make([]string, len(genesisCfg.VrfPublicKeys))
		for k, key := range genesisCfg.VrfPublicKeys {
			vrfPublicKeyList[k] = keyPath + path.Base(key)
		}

		nodeCfg := &config.Configuration{
			IpList:          genesisCfg.IpList,
			PortList:        genesisCfg.PortList,
			PrivateKeys:     privateKeyList,
			PublicKeys:      publicKeyList,
			VrfPrivateKeys:  vrfPrivateKeyList,
			VrfPublicKeys:   vrfPublicKeyList,
			GenesisFilePath: filePath + "genesis.out",
			NodeIndex:       i,
		}

		cfgBytes, err := yaml.Marshal(nodeCfg)
		if err != nil {
			t.Log(err.Error())
			return
		}

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

		conn, ok := sshConnectionPool[ip]
		if !ok {
			conn, err = ssh.Dial("tcp", ip+":22", sshConfig)
			if err != nil {
				t.Logf("failed to dial: %s", err.Error())
				return
			}
			sshConnectionPool[ip] = conn
		}
		sftpConn, ok := sftpConnectionPool[ip]
		if !ok {
			sftpConn, err = sftp.NewClient(conn)
			if err != nil {
				t.Logf("failed to dial: %s", err.Error())
				return
			}
			sftpConnectionPool[ip] = sftpConn
		}
		if index == 1 {
			// each ip only uploads once
			s, err := conn.NewSession()
			if err != nil {
				t.Log(err.Error())
			}
			if err = s.Run("mkdir -p " + prefixPath); err != nil {
				t.Log(err.Error())
			}
			s.Close()
			s, err = conn.NewSession()
			if err != nil {
				t.Log(err.Error())
			}
			if err = s.Run(fmt.Sprintf("echo '%s' > %s", start_sh, prefixPath+"start.sh")); err != nil {
				t.Log(err.Error())
			}
			s.Close()
			s, err = conn.NewSession()
			if err != nil {
				t.Log(err.Error())
			}
			if err = s.Run("chmod +x " + prefixPath + "start.sh"); err != nil {
				t.Log(err.Error())
			}
			s.Close()
			s, err = conn.NewSession()
			if err != nil {
				t.Log(err.Error())
			}
			if err = s.Run(fmt.Sprintf("echo '%s' > %s", genesis_sh, prefixPath+"genesis.sh "+"&& chmod +x "+prefixPath+"genesis.sh")); err != nil {
				t.Log(err.Error())
			}
			s.Close()
		}

		wg.Add(1)
		go func(ip string, index int, conn *ssh.Client, sftpConn *sftp.Client) {
			defer wg.Done()

			sem <- struct{}{}

			if err := createAndUploadFiles(conn, sftpConn, filePath, localFilePath, tarFileName, nodeProgramName, genesisProgramName, cfgBytes, configFileName, num, committeeNum); err != nil {
				t.Logf("upload files to %s node%d failed, %s", ip, index, err.Error())
			} else {
				t.Logf("finish uploading files to %s node%d", ip, index)
			}

			<-sem

		}(ip, index, conn, sftpConn)
	}
	wg.Wait()
	for _, conn := range sshConnectionPool {
		conn.Close()
	}
	for _, conn := range sftpConnectionPool {
		conn.Close()
	}

}

// Mapping of internal IP and public IP
var ipMap = map[string]string{
	"172.31.3.113":  "18.189.141.252",
	"172.31.32.241": "3.144.129.133",
	"172.31.1.239":  "18.118.210.106",
	"172.31.14.126": "18.118.37.244",
	"172.31.1.189":  "3.17.79.71",
	"172.31.0.10":   "3.16.76.212",
	"172.31.5.105":  "3.16.203.242",
	"172.31.12.231": "18.116.14.33",
	"172.31.13.18":  "3.137.219.40",
	"172.31.15.209": "3.15.139.21",
}

func TestUploadFiles2(t *testing.T) {
	genesisCfg, err := gConfig.NewConfiguration("./genesis_100.yaml")
	if err != nil {
		t.Log(err.Error())
		return
	}

	num := 100
	committeeNum := 100
	var userName, prefixPath string
	tarFileName := "sc-s.tar.gz"
	localFilePath := "./" + tarFileName
	nodeProgramName := "scbft"
	genesisProgramName := "genesis"
	configFileName := "sc_bft.yaml"
	prefixPath = "/home/ubuntu/app/scbft-s/"
	userName = "ubuntu"
	sshKey, err := os.ReadFile("./certs/vsm-scp.pem")
	if err != nil {
		t.Log(err.Error())
		return
	}
	signer, err := ssh.ParsePrivateKey(sshKey)
	if err != nil {
		t.Logf("无法解析私钥: %v", err)
		return
	}

	ipNodesMapping := map[string]int{}
	sshConnectionPool := make(map[string]*ssh.Client)
	sftpConnectionPool := make(map[string]*sftp.Client)
	var mapMutex sync.Mutex

	wg := sync.WaitGroup{}
	sem := make(chan struct{}, 10)

	for start := 0; start < num; start += 10 {
		wg.Add(1)
		go func(start int) {
			defer wg.Done()
			sem <- struct{}{}
			defer func() { <-sem }()

			end := start + 10
			if end > num {
				end = num
			}

			for i := start; i < end; i++ {
				ip := genesisCfg.IpList[i]
				publicIp := ipMap[ip]

				mapMutex.Lock()
				index, ok := ipNodesMapping[ip]
				if !ok {
					ipNodesMapping[ip] = 1
				} else {
					ipNodesMapping[ip] = index + 1
				}
				index += 1
				mapMutex.Unlock()

				filePath := fmt.Sprintf("%snode%d/", prefixPath, index)
				keyPath := fmt.Sprintf("%skeys/", filePath)
				publicKeyList := make([]string, len(genesisCfg.PublicKeys))
				for k, key := range genesisCfg.PublicKeys {
					publicKeyList[k] = keyPath + path.Base(key)
				}
				privateKeyList := make([]string, len(genesisCfg.PrivateKeys))
				for k, key := range genesisCfg.PrivateKeys {
					privateKeyList[k] = keyPath + path.Base(key)
				}
				vrfPrivateKeyList := make([]string, len(genesisCfg.VrfPrivateKeys))
				for k, key := range genesisCfg.VrfPrivateKeys {
					vrfPrivateKeyList[k] = keyPath + path.Base(key)
				}
				vrfPublicKeyList := make([]string, len(genesisCfg.VrfPublicKeys))
				for k, key := range genesisCfg.VrfPublicKeys {
					vrfPublicKeyList[k] = keyPath + path.Base(key)
				}

				nodeCfg := &config.Configuration{
					IpList:          genesisCfg.IpList,
					PortList:        genesisCfg.PortList,
					PrivateKeys:     privateKeyList,
					PublicKeys:      publicKeyList,
					VrfPrivateKeys:  vrfPrivateKeyList,
					VrfPublicKeys:   vrfPublicKeyList,
					GenesisFilePath: filePath + "genesis.out",
					NodeIndex:       i,
				}

				cfgBytes, err := yaml.Marshal(nodeCfg)
				if err != nil {
					t.Log(err.Error())
					continue
				}

				sshConfig := &ssh.ClientConfig{
					User: userName,
					Auth: []ssh.AuthMethod{
						ssh.PublicKeys(signer),
					},
					HostKeyCallback: ssh.InsecureIgnoreHostKey(),
				}

				mapMutex.Lock()
				conn, ok := sshConnectionPool[ip]
				if !ok {
					conn, err = ssh.Dial("tcp", publicIp+":22", sshConfig)
					if err != nil {
						t.Logf("failed to dial: %s", err.Error())
						mapMutex.Unlock()
						continue
					}
					sshConnectionPool[ip] = conn
				}
				sftpConn, ok := sftpConnectionPool[ip]
				if !ok {
					sftpConn, err = sftp.NewClient(conn)
					if err != nil {
						t.Logf("failed to dial: %s", err.Error())
						mapMutex.Unlock()
						continue
					}
					sftpConnectionPool[ip] = sftpConn
				}
				mapMutex.Unlock()

				if index == 1 {
					s, err := conn.NewSession()
					if err != nil {
						t.Log(err.Error())
					}
					if err = s.Run("mkdir -p " + prefixPath); err != nil {
						t.Log(err.Error())
					}
					s.Close()
					s, err = conn.NewSession()
					if err != nil {
						t.Log(err.Error())
					}
					if err = s.Run(fmt.Sprintf("echo '%s' > %s", start_sh, prefixPath+"start.sh")); err != nil {
						t.Log(err.Error())
					}
					s.Close()
					s, err = conn.NewSession()
					if err != nil {
						t.Log(err.Error())
					}
					if err = s.Run("chmod +x " + prefixPath + "start.sh"); err != nil {
						t.Log(err.Error())
					}
					s.Close()
					s, err = conn.NewSession()
					if err != nil {
						t.Log(err.Error())
					}
					if err = s.Run(fmt.Sprintf("echo '%s' > %s", genesis_sh, prefixPath+"genesis.sh && chmod +x "+prefixPath+"genesis.sh")); err != nil {
						t.Log(err.Error())
					}
					s.Close()
				}

				st := time.Now()
				if err := createAndUploadFiles(conn, sftpConn, filePath, localFilePath, tarFileName, nodeProgramName, genesisProgramName, cfgBytes, configFileName, num, committeeNum); err != nil {
					t.Logf("upload files to %s node%d failed, %s", publicIp, index, err.Error())
				} else {
					et := time.Now()
					t.Logf("finish uploading files to %s node%d, used time: %.1f s", publicIp, index, et.Sub(st).Seconds())
				}
			}
		}(start)
	}
	wg.Wait()

	for _, conn := range sshConnectionPool {
		conn.Close()
	}
	for _, conn := range sftpConnectionPool {
		conn.Close()
	}
}

func TestChangeIndex(t *testing.T) {
	genesisCfg, err := gConfig.NewConfiguration("./sc_bft.yaml")
	if err != nil {
		t.Log(err.Error())
		return
	}

	num := 50
	var userName, prefixPath string

	configFileName := "sc_bft.yaml"
	prefixPath = "/home/ubuntu/app/scbft-s/"
	userName = "ubuntu"
	sshKey, err := os.ReadFile("./certs/vsm-scp.pem")
	if err != nil {
		t.Log(err.Error())
		return
	}
	signer, err := ssh.ParsePrivateKey(sshKey)
	if err != nil {
		t.Logf("无法解析私钥: %v", err)
		return
	}

	ipNodesMapping := map[string]int{}
	sshConnectionPool := make(map[string]*ssh.Client)

	for start := 0; start < num; start += 10 {

		end := start + 10
		if end > num {
			end = num
		}

		for i := start; i < end; i++ {
			ip := genesisCfg.IpList[i]
			publicIp := ipMap[ip]

			index, ok := ipNodesMapping[ip]
			if !ok {
				ipNodesMapping[ip] = 1
			} else {
				ipNodesMapping[ip] = index + 1
			}
			index += 1

			filePath := fmt.Sprintf("%snode%d/", prefixPath, index)
			keyPath := fmt.Sprintf("%skeys/", filePath)
			publicKeyList := make([]string, len(genesisCfg.PublicKeys))
			for k, key := range genesisCfg.PublicKeys {
				publicKeyList[k] = keyPath + path.Base(key)
			}
			privateKeyList := make([]string, len(genesisCfg.PrivateKeys))
			for k, key := range genesisCfg.PrivateKeys {
				privateKeyList[k] = keyPath + path.Base(key)
			}
			vrfPrivateKeyList := make([]string, len(genesisCfg.VrfPrivateKeys))
			for k, key := range genesisCfg.VrfPrivateKeys {
				vrfPrivateKeyList[k] = keyPath + path.Base(key)
			}
			vrfPublicKeyList := make([]string, len(genesisCfg.VrfPublicKeys))
			for k, key := range genesisCfg.VrfPublicKeys {
				vrfPublicKeyList[k] = keyPath + path.Base(key)
			}

			nodeCfg := &config.Configuration{
				IpList:          genesisCfg.IpList,
				PortList:        genesisCfg.PortList,
				PrivateKeys:     privateKeyList,
				PublicKeys:      publicKeyList,
				VrfPrivateKeys:  vrfPrivateKeyList,
				VrfPublicKeys:   vrfPublicKeyList,
				GenesisFilePath: filePath + "genesis.out",
				NodeIndex:       i,
			}

			cfgBytes, err := yaml.Marshal(nodeCfg)
			if err != nil {
				t.Log(err.Error())
				continue
			}

			sshConfig := &ssh.ClientConfig{
				User: userName,
				Auth: []ssh.AuthMethod{
					ssh.PublicKeys(signer),
				},
				HostKeyCallback: ssh.InsecureIgnoreHostKey(),
			}

			conn, ok := sshConnectionPool[ip]
			if !ok {
				conn, err = ssh.Dial("tcp", publicIp+":22", sshConfig)
				if err != nil {
					t.Logf("failed to dial: %s", err.Error())
					continue
				}
				sshConnectionPool[ip] = conn
			}
			sftpConn, ok := sftpConnectionPool[ip]
			if !ok {
				sftpConn, err = sftp.NewClient(conn)
				if err != nil {
					t.Logf("failed to dial: %s", err.Error())
					continue
				}
				sftpConnectionPool[ip] = sftpConn
			}
			tempFile, err := os.CreateTemp("", "temp_config_*")
			if err != nil {
				t.Log(err.Error())
				return
			}
			defer tempFile.Close()
			tempFile.Write(cfgBytes)
			if err = uploadFile(sftpConn, tempFile.Name(), filePath+configFileName); err != nil {
				t.Log(err)
				continue
			}
			t.Logf("finish uploading config file to %s node%d", publicIp, index)
		}
	}

	for _, conn := range sshConnectionPool {
		conn.Close()
	}
	for _, conn := range sftpConnectionPool {
		conn.Close()
	}
}

func createAndUploadFiles(conn *ssh.Client, sftpConn *sftp.Client, filePath, localFilePath, tarFileName, nodeProgramName, genesisProgramName string, cfgBytes []byte, configFileName string, num, committeeNum int) error {
	s, err := conn.NewSession()
	if err != nil {
		return fmt.Errorf("failed to open session: %v", err)
	}
	defer s.Close()

	if err = s.Run("mkdir -p " + filePath); err != nil {
		return fmt.Errorf("create folder failed: %v", err)
	}

	if err = uploadFile(sftpConn, localFilePath, filePath+tarFileName); err != nil {
		return fmt.Errorf("upload file failed: %v", err)
	}

	if err = runCommand(conn, "tar -xzvf "+filePath+tarFileName+" -C "+filePath); err != nil {
		return fmt.Errorf("unzip file failed: %v", err)
	}

	if err = runCommand(conn, "chmod +x "+filePath+nodeProgramName); err != nil {
		return fmt.Errorf("change permission for node program failed: %v", err)
	}

	if err = runCommand(conn, "chmod +x "+filePath+genesisProgramName); err != nil {
		return fmt.Errorf("change permission for genesis program failed: %v", err)
	}
	tempFile, err := os.CreateTemp("", "temp_config_*")
	if err != nil {
		return fmt.Errorf("create temp file failed: %v", err)
	}
	defer tempFile.Close()
	tempFile.Write(cfgBytes)
	if err = uploadFile(sftpConn, tempFile.Name(), filePath+configFileName); err != nil {
		return fmt.Errorf("upload config failed: %v", err)
	}

	if err = runCommand(conn, fmt.Sprintf("%s%s -n_cnt=%d -cn_cnt=%d -config_path=%s -out_path=%s", filePath, genesisProgramName, num, committeeNum, filePath+configFileName, filePath+"genesis.out")); err != nil {
		return fmt.Errorf("generate genesis failed: %v", err)
	}

	return nil
}

func runCommand(conn *ssh.Client, command string) error {
	const maxRetries = 25
	cmdErr := bytes.Buffer{}
	for i := 0; i < maxRetries; i++ {
		s, err := conn.NewSession()
		if err != nil {
			time.Sleep(time.Second)
			continue
		}
		s.Stderr = &cmdErr
		defer s.Close()

		if err := s.Run(command); err != nil {
			if strings.Contains(err.Error(), "administratively prohibited") {
				time.Sleep(time.Duration(i+1) * time.Second) // retry after waiting
				continue
			}
			return errors.New(cmdErr.String())
		}
		return nil
	}
	return fmt.Errorf("command failed after %d retries", maxRetries)
}

func uploadFile(client *sftp.Client, localFilePath string, remoteFilePath string) error {
	file, err := os.Open(localFilePath)
	if err != nil {
		return fmt.Errorf("unable to open local file: %v", err)
	}
	defer file.Close()

	remoteFile, err := client.Create(remoteFilePath)
	if err != nil {
		return fmt.Errorf("unable to create remote file: %v", err)
	}
	defer remoteFile.Close()
	_, err = io.Copy(remoteFile, file)
	if err != nil {
		return fmt.Errorf("unable to copy file content: %v", err)
	}

	return nil
}

const start_sh = `#!/bin/bash

# Check if the required parameters are provided
if [ "$#" -lt 6 ]; then
    echo "Usage: $0 <node_index> <node_count> <committee_node_count> <start_time> <batch_size> <views>"
    exit 1
fi

# Get parameters
NODE_INDEX=$1
NODE_COUNT=$2
COMMITTEE_NODE_COUNT=$3
START_TIME=$4
BATCH_SIZE=$5
VIEWS=$6

BASE_DIR=~/app/scbft-s

# Start nodes
for ((i=1; i<=NODE_INDEX; i++)); do
    NODE_DIR="$BASE_DIR/node$i"
    SCBFT_EXEC="$NODE_DIR/scbft"
    CONFIG_PATH="$NODE_DIR/sc_bft.yaml" # Assuming each node has its own config file
    LOG_FILE="$NODE_DIR/scbft.log" # Log file for the node

    # Check if the node directory and executable file exist
    if [ -d "$NODE_DIR" ] && [ -f "$SCBFT_EXEC" ] && [ -f "$CONFIG_PATH" ]; then
        echo "Starting $SCBFT_EXEC with config $CONFIG_PATH..."

        # Start the node with nohup and redirect output to the log file
		cd "$NODE_DIR"
        nohup "$SCBFT_EXEC" -config_path="$CONFIG_PATH" -n_cnt="$NODE_COUNT" -cn_cnt="$COMMITTEE_NODE_COUNT" -start_time="$START_TIME" -batch_size="$BATCH_SIZE" -views="$VIEWS" -timeout 100 > "$LOG_FILE" 2>&1 &

        echo "Node $i started, logs are being written to $LOG_FILE"
    else
        echo "Warning: Directory $NODE_DIR, executable file $SCBFT_EXEC, or config file $CONFIG_PATH does not exist"
    fi
done
`

const genesis_sh = `#!/bin/bash
# Check if the required parameters are provided
if [ "$#" -lt 3 ]; then
    echo "Usage: $0 <node_index> <node_count> <committee_node_count>"
    exit 1
fi
NODE_INDEX=$1
NODE_COUNT=$2
COMMITTEE_NODE_COUNT=$3
BASE_DIR=~/app/scbft-s

# Start generating genesis
for ((i=1; i<=NODE_INDEX; i++)); do
    NODE_DIR="$BASE_DIR/node$i"
    GENESIS_EXEC="$NODE_DIR/genesis"
    CONFIG_PATH="$NODE_DIR/sc_bft.yaml" # Assuming each node has its own config file

    # Check if the node directory and executable file exist
    if [ -d "$NODE_DIR" ] && [ -f "$GENESIS_EXEC" ] && [ -f "$CONFIG_PATH" ]; then
        echo "Starting $GENESIS_EXEC with config $CONFIG_PATH..."

        # Start the node and pass the parameters
        "$GENESIS_EXEC" -config_path="$CONFIG_PATH" -n_cnt="$NODE_COUNT" -cn_cnt="$COMMITTEE_NODE_COUNT" -out_path="$NODE_DIR/genesis.out" &
    else
        echo "Warning: Directory $NODE_DIR, executable file $GENESIS_EXEC, or config file $CONFIG_PATH does not exist"
    fi
done`

func TestSshWithCert(t *testing.T) {
	key, err := ioutil.ReadFile("./certs/vsm-scp.pem")
	if err != nil {
		log.Fatalf("无法读取私钥文件: %v", err)
	}

	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		log.Fatalf("无法解析私钥: %v", err)
	}

	config := &ssh.ClientConfig{
		User: "ubuntu",
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	client, err := ssh.Dial("tcp", "3.133.88.3:22", config)
	if err != nil {
		log.Fatalf("无法连接到SSH服务器: %v", err)
	}
	defer client.Close()

	sftpClient, err := sftp.NewClient(client)
	if err != nil {
		log.Fatalf("无法建立SFTP连接: %v", err)
	}
	defer sftpClient.Close()

	files, err := sftpClient.ReadDir("/")
	if err != nil {
		log.Fatalf("无法读取远程目录: %v", err)
	}

	for _, file := range files {
		fmt.Println(file.Name())
	}
}

func TestA(t *testing.T) {
	N := 200
	lambda := 80
	service := N / 10
	CurCom := make([]int, lambda) // Ids of current committee members, Initialized to replica 0 to lambda-1
	for i := range lambda {
		seq := i / service
		CurCom[i] = (i%service)*10 + seq
	}
	t.Log(CurCom)
}