package managerserver

import (
	"strings"
	"testing"

	"github.com/tidbcloud/aws-shared-provider/pkg/tidbmanager"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/clientfake"
)

type testWalkInfo struct {
	endpoints []string
	replicas  int
}

type testTidbManager struct {
	clientfake.FakeTidbManager
	active map[string]testWalkInfo
}

func (tm *testTidbManager) WalkActiveTokenFromAssignCache(f func(token string, endpoints []string, replicas int) bool) {
	for token, info := range tm.active {
		if !f(token, info.endpoints, info.replicas) {
			return
		}
	}
}

func TestCollectActiveTiDBAddresses(t *testing.T) {
	managers := map[string]tidbmanager.TidbManager{
		"tidb-pool": &testTidbManager{
			active: map[string]testWalkInfo{
				"essential-a": {
					endpoints: []string{"pod-a,10.0.0.1"},
					replicas:  1,
				},
				"essential-b/rg1": {
					endpoints: []string{"pod-b,10.0.0.2"},
					replicas:  1,
				},
				"free-c": {
					endpoints: []string{"pod-c,10.0.0.3"},
					replicas:  1,
				},
			},
		},
	}

	// keep only essential tokens, also covers keyspace name trimming.
	keep := func(token string) bool {
		return strings.HasPrefix(utils.TrimToken(token), "essential")
	}

	got := collectActiveTiDBAddresses(managers, keep)

	want := []activeTiDBAddress{
		{
			KeyspaceName: "essential-a",
			Hostname:     "pod-a",
			Port:         activeTiDBPort,
			Host:         "10.0.0.1",
			StatusPort:   activeTiDBStatusPort,
		},
		{
			KeyspaceName: "essential-b",
			Hostname:     "pod-b",
			Port:         activeTiDBPort,
			Host:         "10.0.0.2",
			StatusPort:   activeTiDBStatusPort,
		},
	}

	if len(got) != len(want) {
		t.Fatalf("unexpected response length: got %d want %d (%+v)", len(got), len(want), got)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("unexpected response at %d: got %+v want %+v", i, got[i], want[i])
		}
	}
}