package main

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"math/rand/v2"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/redis/go-redis/v9"
)

// keyCount is the number of seed:<n> keys for a given target/value size.
// Both seed and verify derive it the same way so they agree on the key set
// without storing a manifest (which the "step 2" actions could mutate).
func keyCount(targetBytes, valueBytes int) int64 {
	if valueBytes <= 0 {
		return 0
	}
	return int64((targetBytes + valueBytes - 1) / valueBytes)
}

// genValue deterministically fills buf with incompressible bytes derived from
// (runSeed, idx). It's the single source of truth shared by seed and verify:
// the same (runSeed, idx, len(buf)) always yields identical bytes, which is
// what lets verify regenerate the expected value instead of remembering it.
// ChaCha8 gives a fast, high-quality stream so the payload doesn't compress
// or dedup and thus reflects real memory.
func genValue(runSeed uint64, idx int64, buf []byte) {
	var seed [32]byte
	binary.LittleEndian.PutUint64(seed[0:8], runSeed)
	binary.LittleEndian.PutUint64(seed[8:16], uint64(idx))
	rand.NewChaCha8(seed).Read(buf)
}

// runSeed bulk-loads a single Valkey with deterministic random blobs until the
// on-wire payload reaches targetBytes. Meant for local use over the zcli VPN
// (e.g. --host valkey.zerops), so connection params come from the flags.
//
// Keys are seed:<n> (n in [0, keyCount)) each holding valueBytes generated by
// genValue(runSeed, n). Writes are pipelined in batches of batch keys to keep
// the round-trip count — and VPN latency — from dominating.
//
// throttleBytesPerSec, when > 0, caps the average write rate: after each batch
// the loop sleeps until cumulative bytes / rate seconds have elapsed, so the
// on-wire throughput converges to the cap without bursting a token bucket.
func runSeed(opts *redis.Options, runSeed uint64, targetBytes, valueBytes, batch int, throttleBytesPerSec float64) int {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	rdb := redis.NewClient(opts)
	defer rdb.Close()
	if err := rdb.Ping(ctx).Err(); err != nil {
		fmt.Fprintf(os.Stderr, "seed: connect %s failed: %v\n", opts.Addr, err)
		return 1
	}

	total := keyCount(targetBytes, valueBytes)
	throttle := "unlimited"
	if throttleBytesPerSec > 0 {
		throttle = humanBytes(int(throttleBytesPerSec)) + "/s"
	}
	fmt.Printf("seed: addr=%s db=%d tls=%t target=%s value=%s keys=%d seed=%d batch=%d throttle=%s\n",
		opts.Addr, opts.DB, opts.TLSConfig != nil,
		humanBytes(targetBytes), humanBytes(valueBytes), total, runSeed, batch, throttle)

	start, lastLog := time.Now(), time.Now()
	var written, keys int64

	for keys < total {
		select {
		case <-ctx.Done():
			fmt.Printf("\nseed: interrupted — wrote %d/%d keys (%s)\n", keys, total, time.Since(start).Round(time.Second))
			return 1
		default:
		}

		pipe := rdb.Pipeline()
		for i := 0; i < batch && keys < total; i++ {
			// Fresh buffer per key: go-redis keeps the []byte by reference and
			// only serializes it at Exec, so a reused buffer would store the
			// batch's last value under every key.
			buf := make([]byte, valueBytes)
			genValue(runSeed, keys, buf)
			pipe.Set(ctx, fmt.Sprintf("seed:%d", keys), buf, 0)
			keys++
			written += int64(valueBytes)
		}
		if _, err := pipe.Exec(ctx); err != nil {
			fmt.Fprintf(os.Stderr, "seed: pipeline exec: %v\n", err)
			return 1
		}

		// Throttle on cumulative progress: sleep until enough wall-clock has
		// passed for everything written so far at the target rate. Interruptible
		// so Ctrl+C doesn't wait out a long sleep.
		if throttleBytesPerSec > 0 {
			targetElapsed := time.Duration(float64(written) / throttleBytesPerSec * float64(time.Second))
			if sleep := targetElapsed - time.Since(start); sleep > 0 {
				select {
				case <-ctx.Done():
				case <-time.After(sleep):
				}
			}
		}

		if now := time.Now(); now.Sub(lastLog) >= time.Second {
			rate := float64(written) / now.Sub(start).Seconds()
			fmt.Printf("[%s] %s / %s  (%d/%d keys, %s/s)\n",
				now.Format("15:04:05"), humanBytes(int(written)), humanBytes(targetBytes), keys, total, humanBytes(int(rate)))
			lastLog = now
		}
	}

	elapsed := time.Since(start)
	fmt.Printf("seed: done — wrote %s across %d keys in %s (%s/s). verify with the same --seed-mb/--seed-value-bytes/--seed-seed.\n",
		humanBytes(int(written)), keys, elapsed.Round(time.Millisecond), humanBytes(int(float64(written)/elapsed.Seconds())))
	return 0
}

// runVerify regenerates every seed:<n> value and compares it byte-for-byte to
// what's in Valkey, reporting missing/mismatched keys and any unexpected
// extras (via DBSIZE). It must be run with the same runSeed/targetBytes/
// valueBytes used to seed. Returns non-zero if the data is not 1:1.
func runVerify(opts *redis.Options, runSeed uint64, targetBytes, valueBytes, batch int) int {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	rdb := redis.NewClient(opts)
	defer rdb.Close()
	if err := rdb.Ping(ctx).Err(); err != nil {
		fmt.Fprintf(os.Stderr, "verify: connect %s failed: %v\n", opts.Addr, err)
		return 1
	}

	total := keyCount(targetBytes, valueBytes)
	fmt.Printf("verify: addr=%s db=%d tls=%t expect keys=%d value=%s seed=%d\n",
		opts.Addr, opts.DB, opts.TLSConfig != nil, total, humanBytes(valueBytes), runSeed)

	start, lastLog := time.Now(), time.Now()
	var checked, missing, mismatch int64
	expect := make([]byte, valueBytes)

	for base := int64(0); base < total; base += int64(batch) {
		select {
		case <-ctx.Done():
			fmt.Printf("\nverify: interrupted at %d/%d keys\n", checked, total)
			return 1
		default:
		}

		end := base + int64(batch)
		if end > total {
			end = total
		}
		pipe := rdb.Pipeline()
		cmds := make([]*redis.StringCmd, 0, end-base)
		for idx := base; idx < end; idx++ {
			cmds = append(cmds, pipe.Get(ctx, fmt.Sprintf("seed:%d", idx)))
		}
		if _, err := pipe.Exec(ctx); err != nil && err != redis.Nil {
			fmt.Fprintf(os.Stderr, "verify: pipeline exec: %v\n", err)
			return 1
		}

		for i, cmd := range cmds {
			idx := base + int64(i)
			got, err := cmd.Bytes()
			if err == redis.Nil {
				missing++
				if missing <= 10 {
					fmt.Printf("  MISSING seed:%d\n", idx)
				}
				continue
			}
			if err != nil {
				fmt.Fprintf(os.Stderr, "verify: get seed:%d: %v\n", idx, err)
				return 1
			}
			genValue(runSeed, idx, expect)
			if !bytes.Equal(got, expect) {
				mismatch++
				if mismatch <= 10 {
					fmt.Printf("  MISMATCH seed:%d (len got=%d want=%d)\n", idx, len(got), len(expect))
				}
			}
			checked++
		}

		if now := time.Now(); now.Sub(lastLog) >= time.Second {
			fmt.Printf("[%s] verified %d/%d keys\n", now.Format("15:04:05"), base+int64(len(cmds)), total)
			lastLog = now
		}
	}

	dbsize, err := rdb.DBSize(ctx).Result()
	if err != nil {
		fmt.Fprintf(os.Stderr, "verify: dbsize: %v\n", err)
		return 1
	}
	extras := dbsize - total

	fmt.Printf("verify: %d ok, %d missing, %d mismatched of %d expected; dbsize=%d (%+d vs expected) in %s\n",
		checked-mismatch, missing, mismatch, total, dbsize, extras, time.Since(start).Round(time.Millisecond))

	if missing == 0 && mismatch == 0 && extras == 0 {
		fmt.Println("verify: PASS — data is 1:1 with the seed.")
		return 0
	}
	fmt.Println("verify: FAIL — data diverged from the seed.")
	return 1
}

// runFlush clears the whole database selected by opts.DB with FLUSHDB,
// reporting how many keys it dropped. Only the one logical DB is touched —
// other databases on the same server are left intact.
func runFlush(opts *redis.Options) int {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	rdb := redis.NewClient(opts)
	defer rdb.Close()
	if err := rdb.Ping(ctx).Err(); err != nil {
		fmt.Fprintf(os.Stderr, "flush: connect %s failed: %v\n", opts.Addr, err)
		return 1
	}

	before, err := rdb.DBSize(ctx).Result()
	if err != nil {
		fmt.Fprintf(os.Stderr, "flush: dbsize db%d: %v\n", opts.DB, err)
		return 1
	}
	if err := rdb.FlushDB(ctx).Err(); err != nil {
		fmt.Fprintf(os.Stderr, "flush: flushdb db%d: %v\n", opts.DB, err)
		return 1
	}
	fmt.Printf("flush: FLUSHDB db%d on %s — cleared %d keys\n", opts.DB, opts.Addr, before)
	return 0
}

// runInspect prints every logical database and its key count. The number of
// databases comes from CONFIG GET databases (default 16); per-DB key counts
// come from a single INFO keyspace call, which lists only the non-empty DBs —
// the rest are reported as 0. Empty DBs are still shown so the full layout is
// visible at a glance.
func runInspect(opts *redis.Options) int {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	rdb := redis.NewClient(opts)
	defer rdb.Close()
	if err := rdb.Ping(ctx).Err(); err != nil {
		fmt.Fprintf(os.Stderr, "inspect: connect %s failed: %v\n", opts.Addr, err)
		return 1
	}

	numDBs := 16
	if res, err := rdb.ConfigGet(ctx, "databases").Result(); err == nil {
		if v, ok := res["databases"]; ok {
			if n, err := strconv.Atoi(v); err == nil && n > 0 {
				numDBs = n
			}
		}
	}

	info, err := rdb.Info(ctx, "keyspace").Result()
	if err != nil {
		fmt.Fprintf(os.Stderr, "inspect: info keyspace: %v\n", err)
		return 1
	}
	// Lines look like: db0:keys=12,expires=0,avg_ttl=0
	counts := make(map[int]int64)
	for _, line := range strings.Split(info, "\n") {
		line = strings.TrimSpace(line)
		if !strings.HasPrefix(line, "db") {
			continue
		}
		colon := strings.IndexByte(line, ':')
		if colon < 0 {
			continue
		}
		dbn, err := strconv.Atoi(line[2:colon])
		if err != nil {
			continue
		}
		for _, field := range strings.Split(line[colon+1:], ",") {
			if rest, ok := strings.CutPrefix(field, "keys="); ok {
				if k, err := strconv.ParseInt(rest, 10, 64); err == nil {
					counts[dbn] = k
				}
			}
		}
	}

	fmt.Printf("inspect: addr=%s tls=%t databases=%d\n", opts.Addr, opts.TLSConfig != nil, numDBs)
	var total int64
	for db := 0; db < numDBs; db++ {
		k := counts[db]
		total += k
		fmt.Printf("  db%-2d  %d keys\n", db, k)
	}
	fmt.Printf("inspect: %d keys across %d databases\n", total, numDBs)
	return 0
}

func humanBytes(n int) string {
	const unit = 1024
	if n < unit {
		return fmt.Sprintf("%dB", n)
	}
	div, exp := int64(unit), 0
	for v := n / unit; v >= unit; v /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f%cB", float64(n)/float64(div), "KMGTPE"[exp])
}