package db

import (
	"context"
	"testing"

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

type pageScanner struct {
	steps map[uint64]struct {
		keys []string
		next uint64
	}
	visited []uint64
}

func (s *pageScanner) Scan(ctx context.Context, cursor uint64, _ string, _ int64) *redis.ScanCmd {
	s.visited = append(s.visited, cursor)
	step := s.steps[cursor]
	cmd := redis.NewScanCmd(ctx, nil, "scan")
	cmd.SetVal(step.keys, step.next)
	return cmd
}

func TestKeyPageResumesAcrossMastersWithoutDroppingBatch(t *testing.T) {
	first := &pageScanner{steps: map[uint64]struct {
		keys []string
		next uint64
	}{
		0:  {[]string{"a", "b"}, 17},
		17: {[]string{"c"}, 0},
	}}
	second := &pageScanner{steps: map[uint64]struct {
		keys []string
		next uint64
	}{
		0: {[]string{"d"}, 0},
	}}
	clients := []scanClient{first, second}
	// Simulate a page already at the next cursor of the first shard.
	keys, pos, err := scanKeyPage(context.Background(), clients, keyScanPosition{Pattern: "*", Index: 0, Cursor: 17})
	if err != nil || len(keys) != 2 || keys[0] != "c" || keys[1] != "d" || pos.Index != 2 || pos.Cursor != 0 {
		t.Fatalf("keys=%v pos=%+v err=%v", keys, pos, err)
	}
	if len(first.visited) != 1 || first.visited[0] != 17 || len(second.visited) != 1 || second.visited[0] != 0 {
		t.Fatalf("unexpected scan positions first=%v second=%v", first.visited, second.visited)
	}
}

func TestKeyPageStopsAtBudgetWithResumableCursor(t *testing.T) {
	s := &pageScanner{steps: map[uint64]struct {
		keys []string
		next uint64
	}{
		0: {nil, 1}, 1: {nil, 1},
	}}
	keys, pos, err := scanKeyPage(context.Background(), []scanClient{s}, keyScanPosition{Pattern: "rare*"})
	if err != nil || len(keys) != 0 || len(s.visited) != 20 || pos.Index != 0 || pos.Cursor != 1 {
		t.Fatalf("keys=%v calls=%d pos=%+v err=%v", keys, len(s.visited), pos, err)
	}
}

func TestKeyPageKeepsOversizedScanBatch(t *testing.T) {
	batch := make([]string, redisKeyListCap+3)
	for i := range batch {
		batch[i] = "key"
	}
	s := &pageScanner{steps: map[uint64]struct {
		keys []string
		next uint64
	}{
		0: {batch, 9}, 9: {[]string{"last"}, 0},
	}}
	keys, pos, err := scanKeyPage(context.Background(), []scanClient{s}, keyScanPosition{Pattern: "*"})
	if err != nil || len(keys) != len(batch) || pos.Cursor != 9 || pos.Index != 0 {
		t.Fatalf("first page keys=%d pos=%+v err=%v", len(keys), pos, err)
	}
	nextKeys, end, err := scanKeyPage(context.Background(), []scanClient{s}, pos)
	if err != nil || len(nextKeys) != 1 || nextKeys[0] != "last" || end.Index != 1 {
		t.Fatalf("second page keys=%v pos=%+v err=%v", nextKeys, end, err)
	}
}