package keyspacepoolconfig

import (
	"context"
	"fmt"
	"strings"
	"testing"

	serverlessv1alpha1 "github.com/tidbcloud/aws-shared-provider/apis/serverless/v1alpha1"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

const testNamespace = "keyspace-pool"

type defaultingClient struct {
	client.Client
	created int
}

type countingWaitLimiter struct {
	waits int
}

func (l *countingWaitLimiter) Wait(context.Context) error {
	l.waits++
	return nil
}

func (c *defaultingClient) Create(ctx context.Context, obj client.Object, opts ...client.CreateOption) error {
	if keyspace, ok := obj.(*serverlessv1alpha1.Keyspace); ok {
		c.created++
		keyspace.Name = fmt.Sprintf("kt-created-%d", c.created)
		keyspace.Status.Phase = serverlessv1alpha1.PendingPhase
		keyspace.Status.InPool = true
	}
	return c.Client.Create(ctx, obj, opts...)
}

func TestKeyspacesToCreate(t *testing.T) {
	tests := []struct {
		name       string
		poolSize   int
		inPool     int
		inProgress int
		maxPending int
		want       int
	}{
		{
			name:       "pool already meets target",
			poolSize:   25,
			inPool:     25,
			maxPending: 15,
		},
		{
			name:       "in-progress limit already reached",
			poolSize:   100000,
			inPool:     25,
			inProgress: 15,
			maxPending: 15,
		},
		{
			name:       "fills remaining in-progress slots",
			poolSize:   100000,
			inPool:     25,
			inProgress: 12,
			maxPending: 15,
			want:       3,
		},
		{
			name:       "does not exceed pool shortfall",
			poolSize:   30,
			inPool:     28,
			maxPending: 15,
			want:       2,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := keyspacesToCreate(tt.poolSize, tt.inPool, tt.inProgress, tt.maxPending)
			if got != tt.want {
				t.Fatalf("keyspacesToCreate() = %d, want %d", got, tt.want)
			}
		})
	}
}

func TestReconcileLimitsInProgressKeyspaces(t *testing.T) {
	ctx := context.Background()
	createRateStore := NewCreateRateStore(0)
	objects := []client.Object{
		&corev1.ConfigMap{
			ObjectMeta: metav1.ObjectMeta{Name: KeyspacePoolConfigName, Namespace: testNamespace},
			Data: map[string]string{
				KeyspacePoolSizeKey:       "12",
				KeyspacePoolMaxPendingKey: "5",
			},
		},
	}
	objects = append(objects,
		keyspacesWithPhase(serverlessv1alpha1.AvailablePhase, 2)...)
	objects = append(objects,
		keyspacesWithPhase(serverlessv1alpha1.PendingPhase, 1)...)
	objects = append(objects,
		keyspacesWithPhase(serverlessv1alpha1.CreatingPhase, 1)...)
	objects = append(objects,
		keyspacesWithPhase(serverlessv1alpha1.ActivatingPhase, 1)...)

	kubeClient := fake.NewClientBuilder().
		WithScheme(scheme.Scheme).
		WithObjects(objects...).
		WithIndex(&serverlessv1alpha1.Keyspace{}, "status.inPool", func(obj client.Object) []string {
			if obj.(*serverlessv1alpha1.Keyspace).Status.InPool {
				return []string{"true"}
			}
			return nil
		}).
		Build()
	defaultingClient := &defaultingClient{Client: kubeClient}
	r := &Reconciler{
		Client:          defaultingClient,
		Namespace:       testNamespace,
		KubeClusterName: "staging-us-east-1-f01",
		CreateRateStore: createRateStore,
	}
	req := ctrl.Request{NamespacedName: client.ObjectKey{
		Name:      KeyspacePoolConfigName,
		Namespace: testNamespace,
	}}

	if _, err := r.Reconcile(ctx, req); err != nil {
		t.Fatalf("first Reconcile() error = %v", err)
	}
	if defaultingClient.created != 2 {
		t.Fatalf("first Reconcile() created %d keyspaces, want 2", defaultingClient.created)
	}

	if _, err := r.Reconcile(ctx, req); err != nil {
		t.Fatalf("second Reconcile() error = %v", err)
	}
	if defaultingClient.created != 2 {
		t.Fatalf("second Reconcile() created %d total keyspaces, want 2", defaultingClient.created)
	}

	setCreatedKeyspacePhase(t, ctx, kubeClient, 1, 2, serverlessv1alpha1.AvailablePhase)

	if _, err := r.Reconcile(ctx, req); err != nil {
		t.Fatalf("third Reconcile() error = %v", err)
	}
	if defaultingClient.created != 4 {
		t.Fatalf("third Reconcile() created %d total keyspaces, want 4", defaultingClient.created)
	}

	setCreatedKeyspacePhase(t, ctx, kubeClient, 3, 4, serverlessv1alpha1.AvailablePhase)

	if _, err := r.Reconcile(ctx, req); err != nil {
		t.Fatalf("fourth Reconcile() error = %v", err)
	}
	if defaultingClient.created != 6 {
		t.Fatalf("fourth Reconcile() created %d total keyspaces, want 6", defaultingClient.created)
	}

	setCreatedKeyspacePhase(t, ctx, kubeClient, 5, 6, serverlessv1alpha1.AvailablePhase)

	if _, err := r.Reconcile(ctx, req); err != nil {
		t.Fatalf("fifth Reconcile() error = %v", err)
	}
	if defaultingClient.created != 7 {
		t.Fatalf("fifth Reconcile() created %d total keyspaces, want 7", defaultingClient.created)
	}

	if _, err := r.Reconcile(ctx, req); err != nil {
		t.Fatalf("sixth Reconcile() error = %v", err)
	}
	if defaultingClient.created != 7 {
		t.Fatalf("sixth Reconcile() created %d total keyspaces, want 7", defaultingClient.created)
	}
}

func setCreatedKeyspacePhase(
	t *testing.T,
	ctx context.Context,
	kubeClient client.Client,
	first, last int,
	phase serverlessv1alpha1.PhaseStatus,
) {
	t.Helper()
	for i := first; i <= last; i++ {
		keyspace := &serverlessv1alpha1.Keyspace{}
		key := client.ObjectKey{Name: fmt.Sprintf("kt-created-%d", i), Namespace: testNamespace}
		if err := kubeClient.Get(ctx, key, keyspace); err != nil {
			t.Fatalf("Get(%s) error = %v", key.Name, err)
		}
		keyspace.Status.Phase = phase
		if err := kubeClient.Update(ctx, keyspace); err != nil {
			t.Fatalf("Update(%s) error = %v", key.Name, err)
		}
	}
}

func TestParseKeyspacePoolConfig(t *testing.T) {
	tests := []struct {
		name                    string
		data                    map[string]string
		wantSize                int
		wantMaxPending          int
		wantCreateRatePerSecond float64
		wantErr                 string
	}{
		{
			name: "uses configured value",
			data: map[string]string{
				KeyspacePoolSizeKey:       "25",
				KeyspacePoolMaxPendingKey: "5",
			},
			wantSize:                25,
			wantMaxPending:          5,
			wantCreateRatePerSecond: 2.5,
		},
		{
			name:                    "defaults when field is missing",
			data:                    map[string]string{KeyspacePoolSizeKey: "25"},
			wantSize:                25,
			wantMaxPending:          100,
			wantCreateRatePerSecond: 0,
		},
		{
			name:    "rejects invalid size",
			data:    map[string]string{KeyspacePoolSizeKey: "invalid"},
			wantErr: "parse size",
		},
		{
			name:    "rejects invalid integer",
			data:    map[string]string{KeyspacePoolSizeKey: "25", KeyspacePoolMaxPendingKey: "invalid"},
			wantErr: "parse max-pending",
		},
		{
			name:    "rejects zero",
			data:    map[string]string{KeyspacePoolSizeKey: "25", KeyspacePoolMaxPendingKey: "0"},
			wantErr: "max-pending must be greater than 0",
		},
		{
			name:    "rejects negative value",
			data:    map[string]string{KeyspacePoolSizeKey: "25", KeyspacePoolMaxPendingKey: "-1"},
			wantErr: "max-pending must be greater than 0",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := parseKeyspacePoolConfig(&corev1.ConfigMap{Data: tt.data}, tt.wantCreateRatePerSecond)
			if tt.wantErr != "" {
				if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
					t.Fatalf("parseKeyspacePoolConfig() error = %v, want containing %q", err, tt.wantErr)
				}
				return
			}
			if err != nil {
				t.Fatalf("parseKeyspacePoolConfig() error = %v", err)
			}
			if got.poolSize != tt.wantSize || got.maxPending != tt.wantMaxPending ||
				got.createRatePerSecond != tt.wantCreateRatePerSecond {
				t.Fatalf("parseKeyspacePoolConfig() = (%d, %d, %v), want (%d, %d, %v)",
					got.poolSize, got.maxPending, got.createRatePerSecond,
					tt.wantSize, tt.wantMaxPending, tt.wantCreateRatePerSecond)
			}
		})
	}
}

func TestParseCreateRatePerSecond(t *testing.T) {
	tests := []struct {
		name     string
		data     map[string]string
		wantRate float64
		wantErr  string
	}{
		{
			name: "uses configured value",
			data: map[string]string{
				KeyspacePoolCreateRatePerSecondKey: "2.5",
			},
			wantRate: 2.5,
		},
		{
			name:     "defaults when field is missing",
			data:     map[string]string{},
			wantRate: 0,
		},
		{
			name: "rejects invalid create rate",
			data: map[string]string{
				KeyspacePoolCreateRatePerSecondKey: "invalid",
			},
			wantErr: "parse create-rate-per-second",
		},
		{
			name: "rejects negative create rate",
			data: map[string]string{
				KeyspacePoolCreateRatePerSecondKey: "-1",
			},
			wantErr: "create-rate-per-second must be greater than or equal to 0",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := parseCreateRatePerSecond(&corev1.ConfigMap{Data: tt.data})
			if tt.wantErr != "" {
				if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
					t.Fatalf("parseCreateRatePerSecond() error = %v, want containing %q", err, tt.wantErr)
				}
				return
			}
			if err != nil {
				t.Fatalf("parseCreateRatePerSecond() error = %v", err)
			}
			if got != tt.wantRate {
				t.Fatalf("parseCreateRatePerSecond() = %v, want %v", got, tt.wantRate)
			}
		})
	}
}

func TestReconcileWaitsForEachCreatedKeyspace(t *testing.T) {
	ctx := context.Background()
	// Seed the runtime store with the startup rate; the ConfigMap value should
	// not override what the controller already has in memory.
	createRateStore := NewCreateRateStore(2.5)
	objects := []client.Object{
		&corev1.ConfigMap{
			ObjectMeta: metav1.ObjectMeta{Name: KeyspacePoolConfigName, Namespace: testNamespace},
			Data: map[string]string{
				KeyspacePoolSizeKey:                "3",
				KeyspacePoolMaxPendingKey:          "3",
				KeyspacePoolCreateRatePerSecondKey: "99",
			},
		},
	}

	kubeClient := fake.NewClientBuilder().
		WithScheme(scheme.Scheme).
		WithObjects(objects...).
		WithIndex(&serverlessv1alpha1.Keyspace{}, "status.inPool", func(obj client.Object) []string {
			if obj.(*serverlessv1alpha1.Keyspace).Status.InPool {
				return []string{"true"}
			}
			return nil
		}).
		Build()
	defaultingClient := &defaultingClient{Client: kubeClient}
	countingLimiter := &countingWaitLimiter{}
	createLimiter := newCreateRateLimiter()
	var gotRateLimit float64
	createLimiter.newLimiter = func(limit float64) waitLimiter {
		gotRateLimit = limit
		return countingLimiter
	}

	r := &Reconciler{
		Client:          defaultingClient,
		Namespace:       testNamespace,
		KubeClusterName: "staging-us-east-1-f01",
		CreateRateStore: createRateStore,
		rateLimiter:     createLimiter,
	}
	req := ctrl.Request{NamespacedName: client.ObjectKey{
		Name:      KeyspacePoolConfigName,
		Namespace: testNamespace,
	}}

	if _, err := r.Reconcile(ctx, req); err != nil {
		t.Fatalf("Reconcile() error = %v", err)
	}
	if defaultingClient.created != 3 {
		t.Fatalf("Reconcile() created %d keyspaces, want 3", defaultingClient.created)
	}
	if countingLimiter.waits != 3 {
		t.Fatalf("rate limiter Wait() called %d times, want 3", countingLimiter.waits)
	}
	if gotRateLimit != 2.5 {
		t.Fatalf("rate limiter configured with %v, want 2.5", gotRateLimit)
	}
}

func TestReconcileUsesInMemoryCreateRate(t *testing.T) {
	ctx := context.Background()
	// This test proves the runtime store wins even if the ConfigMap contains a
	// different value.
	createRateStore := NewCreateRateStore(1.25)
	createRateStore.Set(2.5)
	objects := []client.Object{
		&corev1.ConfigMap{
			ObjectMeta: metav1.ObjectMeta{Name: KeyspacePoolConfigName, Namespace: testNamespace},
			Data: map[string]string{
				KeyspacePoolSizeKey:                "1",
				KeyspacePoolMaxPendingKey:          "1",
				KeyspacePoolCreateRatePerSecondKey: "99",
			},
		},
	}

	kubeClient := fake.NewClientBuilder().
		WithScheme(scheme.Scheme).
		WithObjects(objects...).
		WithIndex(&serverlessv1alpha1.Keyspace{}, "status.inPool", func(obj client.Object) []string {
			if obj.(*serverlessv1alpha1.Keyspace).Status.InPool {
				return []string{"true"}
			}
			return nil
		}).
		Build()
	defaultingClient := &defaultingClient{Client: kubeClient}
	countingLimiter := &countingWaitLimiter{}
	createLimiter := newCreateRateLimiter()
	var gotRateLimit float64
	createLimiter.newLimiter = func(limit float64) waitLimiter {
		gotRateLimit = limit
		return countingLimiter
	}

	r := &Reconciler{
		Client:          defaultingClient,
		Namespace:       testNamespace,
		KubeClusterName: "staging-us-east-1-f01",
		CreateRateStore: createRateStore,
		rateLimiter:     createLimiter,
	}
	req := ctrl.Request{NamespacedName: client.ObjectKey{
		Name:      KeyspacePoolConfigName,
		Namespace: testNamespace,
	}}

	if _, err := r.Reconcile(ctx, req); err != nil {
		t.Fatalf("Reconcile() error = %v", err)
	}
	if gotRateLimit != 2.5 {
		t.Fatalf("rate limiter configured with %v, want 2.5", gotRateLimit)
	}
	if countingLimiter.waits != 1 {
		t.Fatalf("rate limiter Wait() called %d times, want 1", countingLimiter.waits)
	}
}

func keyspacesWithPhase(phase serverlessv1alpha1.PhaseStatus, count int) []client.Object {
	keyspaces := make([]client.Object, 0, count)
	phaseName := strings.ToLower(string(phase))
	for i := 0; i < count; i++ {
		keyspaces = append(keyspaces, &serverlessv1alpha1.Keyspace{
			ObjectMeta: metav1.ObjectMeta{
				Name:      fmt.Sprintf("kt-%s-%d", phaseName, i),
				Namespace: testNamespace,
			},
			Spec: serverlessv1alpha1.KeyspaceSpec{
				KeyspaceName: fmt.Sprintf("token-%s-%d", phaseName, i),
			},
			Status: serverlessv1alpha1.KeyspaceStatus{
				Phase:  phase,
				InPool: true,
			},
		})
	}
	return keyspaces
}