// Command gl-redistribution-loadtest proves balance conservation while applying
// a fixed number of random transfers between ten isolated test users.
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	mathrand "math/rand"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	basev1 "gl/gen/base/v1"
	ledgerv1 "gl/gen/ledger/v1"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/protobuf/types/known/timestamppb"
)

const (
	personCount         = 10
	initialSupply       = int64(10_000)
	minimumBalance      = int64(45)
	defaultTransactions = uint64(10_000_000)
	seedTransactions    = uint64(personCount - 1)
)

type client interface {
	Health(context.Context, *basev1.Empty, ...grpc.CallOption) (*ledgerv1.HealthResponse, error)
	AppendJournal(context.Context, *ledgerv1.AppendJournalRequest, ...grpc.CallOption) (*ledgerv1.AppendJournalResponse, error)
	GetBalance(context.Context, *ledgerv1.GetBalanceRequest, ...grpc.CallOption) (*ledgerv1.GetBalanceResponse, error)
}

type config struct {
	RunID          string
	AssetID        int64
	Transactions   uint64
	Concurrency    int
	RequestTimeout time.Duration
	RandomSeed     int64
}

type transfer struct {
	Sender   int
	Receiver int
	Amount   int64
}

type balanceState struct {
	mu       sync.Mutex
	balances [personCount]int64
}

type metrics struct {
	attempted    atomic.Uint64
	succeeded    atomic.Uint64
	errors       atomic.Uint64
	latencyNanos atomic.Int64
	maxNanos     atomic.Int64
}

func main() {
	os.Exit(runCLI(os.Args[1:], os.Stdin, os.Stdout, os.Stderr))
}

func runCLI(args []string, stdin io.Reader, stdout, stderr io.Writer) int {
	flags := flag.NewFlagSet("gl-redistribution-loadtest", flag.ContinueOnError)
	flags.SetOutput(stderr)
	target := flags.String("target", "127.0.0.1:8600", "GL gRPC target")
	assetID := flags.Int64("asset-id", 0, "asset number; prompted when omitted")
	transactions := flags.Uint64("transactions", defaultTransactions, "total user-to-user transfers")
	concurrency := flags.Int("concurrency", 32, "number of concurrent transfer workers")
	requestTimeout := flags.Duration("request-timeout", 10*time.Second, "timeout for each gRPC call")
	runID := flags.String("run-id", "", "isolated idempotency namespace; generated when empty")
	randomSeed := flags.Int64("seed", 0, "random seed; generated when zero")
	if err := flags.Parse(args); err != nil {
		return 2
	}
	if *transactions < seedTransactions || *concurrency < 1 || *requestTimeout <= 0 {
		fmt.Fprintf(stderr, "transactions must be at least %d; concurrency and request-timeout must be positive\n", seedTransactions)
		return 2
	}

	resolvedAssetID, err := resolveAssetID(*assetID, stdin, stdout)
	if err != nil {
		fmt.Fprintf(stderr, "asset number: %v\n", err)
		return 2
	}
	if *runID == "" {
		*runID = "gl-redistribution-" + strconv.FormatInt(time.Now().UTC().UnixNano(), 36)
	}
	if *randomSeed == 0 {
		*randomSeed = time.Now().UnixNano()
	}

	connection, err := grpc.NewClient(*target, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		fmt.Fprintf(stderr, "create gRPC client: %v\n", err)
		return 2
	}
	defer connection.Close()
	ledger := ledgerv1.NewGeneralLedgerServiceClient(connection)

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	cfg := config{
		RunID:          *runID,
		AssetID:        resolvedAssetID,
		Transactions:   *transactions,
		Concurrency:    *concurrency,
		RequestTimeout: *requestTimeout,
		RandomSeed:     *randomSeed,
	}
	if err := execute(ctx, ledger, cfg, stdout); err != nil {
		fmt.Fprintf(stderr, "redistribution load test failed: %v\n", err)
		return 1
	}
	return 0
}

func resolveAssetID(value int64, input io.Reader, output io.Writer) (int64, error) {
	if value > 0 {
		return value, nil
	}
	if value < 0 {
		return 0, fmt.Errorf("must be positive")
	}
	fmt.Fprint(output, "Asset number: ")
	var text string
	if _, err := fmt.Fscan(input, &text); err != nil {
		return 0, fmt.Errorf("read input: %w", err)
	}
	assetID, err := strconv.ParseInt(strings.TrimSpace(text), 10, 64)
	if err != nil || assetID <= 0 {
		return 0, fmt.Errorf("must be a positive integer")
	}
	return assetID, nil
}

func execute(ctx context.Context, ledger client, cfg config, output io.Writer) error {
	callCtx, cancel := context.WithTimeout(ctx, cfg.RequestTimeout)
	health, err := ledger.Health(callCtx, &basev1.Empty{})
	cancel()
	if err != nil {
		return fmt.Errorf("health check: %w", err)
	}
	if !health.GetServing() || !health.GetDatabaseReady() {
		return fmt.Errorf("GL is not ready: %v", health)
	}

	people := personIDs(cfg.RunID)
	startingBalances, err := readBalances(ctx, ledger, cfg, people)
	if err != nil {
		return fmt.Errorf("read unfunded balances: %w", err)
	}
	if err := verifyUnfunded(startingBalances); err != nil {
		return err
	}
	if err := fundFirstPerson(ctx, ledger, cfg, people[0]); err != nil {
		return fmt.Errorf("fund first person: %w", err)
	}
	initialBalances, err := readBalances(ctx, ledger, cfg, people)
	if err != nil {
		return fmt.Errorf("read initial balances: %w", err)
	}
	if err := verifyInitial(initialBalances); err != nil {
		return err
	}

	fmt.Fprintf(output, "GL redistribution load test: target users=%d asset=%d transfers=%d workers=%d seed=%d run=%s\n",
		personCount, cfg.AssetID, cfg.Transactions, cfg.Concurrency, cfg.RandomSeed, cfg.RunID)
	fmt.Fprintln(output, "initial balances: [10000 0 0 0 0 0 0 0 0 0]")

	stats := &metrics{}
	started := time.Now()
	balances := initialBalances
	for person := 1; person < personCount; person++ {
		value := transfer{Sender: 0, Receiver: person, Amount: minimumBalance}
		startedCall := time.Now()
		err := appendTransfer(ctx, ledger, cfg, uint64(person), people, value)
		stats.observe(time.Since(startedCall), err)
		if err != nil {
			return fmt.Errorf("seed person %d: %w", person+1, err)
		}
		balances[0] -= minimumBalance
		balances[person] += minimumBalance
	}

	state := &balanceState{balances: balances}
	if err := runRandomTransfers(ctx, ledger, cfg, people, state, stats, output); err != nil {
		printSummary(output, stats, time.Since(started))
		return err
	}
	elapsed := time.Since(started)
	printSummary(output, stats, elapsed)

	finalBalances, err := readBalances(ctx, ledger, cfg, people)
	if err != nil {
		return fmt.Errorf("read final balances: %w", err)
	}
	if err := verifyFinal(finalBalances); err != nil {
		return err
	}
	fmt.Fprintf(output, "final balances: %v\n", finalBalances)
	fmt.Fprintf(output, "PASS: all %d people have at least %d; total=%d\n", personCount, minimumBalance, initialSupply)
	return nil
}

func personIDs(runID string) [personCount]string {
	var people [personCount]string
	for index := range people {
		people[index] = fmt.Sprintf("%s-person-%02d", runID, index+1)
	}
	return people
}

func readBalances(ctx context.Context, ledger client, cfg config, people [personCount]string) ([personCount]int64, error) {
	var balances [personCount]int64
	for index, person := range people {
		callCtx, cancel := context.WithTimeout(ctx, cfg.RequestTimeout)
		response, err := ledger.GetBalance(callCtx, &ledgerv1.GetBalanceRequest{Account: personAccount(person, cfg.AssetID)})
		cancel()
		if err != nil {
			return balances, fmt.Errorf("person %d: %w", index+1, err)
		}
		if response == nil {
			return balances, fmt.Errorf("person %d: empty response", index+1)
		}
		balance, err := strconv.ParseInt(response.GetBalance(), 10, 64)
		if err != nil {
			return balances, fmt.Errorf("person %d: invalid integer balance %q", index+1, response.GetBalance())
		}
		balances[index] = balance
	}
	return balances, nil
}

func verifyUnfunded(balances [personCount]int64) error {
	for index, balance := range balances {
		if balance != 0 {
			return fmt.Errorf("person %d must start unfunded; got %d (use a fresh run-id)", index+1, balance)
		}
	}
	return nil
}

func verifyInitial(balances [personCount]int64) error {
	if balances[0] != initialSupply {
		return fmt.Errorf("person 1 initial balance: got %d, want %d", balances[0], initialSupply)
	}
	for index := 1; index < personCount; index++ {
		if balances[index] != 0 {
			return fmt.Errorf("person %d initial balance: got %d, want 0", index+1, balances[index])
		}
	}
	return nil
}

func verifyFinal(balances [personCount]int64) error {
	var total int64
	for index, balance := range balances {
		if balance < minimumBalance {
			return fmt.Errorf("person %d final balance %d is below %d", index+1, balance, minimumBalance)
		}
		total += balance
	}
	if total != initialSupply {
		return fmt.Errorf("final balance sum: got %d, want %d", total, initialSupply)
	}
	return nil
}

func fundFirstPerson(ctx context.Context, ledger client, cfg config, firstPerson string) error {
	request := &ledgerv1.AppendJournalRequest{
		SourceService:       "gl-redistribution-loadtest",
		IdempotencyKey:      "redistribution:" + cfg.RunID + ":fund:v1",
		SourceTransactionId: cfg.RunID + "-fund",
		TrackingCode:        cfg.RunID + "-fund",
		EffectKind:          "load-test-funding",
		EventVersion:        1,
		OccurredAt:          timestamppb.Now(),
		ActorId:             cfg.RunID,
		Metadata:            map[string]string{"load_test": cfg.RunID, "phase": "funding"},
		Entries: []*ledgerv1.JournalEntry{
			{LineNumber: 1, Account: treasuryAccount(cfg.RunID, cfg.AssetID), Amount: "-10000", Description: "load-test source"},
			{LineNumber: 2, Account: personAccount(firstPerson, cfg.AssetID), Amount: "10000", Description: "initial person balance"},
		},
	}
	callCtx, cancel := context.WithTimeout(ctx, cfg.RequestTimeout)
	defer cancel()
	_, err := ledger.AppendJournal(callCtx, request)
	return err
}

func runRandomTransfers(ctx context.Context, ledger client, cfg config, people [personCount]string, state *balanceState, stats *metrics, output io.Writer) error {
	if cfg.Transactions == seedTransactions {
		return nil
	}
	workerCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	var sequence atomic.Uint64
	sequence.Store(seedTransactions)
	var workers sync.WaitGroup
	var firstErr error
	var errorOnce sync.Once
	workers.Add(cfg.Concurrency)
	for worker := range cfg.Concurrency {
		go func(worker int) {
			defer workers.Done()
			random := mathrand.New(mathrand.NewSource(cfg.RandomSeed + int64(worker+1)*7919))
			for workerCtx.Err() == nil {
				number := sequence.Add(1)
				if number > cfg.Transactions {
					return
				}
				value := state.reserve(random)
				started := time.Now()
				err := appendTransfer(workerCtx, ledger, cfg, number, people, value)
				stats.observe(time.Since(started), err)
				if err != nil {
					state.rollback(value)
					errorOnce.Do(func() {
						firstErr = fmt.Errorf("transfer %d: %w", number, err)
						cancel()
					})
					return
				}
				succeeded := stats.succeeded.Load()
				if succeeded%1_000_000 == 0 {
					fmt.Fprintf(output, "progress: %d/%d transfers\n", succeeded, cfg.Transactions)
				}
			}
		}(worker)
	}
	workers.Wait()
	if firstErr != nil {
		return firstErr
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	if succeeded := stats.succeeded.Load(); succeeded != cfg.Transactions {
		return fmt.Errorf("completed %d transfers, want %d", succeeded, cfg.Transactions)
	}
	return nil
}

func (s *balanceState) reserve(random *mathrand.Rand) transfer {
	s.mu.Lock()
	defer s.mu.Unlock()

	eligible := make([]int, 0, personCount)
	for index, balance := range s.balances {
		if balance > minimumBalance {
			eligible = append(eligible, index)
		}
	}
	sender := eligible[random.Intn(len(eligible))]
	receiver := random.Intn(personCount - 1)
	if receiver >= sender {
		receiver++
	}
	spendable := s.balances[sender] - minimumBalance
	amount := random.Int63n(spendable) + 1
	s.balances[sender] -= amount
	s.balances[receiver] += amount
	return transfer{Sender: sender, Receiver: receiver, Amount: amount}
}

func (s *balanceState) rollback(value transfer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.balances[value.Sender] += value.Amount
	s.balances[value.Receiver] -= value.Amount
}

func appendTransfer(ctx context.Context, ledger client, cfg config, sequence uint64, people [personCount]string, value transfer) error {
	number := strconv.FormatUint(sequence, 10)
	amount := strconv.FormatInt(value.Amount, 10)
	transactionID := cfg.RunID + "-transfer-" + number
	request := &ledgerv1.AppendJournalRequest{
		SourceService:       "gl-redistribution-loadtest",
		IdempotencyKey:      "redistribution:" + cfg.RunID + ":transfer:" + number + ":v1",
		SourceTransactionId: transactionID,
		TrackingCode:        transactionID,
		EffectKind:          "random-internal-transfer",
		EventVersion:        1,
		OccurredAt:          timestamppb.Now(),
		CorrelationId:       cfg.RunID,
		ActorId:             people[value.Sender],
		Metadata:            map[string]string{"load_test": cfg.RunID, "sequence": number},
		Entries: []*ledgerv1.JournalEntry{
			{LineNumber: 1, Account: personAccount(people[value.Sender], cfg.AssetID), Amount: "-" + amount, Description: "random transfer debit"},
			{LineNumber: 2, Account: personAccount(people[value.Receiver], cfg.AssetID), Amount: amount, Description: "random transfer credit"},
		},
	}
	callCtx, cancel := context.WithTimeout(ctx, cfg.RequestTimeout)
	defer cancel()
	_, err := ledger.AppendJournal(callCtx, request)
	return err
}

func personAccount(ownerID string, assetID int64) *ledgerv1.AccountReference {
	return &ledgerv1.AccountReference{
		AccountClass: ledgerv1.AccountClass_ACCOUNT_CLASS_USER_AVAILABLE,
		OwnerType:    "user",
		OwnerId:      ownerID,
		AssetId:      assetID,
	}
}

func treasuryAccount(runID string, assetID int64) *ledgerv1.AccountReference {
	return &ledgerv1.AccountReference{
		AccountClass: ledgerv1.AccountClass_ACCOUNT_CLASS_TREASURY,
		OwnerType:    "load-test",
		OwnerId:      runID,
		AssetId:      assetID,
	}
}

func (m *metrics) observe(latency time.Duration, err error) {
	m.attempted.Add(1)
	m.latencyNanos.Add(latency.Nanoseconds())
	for {
		current := m.maxNanos.Load()
		if latency.Nanoseconds() <= current || m.maxNanos.CompareAndSwap(current, latency.Nanoseconds()) {
			break
		}
	}
	if err != nil {
		m.errors.Add(1)
		return
	}
	m.succeeded.Add(1)
}

func printSummary(output io.Writer, stats *metrics, elapsed time.Duration) {
	attempted := stats.attempted.Load()
	average := time.Duration(0)
	if attempted > 0 {
		average = time.Duration(stats.latencyNanos.Load() / int64(attempted))
	}
	rate := float64(0)
	if elapsed > 0 {
		rate = float64(stats.succeeded.Load()) / elapsed.Seconds()
	}
	fmt.Fprintf(output, "transfers: attempted=%d succeeded=%d errors=%d elapsed=%s rate=%.1f tx/s avg=%s max=%s\n",
		attempted, stats.succeeded.Load(), stats.errors.Load(), elapsed.Round(time.Millisecond), rate, average, time.Duration(stats.maxNanos.Load()))
}