package main

import (
	"crypto/sha256"
	"encoding/binary"
	"encoding/gob"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"ripple/state"
	"ripple/tests/internal/testclient"
	"ripple/types"
)

type User struct {
	Username      string `json:"username"`
	ServerAddress string `json:"serverAddress"`
	Port          int    `json:"port"`
	SecretKey     string `json:"secretKey"`
	Counter       uint32
	Local         bool `json:"local"`
}

type Step struct {
	Title string   `json:"title"`
	User  string   `json:"user"`
	Cmd   string   `json:"cmd"`
	Args  []string `json:"args"`
}

type Scenario struct {
	Users map[string]User `json:"users"`
	Steps []Step          `json:"steps"`
}

type Server struct {
	tmpDir string
	cmd    *exec.Cmd
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: go run test_program.go <scenario.json>")
		os.Exit(1)
	}
	scenarioFile := os.Args[1]
	data, err := ioutil.ReadFile(scenarioFile)
	if err != nil {
		panic(err)
	}
	var scenario Scenario
	if err := json.Unmarshal(data, &scenario); err != nil {
		panic(err)
	}
	localServers := make(map[string]Server)
	for userKey, usr := range scenario.Users {
		if usr.Local {
			tmpDir, err := os.MkdirTemp("", "ripple-")
			if err != nil {
				panic(err)
			}
			configPath := tmpDir + "/config.json"
			if err := writeConfigFile(configPath, usr); err != nil {
				panic(err)
			}
			initCmd := exec.Command("./ripple", "init", configPath)
			initCmd.Env = append(os.Environ(), "HOME="+tmpDir)
			if err := initCmd.Run(); err != nil {
				panic(err)
			}
			serverCmd := exec.Command("./ripple")
			serverCmd.Env = append(os.Environ(), "HOME="+tmpDir)
			serverCmd.Stdout = os.Stdout
			serverCmd.Stderr = os.Stderr
			if err := serverCmd.Start(); err != nil {
				panic(err)
			}
			localServers[userKey] = Server{tmpDir, serverCmd}
		}
	}
	time.Sleep(200 * time.Millisecond)
	for i, step := range scenario.Steps {
		fmt.Printf("\n=== Step %d: %s ===\n", i+1, step.Title)
		user, ok := scenario.Users[step.User]
		if !ok {
			fmt.Printf("Error: no user named '%s'\n", step.User)
			os.Exit(1)
		}
		expandedArgs := make([]string, len(step.Args))
		for idx, arg := range step.Args {
			val, err := expandArg(arg, scenario)
			if err != nil {
				fmt.Printf("Error expanding arg '%s': %v\n", arg, err)
				os.Exit(1)
			}
			expandedArgs[idx] = val
		}
		user.Counter++
		scenario.Users[step.User] = user
		secretKeyBytes, err := hex.DecodeString(user.SecretKey)
		if err != nil {
			fmt.Printf("invalid user.SecretKey hex: %v\n", err)
			os.Exit(1)
		}
		var secretKey [32]byte
		copy(secretKey[:], secretKeyBytes)
		cmd, args, err := buildUserRequest(step.Cmd, expandedArgs, oneTimeKey(secretKey, user.Counter))
		if err != nil {
			fmt.Printf("Error building request: %v\n", err)
			os.Exit(1)
		}
		userRequest := testclient.UserRequest{
			Instruction: types.Instruction{
				Command:   cmd,
				Arguments: args,
			},
			Counter: user.Counter,
		}
		serverAddr := fmt.Sprintf("%s:%d", user.ServerAddress, user.Port)
		client := testclient.NewClient(serverAddr, secretKey)
		defer client.Close()
		status, data, err := client.SendRequest(userRequest)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		fmt.Printf("Client => %s %s\n", statusName(status), data)
		time.Sleep(500 * time.Millisecond)
	}
	sleepWithSpinner(20)
	for userKey, srv := range localServers {
		if storage, err := readState(srv.tmpDir); err == nil {
			printState(userKey, storage)
		} else {
			fmt.Printf("Could not read final state for %s: %v\n", userKey, err)
		}
		if srv.cmd.Process != nil {
			_ = srv.cmd.Process.Kill()
		}
		_ = os.RemoveAll(srv.tmpDir)
		fmt.Printf("Shut down server for %s and removed %s\n", userKey, srv.tmpDir)
	}
	fmt.Println("\nAll steps completed successfully!")
}

func writeConfigFile(path string, u User) error {
	cfg := map[string]interface{}{
		"port":          u.Port,
		"username":      u.Username,
		"serverAddress": u.ServerAddress,
		"secretKey":     u.SecretKey,
	}
	data, err := json.MarshalIndent(cfg, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, data, 0644)
}

func expandArg(arg string, scenario Scenario) (string, error) {
	if !strings.HasPrefix(arg, "@") {
		return arg, nil
	}
	ref := arg[1:]
	parts := strings.Split(ref, ".")
	if len(parts) != 2 {
		return "", fmt.Errorf("invalid reference '%s', must be @user.field", arg)
	}
	userKey, field := parts[0], parts[1]
	user, ok := scenario.Users[userKey]
	if !ok {
		return "", fmt.Errorf("no user '%s' in scenario", userKey)
	}
	switch field {
	case "username":
		return user.Username, nil
	case "serverAddress":
		return user.ServerAddress, nil
	case "port":
		return fmt.Sprintf("%d", user.Port), nil
	case "secretKey":
		return user.SecretKey, nil
	}
	return "", fmt.Errorf("unknown field '%s' in user '%s'", field, userKey)
}

func oneTimeKey(secretKey [32]byte, counter uint32) [32]byte {
	buf := make([]byte, 36)
	copy(buf[0:32], secretKey[:])
	binary.BigEndian.PutUint32(buf[32:], counter)
	return sha256.Sum256(buf)
}

func xor32(a, b [32]byte) [32]byte {
	var out [32]byte
	for i := 0; i < 32; i++ {
		out[i] = a[i] ^ b[i]
	}
	return out
}

func buildUserRequest(cmdStr string, argsStr []string, oneTimeKey [32]byte) (byte, []byte, error) {
	var cmd byte
	var args []byte
	var err error
	switch cmdStr {
	case "add_account":
		args, err = buildAddAccountArgs(argsStr, oneTimeKey)
		cmd = 0x00
	case "set_trustline":
		args, err = buildSetTrustlineArgs(argsStr)
		cmd = 0x01
	case "new_payment":
		args, err = buildNewPaymentArgs(argsStr, oneTimeKey)
		cmd = 0x02
	case "start_payment":
		cmd = 0x03
	default:
		err = fmt.Errorf("unknown cmd '%s'", cmdStr)
	}
	return cmd, args, err
}

func buildAddAccountArgs(args []string, oneTimeKey [32]byte) ([]byte, error) {
	if len(args) < 5 {
		return nil, fmt.Errorf("add_account requires 5 args: (username, serverAddress, port, secretKey, turnBit)")
	}
	username := args[0]
	serverAddress := args[1]
	portStr := args[2]
	secretKeyHex := args[3]
	turnBitStr := args[4]
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, err
	}
	secretKeyBytes, err := hex.DecodeString(secretKeyHex)
	if err != nil {
		return nil, err
	}
	var secretKey [32]byte
	copy(secretKey[0:32], secretKeyBytes)
	turnBit, err := strconv.Atoi(turnBitStr)
	if err != nil {
		return nil, err
	}
	ciphertext := xor32(secretKey, oneTimeKey)
	out := make([]byte, 99)
	copy(out[0:32], []byte(username))
	copy(out[32:64], []byte(serverAddress))
	binary.BigEndian.PutUint16(out[64:66], uint16(port))
	copy(out[66:98], ciphertext[:])
	out[98] = byte(turnBit)
	return out, nil
}

func buildSetTrustlineArgs(args []string) ([]byte, error) {
	if len(args) < 3 {
		return nil, fmt.Errorf("set_trustline requires 3 args: (username, serverAddress, value)")
	}
	username := args[0]
	serverAddress := args[1]
	valueStr := args[2]
	value, err := strconv.ParseUint(valueStr, 10, 64)
	if err != nil {
		return nil, err
	}
	out := make([]byte, 72)
	copy(out[0:32], []byte(username))
	copy(out[32:64], []byte(serverAddress))
	binary.BigEndian.PutUint64(out[64:72], value)
	return out, nil
}

func buildNewPaymentArgs(args []string, oneTimeKey [32]byte) ([]byte, error) {
	if len(args) < 5 {
		return nil, fmt.Errorf("new_payment requires 5 args: (username, serverAddress, portStr, secretKeyHex, amountStr)")
	}
	username := args[0]
	serverAddress := args[1]
	portStr := args[2]
	secretKeyHex := args[3]
	amountStr := args[4]
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, err
	}
	secretKeyBytes, err := hex.DecodeString(secretKeyHex)
	if err != nil {
		return nil, err
	}
	var secretKey [32]byte
	copy(secretKey[0:32], secretKeyBytes)
	amount, err := strconv.ParseInt(amountStr, 10, 64)
	if err != nil {
		return nil, err
	}
	ciphertext := xor32(secretKey, oneTimeKey)
	out := make([]byte, 106)
	copy(out[0:32], []byte(username))
	copy(out[32:64], []byte(serverAddress))
	binary.BigEndian.PutUint16(out[64:66], uint16(port))
	copy(out[66:98], ciphertext[:])
	binary.BigEndian.PutUint64(out[98:106], uint64(amount))
	return out, nil
}

func readState(tmpDir string) (*state.Storage, error) {
	f, err := os.Open(filepath.Join(tmpDir, ".ripple", "data.gob"))
	if err != nil {
		return nil, err
	}
	defer f.Close()
	dec := gob.NewDecoder(f)
	var s state.Storage
	if err := dec.Decode(&s); err != nil {
		return nil, err
	}
	return &s, nil
}

func printState(userKey string, s *state.Storage) {
	fmt.Printf("\n=== Final state for %s ===\n", userKey)
	for accID, acc := range s.Accounts {
		fmt.Printf(" Account %v\n", accID)
		fmt.Printf("  TurnBit=%d TurnCounter=%d Creditline=%d\n", acc.TurnBit, acc.TurnCounter, acc.Creditline)
		fmt.Printf("  TrustlineIn=%d TrustlineOut=%d\n", acc.TrustlineIn, acc.TrustlineOut)
		if len(acc.Pending) > 0 {
			fmt.Printf("  Pending: %v\n", acc.Pending)
		}
	}
	for payID, pay := range s.Payments {
		fmt.Printf(" Payment %x => Amount=%d, Outgoing=%v, Incoming=%v, Cancel=%v\n",
			payID, pay.Amount, pay.Outgoing, pay.Incoming, pay.Cancel)
	}
	fmt.Println(" Receipts:")
	for i, r := range s.Receipts {
		if r.Identifier == [32]byte{} && r.Timestamp == 0 {
			continue
		}
		fmt.Printf("  #%d => ID=%x, Counterpart=%v, Amount=%d, Timestamp=%d\n",
			i, r.Identifier, r.Counterpart, r.Amount, r.Timestamp)
	}
}
func sleepWithSpinner(seconds int) {
    spinnerChars := []rune{'|', '/', '-', '\\'}

	for i := 0; i < seconds; i++ {
		for _, c := range spinnerChars {
			fmt.Printf("\rGive server time to do payment... %02d seconds left %c", seconds-i, c)
			time.Sleep(250 * time.Millisecond)
		}
	}
    fmt.Print("\r \n")
}

func statusName(status byte) string {
    switch status {
    case 0x00:
        return "SUCCESS"
    case 0x01:
        return "ERROR"
    default:
        return "UNKNOWN"
    }
}