package main

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"github.com/alecthomas/kong"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/pelletier/go-toml"
	"github.com/tyler-smith/go-bip39"
	"golang.org/x/crypto/sha3"
	"log"
	"os"
)

var cli struct {
	Seeds     []string `arg:"" optional:"" type:"string" help:"Deterministic seeds or BIP-39 mnemonics to generate keys."`
	Json      bool     `optional:"" short:"j" help:"Output results in JSON format"`
	ConfigFile string  `optional:"" short:"c" type:"path" help:"Path to TOML config file containing seeds"`
}

type Config struct {
	Seeds []string `toml:"seeds"`
}

type WalletInfo struct {
	PrivateKey     string `json:"private_key"`
	PublicKey      string `json:"public_key"`
	EthereumAddress string `json:"ethereum_address"`
}

func readTomlConfig(filePath string) ([]string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return nil, fmt.Errorf("error opening config file: %w", err)
	}
	defer file.Close()

	config := &Config{}
	if err := toml.NewDecoder(file).Decode(config); err != nil {
		return nil, fmt.Errorf("error decoding TOML: %w", err)
	}

	return config.Seeds, nil
}

func generateWallet(seed []byte) (WalletInfo, error) {
	hash := sha3.NewLegacyKeccak256()
	hash.Write(seed)
	buf := hash.Sum(nil)

	privateKey, err := crypto.ToECDSA(buf)
	if err != nil {
		return WalletInfo{}, fmt.Errorf("failed to create private key: %w", err)
	}

	publicKey := privateKey.Public()
	publicKeyECDSA, ok := publicKey.(*crypto.PublicKeyECDSA)
	if !ok {
		return WalletInfo{}, fmt.Errorf("error casting public key to ECDSA")
	}

	address := crypto.PubkeyToAddress(*publicKeyECDSA).Hex()
	return WalletInfo{
		PrivateKey:     hex.EncodeToString(crypto.FromECDSA(privateKey)),
		PublicKey:      hex.EncodeToString(crypto.FromECDSAPub(publicKeyECDSA)),
		EthereumAddress: address,
	}, nil
}

func main() {
	kong.Parse(&cli)

	var seeds []string
	var err error

	if cli.ConfigFile != "" {
		seeds, err = readTomlConfig(cli.ConfigFile)
		if err != nil {
			log.Fatalf("Failed to read TOML config: %v", err)
		}
	} else {
		seeds = cli.Seeds
	}

	var walletInfos []WalletInfo

	for _, seedStr := range seeds {
		var seed []byte
		if bip39.IsMnemonicValid(seedStr) {
			seed, err = bip39.NewSeedWithErrorChecking(seedStr, "")
			if err != nil {
				log.Printf("Failed to create seed from mnemonic: %v", err)
				continue
			}
		} else {
			seed = []byte(seedStr)
		}

		walletInfo, err := generateWallet(seed)
		if err != nil {
			log.Printf("Error generating wallet: %v", err)
			continue
		}
		walletInfos = append(walletInfos, walletInfo)
	}

	if cli.Json {
		jsonOutput, err := json.Marshal(walletInfos)
		if err != nil {
			log.Fatalf("JSON Marshalling failed: %v", err)
		}
		fmt.Println(string(jsonOutput))
	} else {
		for _, walletInfo := range walletInfos {
			fmt.Printf("Private Key: %s\nPublic Key: %s\nEthereum Address: %s\n",
				walletInfo.PrivateKey, walletInfo.PublicKey, walletInfo.EthereumAddress)
		}
	}
}