package main

import (
	"flag"
	"fmt"
	"os"
	"sort"
	"strings"
	"sync"
	"time"
)

type Mode string

const (
	ModePrivateKeys Mode = "private-keys"
	ModeAxiom       Mode = "axiom"
	ModeBoth        Mode = "both"
)

type CheckResult struct {
	Line    string
	Balance float64
}

type Statistics struct {
	TotalBalance    float64
	KeysWithBalance int
	EmptyKeys       int
	TotalKeys       int
	Duration        time.Duration
	Mode            string
}

var (
	results      []CheckResult
	resultsMutex sync.Mutex
	stats        Statistics
)

func addResult(line string, balance float64) {
	resultsMutex.Lock()
	defer resultsMutex.Unlock()
	results = append(results, CheckResult{Line: line, Balance: balance})

	stats.TotalBalance += balance
	if balance > 0 {
		stats.KeysWithBalance++
	} else {
		stats.EmptyKeys++
	}
}

func main() {
	addressFlag := flag.String("address", "", "Check specific address")
	flag.Parse()

	printHeader()

	cfg, err := loadConfig()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to load config: %v\n", err)
		os.Exit(1)
	}

	initProxies()

	if *addressFlag != "" {
		checkSingleAddress(*addressFlag, cfg)
		return
	}

	for {
		choice, err := showMainMenu(cfg)
		if err != nil {
			fmt.Fprintf(os.Stderr, "\nCancelled\n")
			os.Exit(1)
		}

		switch choice {
		case "settings":
			if err := showSettingsMenu(cfg); err != nil {
				fmt.Fprintf(os.Stderr, "\nSettings cancelled\n")
				continue
			}

		case "start":
			initApiKeyWorkers(cfg)

			fmt.Println()
			printConfigSummary(cfg)

			moduleManager := NewModuleManager(cfg)
			enabledModules := moduleManager.GetEnabledModules()

			if cfg.EnableSolanaCheck || len(enabledModules) > 0 {
				colorPrint(ColorGreen, "Active modules:      ")
				modules := []string{}
				if cfg.EnableSolanaCheck {
					modules = append(modules, "Solana")
				}
				for _, m := range enabledModules {
					modules = append(modules, m.Name())
				}
				for i, m := range modules {
					if i > 0 {
						fmt.Print(", ")
					}
					colorPrint(ColorCyan, m)
				}
				fmt.Println()
			}
			fmt.Printf("Loaded proxies:      %d\n", len(proxies))

			keys := loadPrivateKeys()
			if len(keys) == 0 {
				fmt.Fprintln(os.Stderr, "\nNo keys found!\nAdd to: private_keys.txt or keys/*.txt")
				continue
			}
			stats.TotalKeys = len(keys)

			runChecker(keys, cfg, moduleManager)

			fmt.Println("\nPress Enter to return to main menu...")
			fmt.Scanln()
			printHeader()

		case "exit":
			colorPrintln(ColorYellow, "\nGoodbye!")
			os.Exit(0)
		}
	}
}

func printHeader() {
	colorPrintln(ColorCyan, "")
	colorPrintln(ColorCyan, "   ___      _                  ___ _           _           ")
	colorPrintln(ColorCyan, "  / _ \\__ _(_)___ _ __  ___   / __| |_  ___ __| |_____ _ _ ")
	colorPrintln(ColorCyan, " | (_| \\ \\ / / _ \\ '  \\/ -_) | (__| ' \\/ -_) _| / / -_) '_|")
	colorPrintln(ColorCyan, "  \\___/_\\_\\_\\\\___/_|_|_\\___|  \\___|_||_\\___\\__|_\\_\\___|_|  ")
	fmt.Print("                                              ")
	colorPrintln(ColorYellow, "by @septum")
	colorPrintln(ColorCyan, "")
}

func checkSingleAddress(address string, cfg *Config) {
	if strings.HasPrefix(address, "0x") || strings.HasPrefix(address, "0X") {
		balance, details, err := getDeBankBalance(address)
		if err != nil {
			fmt.Printf("[evm address][debank] %s - Error: %v\n", address, err)
			return
		}

		fmt.Printf("[evm address][debank] %s - $%.2f\n", address, balance)
		if details != "" {
			fmt.Print(details)
		}
	} else {
		portfolio, err := getPortfolioValue(address, cfg)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
			os.Exit(1)
		}

		tokenStr := ""
		if portfolio.Value > 0 {
			if portfolio.TokenCount > 10 {
				tokenStr = " (>10 tokens)"
			} else {
				tokenStr = fmt.Sprintf(" (%d tokens)", portfolio.TokenCount)
			}
		}

		valueStr := "0.00"
		if portfolio.Value > 0 {
			valueStr = fmt.Sprintf("%.2f", portfolio.Value)
		}

		fmt.Printf("[solana address][jup.ag] %s - $%s%s\n", address, valueStr, tokenStr)
	}
}

func runChecker(keys []string, cfg *Config, moduleManager *ModuleManager) {
	outputFile := fmt.Sprintf("results_%s.txt", time.Now().Format("2006-01-02_15-04-05"))

	fmt.Printf("\n%s\n", strings.Repeat("━", 60))
	fmt.Printf("Keys to check: %d\n", len(keys))
	fmt.Printf("Output file: %s\n", outputFile)
	fmt.Printf("%s\n\n", strings.Repeat("━", 60))

	startTime := time.Now()
	processKeys(keys, cfg, moduleManager)
	stats.Duration = time.Since(startTime)

	sort.Slice(results, func(i, j int) bool { return results[i].Balance > results[j].Balance })

	if err := writeResults(outputFile); err != nil {
		fmt.Fprintf(os.Stderr, "Failed to write results: %v\n", err)
		return
	}

	fmt.Printf("\n\n")
	colorPrintln(ColorGreen, fmt.Sprintf("✅ Results saved to: %s", outputFile))
	printSummary()
}

func processKeys(keys []string, cfg *Config, moduleManager *ModuleManager) {
	tracker := NewProgressTracker(len(keys))
	display := NewStatusDisplay(tracker)
	display.Start()
	defer display.Stop()

	for i := 0; i < len(keys); i += cfg.Threads {
		end := min(i+cfg.Threads, len(keys))
		var wg sync.WaitGroup
		for _, key := range keys[i:end] {
			wg.Add(1)
			go func(pk string) {
				defer wg.Done()
				checkKey(pk, cfg, moduleManager, tracker)
			}(key)
		}
		wg.Wait()
	}
}

func checkKey(pk string, cfg *Config, moduleManager *ModuleManager, tracker *ProgressTracker) {
	defer func() {
		if r := recover(); r != nil {
			line := fmt.Sprintf("Error: %v", r)
			addResult(line, 0)
			tracker.Increment(0)
		}
	}()

	var line string
	totalBalance := 0.0

	if cfg.EnableSolanaCheck {
		kp, err := keypairFromSecretKey(pk)
		if err != nil {
			addResult(fmt.Sprintf("Error: %v", err), 0)
			tracker.Increment(0)
			return
		}
		addr := kp.PublicKeyBase58()
		pkBalance, _ := getPortfolioValue(addr, cfg)

		line += fmt.Sprintf("[Private Key] %s\n", pk)
		line += fmt.Sprintf("  [SOL] %s - $%.2f", addr, pkBalance.Value)
		if pkBalance.TokenCount > 0 {
			line += fmt.Sprintf(" (%d tokens)", pkBalance.TokenCount)
		}
		line += "\n"

		totalBalance += pkBalance.Value
	} else {
		line += fmt.Sprintf("[Private Key] %s\n", pk)
	}

	moduleBalance, moduleDetails := moduleManager.CheckAll(pk)
	if moduleDetails != "" {
		line += moduleDetails
		totalBalance += moduleBalance
	}

	if totalBalance > 0 {
		line += fmt.Sprintf("  %s[TOTAL: $%.2f]%s\n", ColorYellow, totalBalance, ColorReset)
	}

	addResult(line, totalBalance)
	tracker.Increment(totalBalance)
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func printSummary() {
	fmt.Println("\n" + strings.Repeat("=", 60))
	fmt.Println("SUMMARY")
	fmt.Println(strings.Repeat("=", 60))
	fmt.Printf("Mode:               %s\n", stats.Mode)
	fmt.Printf("Total Balance:      $%.2f\n", stats.TotalBalance)
	fmt.Printf("Keys with Balance:  %d\n", stats.KeysWithBalance)
	fmt.Printf("Empty Keys:         %d\n", stats.EmptyKeys)
	fmt.Printf("Total Keys:         %d\n", stats.TotalKeys)
	fmt.Printf("Duration:           %s\n", stats.Duration.Round(time.Second))
	fmt.Println(strings.Repeat("=", 60))
}

func writeResults(filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()

	fmt.Fprintf(f, "RESULTS\n")
	fmt.Fprintf(f, "Generated: %s\n\n", time.Now().Format("2006-01-02 15:04:05"))
	fmt.Fprintf(f, "Mode:               %s\n", stats.Mode)
	fmt.Fprintf(f, "Total Balance:      $%.2f\n", stats.TotalBalance)
	fmt.Fprintf(f, "Accounts with Balance:  %d\n", stats.KeysWithBalance)
	fmt.Fprintf(f, ":0.01$> accounts%d\n", stats.EmptyKeys)
	fmt.Fprintf(f, "Total Keys:         %d\n", stats.TotalKeys)
	fmt.Fprintf(f, "Duration:           %s\n", stats.Duration.Round(time.Second))
	fmt.Fprintf(f, "\n%s\n\n", strings.Repeat("=", 80))

	// Write sorted results
	for _, result := range results {
		fmt.Fprintln(f, result.Line)
	}

	return nil
}

func check(pk string, mode Mode, cfg *Config) {
	defer func() {
		if r := recover(); r != nil {
			line := fmt.Sprintf("Error: %v", r)
			fmt.Println(line)
			addResult(line, 0)
		}
	}()

	if mode == ModePrivateKeys {
		kp, err := keypairFromSecretKey(pk)
		if err != nil {
			line := fmt.Sprintf("Error: %v", err)
			fmt.Println(line)
			addResult(line, 0)
			return
		}

		addr := kp.PublicKeyBase58()
		portfolio, err := getPortfolioValue(addr, cfg)
		if err != nil {
			line := fmt.Sprintf("Error: %v", err)
			fmt.Println(line)
			addResult(line, 0)
			return
		}

		tokenStr := ""
		if portfolio.Value > 0 {
			if portfolio.TokenCount > 10 {
				tokenStr = " (>10 tokens)"
			} else {
				tokenStr = fmt.Sprintf(" (%d tokens)", portfolio.TokenCount)
			}
		}

		valueStr := "0.00"
		if portfolio.Value > 0 {
			valueStr = fmt.Sprintf("%.2f", portfolio.Value)
		}

		line := fmt.Sprintf("[private key][jup.ag] %s | %s - $%s%s", pk, addr, valueStr, tokenStr)
		fmt.Println(line)
		addResult(line, portfolio.Value)
		return
	}

	ax, err := getAxiomWallet(pk)
	if err != nil {
		line := fmt.Sprintf("Error: %v", err)
		fmt.Println(line)
		addResult(line, 0)
		return
	}

	if ax.IsNewUser {
		line := fmt.Sprintf("[axiom][jup.ag] %s | %s - NEW USER (skipped)", pk, ax.Sol)
		fmt.Println(line)
		addResult(line, 0)
		return
	}

	if mode == ModeAxiom {
		portfolio, err := getPortfolioValue(ax.Sol, cfg)
		if err != nil {
			line := fmt.Sprintf("Error: %v", err)
			fmt.Println(line)
			addResult(line, 0)
			return
		}

		tokenStr := ""
		if portfolio.Value > 0 {
			if portfolio.TokenCount > 10 {
				tokenStr = " (>10 tokens)"
			} else {
				tokenStr = fmt.Sprintf(" (%d tokens)", portfolio.TokenCount)
			}
		}

		valueStr := "0.00"
		if portfolio.Value > 0 {
			valueStr = fmt.Sprintf("%.2f", portfolio.Value)
		}

		line := fmt.Sprintf("[axiom][jup.ag] %s | %s - $%s%s", pk, ax.Sol, valueStr, tokenStr)
		fmt.Println(line)
		addResult(line, portfolio.Value)
		return
	}

	if mode == ModeBoth {
		kp, err := keypairFromSecretKey(pk)
		if err != nil {
			line := fmt.Sprintf("Error: %v", err)
			fmt.Println(line)
			addResult(line, 0)
			return
		}

		addr := kp.PublicKeyBase58()

		var p1, p2 *PortfolioValue
		var wg sync.WaitGroup
		var err1, err2 error

		wg.Add(2)

		go func() {
			defer wg.Done()
			p1, err1 = getPortfolioValue(addr, cfg)
		}()

		go func() {
			defer wg.Done()
			p2, err2 = getPortfolioValue(ax.Sol, cfg)
		}()

		wg.Wait()

		if err1 != nil || err2 != nil {
			line := "Error fetching portfolios"
			fmt.Println(line)
			addResult(line, 0)
			return
		}

		total := p1.Value + p2.Value
		count := p1.TokenCount + p2.TokenCount

		tokenStr := ""
		if total > 0 {
			if count > 10 {
				tokenStr = " (>10 tokens)"
			} else {
				tokenStr = fmt.Sprintf(" (%d tokens)", count)
			}
		}

		valueStr := "0.00"
		if total > 0 {
			valueStr = fmt.Sprintf("%.2f", total)
		}

		line := fmt.Sprintf("[pk+axiom][jup.ag] %s | %s | %s - $%s%s", pk, addr, ax.Sol, valueStr, tokenStr)
		fmt.Println(line)
		addResult(line, total)
	}
}

func tern(cond bool, t, f string) string {
	if cond {
		return t
	}
	return f
}