internal/storage/sql/apikey_test.go

package sql_test

import (
	"context"
	"testing"

	"ccLoad/internal/model"
)

func TestAPIKey_CreateAndGet(t *testing.T) {
	t.Parallel()

	store := newTestStore(t, "apikey.db")

	ctx := context.Background()
	channelID := createTestChannel(t, ctx, store, "apikey-test-channel")

	// 批量创建 API Keys
	keys := []*model.APIKey{
		{ChannelID: channelID, KeyIndex: 0, APIKey: "sk-key-0", KeyStrategy: model.KeyStrategySequential},
		{ChannelID: channelID, KeyIndex: 1, APIKey: "sk-key-1", KeyStrategy: model.KeyStrategySequential},
		{ChannelID: channelID, KeyIndex: 2, APIKey: "sk-key-2", KeyStrategy: model.KeyStrategySequential},
	}
	if err := store.CreateAPIKeysBatch(ctx, keys); err != nil {
		t.Fatalf("create api keys batch: %v", err)
	}

	// 获取单个 API Key
	key, err := store.GetAPIKey(ctx, channelID, 1)
	if err != nil {
		t.Fatalf("get api key: %v", err)
	}
	if key.APIKey != "sk-key-1" {
		t.Errorf("api key: got %q, want %q", key.APIKey, "sk-key-1")
	}
	if key.KeyIndex != 1 {
		t.Errorf("key index: got %d, want %d", key.KeyIndex, 1)
	}

	// 获取渠道所有 API Keys
	allKeys, err := store.GetAPIKeys(ctx, channelID)
	if err != nil {
		t.Fatalf("get api keys: %v", err)
	}
	if len(allKeys) != 3 {
		t.Errorf("expected 3 keys, got %d", len(allKeys))
	}
}

func TestAPIKey_UpdateStrategy(t *testing.T) {
	t.Parallel()

	store := newTestStore(t, "strategy.db")

	ctx := context.Background()
	channelID := createTestChannel(t, ctx, store, "strategy-test-channel")

	// 创建 API Key
	keys := []*model.APIKey{
		{ChannelID: channelID, KeyIndex: 0, APIKey: "sk-key", KeyStrategy: model.KeyStrategySequential},
	}
	if err := store.CreateAPIKeysBatch(ctx, keys); err != nil {
		t.Fatalf("create api keys batch: %v", err)
	}

	// 更新策略
	if err := store.UpdateAPIKeysStrategy(ctx, channelID, model.KeyStrategyRoundRobin); err != nil {
		t.Fatalf("update api keys strategy: %v", err)
	}

	// 验证更新
	key, err := store.GetAPIKey(ctx, channelID, 0)
	if err != nil {
		t.Fatalf("get api key: %v", err)
	}
	if key.KeyStrategy != model.KeyStrategyRoundRobin {
		t.Errorf("strategy: got %q, want %q", key.KeyStrategy, model.KeyStrategyRoundRobin)
	}
}

func TestAPIKey_Delete(t *testing.T) {
	t.Parallel()

	store := newTestStore(t, "delete.db")

	ctx := context.Background()
	channelID := createTestChannel(t, ctx, store, "delete-test-channel")

	// 创建多个 API Keys
	keys := []*model.APIKey{
		{ChannelID: channelID, KeyIndex: 0, APIKey: "sk-key-0", KeyStrategy: model.KeyStrategySequential},
		{ChannelID: channelID, KeyIndex: 1, APIKey: "sk-key-1", KeyStrategy: model.KeyStrategySequential},
		{ChannelID: channelID, KeyIndex: 2, APIKey: "sk-key-2", KeyStrategy: model.KeyStrategySequential},
	}
	if err := store.CreateAPIKeysBatch(ctx, keys); err != nil {
		t.Fatalf("create api keys batch: %v", err)
	}

	// 删除中间的 key
	if err := store.DeleteAPIKey(ctx, channelID, 1); err != nil {
		t.Fatalf("delete api key: %v", err)
	}

	// 验证删除
	allKeys, err := store.GetAPIKeys(ctx, channelID)
	if err != nil {
		t.Fatalf("get api keys: %v", err)
	}
	if len(allKeys) != 2 {
		t.Errorf("expected 2 keys after delete, got %d", len(allKeys))
	}
}

func TestAPIKey_CompactIndices(t *testing.T) {
	t.Parallel()

	store := newTestStore(t, "compact.db")

	ctx := context.Background()
	channelID := createTestChannel(t, ctx, store, "compact-test-channel")

	// 创建 3 个 API Keys: indices 0, 1, 2
	keys := []*model.APIKey{
		{ChannelID: channelID, KeyIndex: 0, APIKey: "sk-key-0", KeyStrategy: model.KeyStrategySequential},
		{ChannelID: channelID, KeyIndex: 1, APIKey: "sk-key-1", KeyStrategy: model.KeyStrategySequential},
		{ChannelID: channelID, KeyIndex: 2, APIKey: "sk-key-2", KeyStrategy: model.KeyStrategySequential},
	}
	if err := store.CreateAPIKeysBatch(ctx, keys); err != nil {
		t.Fatalf("create api keys batch: %v", err)
	}

	// 删除 index=1 的 key
	if err := store.DeleteAPIKey(ctx, channelID, 1); err != nil {
		t.Fatalf("delete api key: %v", err)
	}

	// 压缩索引：将 index=2 移动到 index=1
	if err := store.CompactKeyIndices(ctx, channelID, 1); err != nil {
		t.Fatalf("compact key indices: %v", err)
	}

	// 验证压缩结果
	allKeys, err := store.GetAPIKeys(ctx, channelID)
	if err != nil {
		t.Fatalf("get api keys: %v", err)
	}
	if len(allKeys) != 2 {
		t.Errorf("expected 2 keys, got %d", len(allKeys))
	}

	// 检查索引是连续的
	for i, key := range allKeys {
		if key.KeyIndex != i {
			t.Errorf("key %d: expected index %d, got %d", i, i, key.KeyIndex)
		}
	}
}

func TestAPIKey_DeleteAll(t *testing.T) {
	t.Parallel()

	store := newTestStore(t, "delete_all.db")

	ctx := context.Background()
	channelID := createTestChannel(t, ctx, store, "delete-all-test-channel")

	// 创建多个 API Keys
	keys := []*model.APIKey{
		{ChannelID: channelID, KeyIndex: 0, APIKey: "sk-key-0", KeyStrategy: model.KeyStrategySequential},
		{ChannelID: channelID, KeyIndex: 1, APIKey: "sk-key-1", KeyStrategy: model.KeyStrategySequential},
	}
	if err := store.CreateAPIKeysBatch(ctx, keys); err != nil {
		t.Fatalf("create api keys batch: %v", err)
	}

	// 删除所有
	if err := store.DeleteAllAPIKeys(ctx, channelID); err != nil {
		t.Fatalf("delete all api keys: %v", err)
	}

	// 验证全部删除
	allKeys, err := store.GetAPIKeys(ctx, channelID)
	if err != nil {
		t.Fatalf("get api keys: %v", err)
	}
	if len(allKeys) != 0 {
		t.Errorf("expected 0 keys, got %d", len(allKeys))
	}
}

func TestAPIKey_GetAllAPIKeys(t *testing.T) {
	t.Parallel()

	store := newTestStore(t, "get_all.db")

	ctx := context.Background()

	// 创建两个渠道
	channelID1 := createTestChannel(t, ctx, store, "channel-1")
	channelID2 := createTestChannel(t, ctx, store, "channel-2")

	// 为每个渠道创建 API Keys
	keys1 := []*model.APIKey{
		{ChannelID: channelID1, KeyIndex: 0, APIKey: "sk-ch1-key-0", KeyStrategy: model.KeyStrategySequential},
		{ChannelID: channelID1, KeyIndex: 1, APIKey: "sk-ch1-key-1", KeyStrategy: model.KeyStrategySequential},
	}
	if err := store.CreateAPIKeysBatch(ctx, keys1); err != nil {
		t.Fatalf("create api keys for channel 1: %v", err)
	}

	keys2 := []*model.APIKey{
		{ChannelID: channelID2, KeyIndex: 0, APIKey: "sk-ch2-key-0", KeyStrategy: model.KeyStrategyRoundRobin},
	}
	if err := store.CreateAPIKeysBatch(ctx, keys2); err != nil {
		t.Fatalf("create api keys for channel 2: %v", err)
	}

	// 获取所有 API Keys（返回 map[channelID][]*APIKey）
	allKeys, err := store.GetAllAPIKeys(ctx)
	if err != nil {
		t.Fatalf("get all api keys: %v", err)
	}
	if countAPIKeys(allKeys) != 3 {
		t.Errorf("expected 3 total keys, got %d", countAPIKeys(allKeys))
	}
}

func TestAPIKey_ImportChannelBatch(t *testing.T) {
	t.Parallel()

	store := newTestStore(t, "import.db")

	ctx := context.Background()

	// 批量导入渠道（包含渠道配置和 API Keys）
	channels := []*model.ChannelWithKeys{
		{
			Config: &model.Config{
				Name:                  "imported-channel-1",
				URL:                   "https://api1.example.com",
				Priority:              10,
				Enabled:               true,
				ScheduledCheckEnabled: true,
				ChannelType:           "openai",
				ModelEntries: []model.ModelEntry{
					{Model: "gpt-4"},
					{Model: "gpt-3.5-turbo"},
				},
			},
			APIKeys: []model.APIKey{
				{KeyIndex: 0, APIKey: "sk-import-key-1", KeyStrategy: model.KeyStrategySequential},
				{KeyIndex: 1, APIKey: "sk-import-key-2", KeyStrategy: model.KeyStrategySequential},
			},
		},
		{
			Config: &model.Config{
				Name:                  "imported-channel-2",
				URL:                   "https://api2.example.com",
				Priority:              20,
				Enabled:               true,
				ScheduledCheckEnabled: false,
				ChannelType:           "anthropic",
				ModelEntries: []model.ModelEntry{
					{Model: "claude-3"},
				},
			},
			APIKeys: []model.APIKey{
				{KeyIndex: 0, APIKey: "sk-anthropic-key", KeyStrategy: model.KeyStrategyRoundRobin},
			},
		},
	}

	created, updated, err := store.ImportChannelBatch(ctx, channels)
	if err != nil {
		t.Fatalf("import channel batch: %v", err)
	}
	if created != 2 {
		t.Errorf("expected 2 created, got %d", created)
	}
	if updated != 0 {
		t.Errorf("expected 0 updated, got %d", updated)
	}

	// 验证导入结果
	configs, err := store.ListConfigs(ctx)
	if err != nil {
		t.Fatalf("list configs: %v", err)
	}
	if len(configs) != 2 {
		t.Errorf("expected 2 channels, got %d", len(configs))
	}

	configsByName := make(map[string]*model.Config, len(configs))
	for _, cfg := range configs {
		configsByName[cfg.Name] = cfg
	}
	if !configsByName["imported-channel-1"].ScheduledCheckEnabled {
		t.Fatalf("expected imported-channel-1 scheduled check enabled")
	}
	if configsByName["imported-channel-2"].ScheduledCheckEnabled {
		t.Fatalf("expected imported-channel-2 scheduled check disabled")
	}

	// 验证 API Keys 也被导入（渠道1=2个，渠道2=1个）
	allKeys, err := store.GetAllAPIKeys(ctx)
	if err != nil {
		t.Fatalf("get all api keys: %v", err)
	}
	if len(allKeys) != 2 {
		t.Fatalf("expected keys for 2 channels, got %d", len(allKeys))
	}
	if countAPIKeys(allKeys) != 3 {
		t.Fatalf("expected 3 api keys total, got %d", countAPIKeys(allKeys))
	}

	idsByName := make(map[string]int64, len(configs))
	for _, cfg := range configs {
		idsByName[cfg.Name] = cfg.ID
	}
	id1 := idsByName["imported-channel-1"]
	id2 := idsByName["imported-channel-2"]
	if id1 == 0 || id2 == 0 {
		t.Fatalf("expected imported channels to have non-zero ids, got id1=%d id2=%d", id1, id2)
	}

	keys1, err := store.GetAPIKeys(ctx, id1)
	if err != nil {
		t.Fatalf("get api keys for imported-channel-1: %v", err)
	}
	if len(keys1) != 2 {
		t.Fatalf("expected 2 keys for imported-channel-1, got %d", len(keys1))
	}
	keys2, err := store.GetAPIKeys(ctx, id2)
	if err != nil {
		t.Fatalf("get api keys for imported-channel-2: %v", err)
	}
	if len(keys2) != 1 {
		t.Fatalf("expected 1 key for imported-channel-2, got %d", len(keys2))
	}
}

func TestAPIKey_ImportChannelBatchPreservesScheduledCheckWithExplicitID(t *testing.T) {
	t.Parallel()

	store := newTestStore(t, "import-preserve-id.db")
	ctx := context.Background()

	channel := &model.ChannelWithKeys{
		Config: &model.Config{
			ID:                    42,
			Name:                  "imported-channel-explicit-id",
			URL:                   "https://api.example.com",
			Priority:              5,
			Enabled:               true,
			ScheduledCheckEnabled: true,
			ChannelType:           "openai",
			ModelEntries: []model.ModelEntry{
				{Model: "gpt-4o-mini"},
			},
		},
		APIKeys: []model.APIKey{
			{KeyIndex: 0, APIKey: "sk-import-key-explicit", KeyStrategy: model.KeyStrategySequential},
		},
	}

	created, updated, err := store.ImportChannelBatch(ctx, []*model.ChannelWithKeys{channel})
	if err != nil {
		t.Fatalf("import channel batch with explicit id: %v", err)
	}
	if created != 1 || updated != 0 {
		t.Fatalf("expected created=1 updated=0, got created=%d updated=%d", created, updated)
	}

	config, err := store.GetConfig(ctx, channel.Config.ID)
	if err != nil {
		t.Fatalf("get imported config: %v", err)
	}
	if !config.ScheduledCheckEnabled {
		t.Fatalf("expected scheduled_check_enabled to persist for explicit id import")
	}
}

func TestAPIKey_ImportChannelBatchPreservesModelEntryOrder(t *testing.T) {
	t.Parallel()

	store := newTestStore(t, "import-preserve-model-order.db")
	ctx := context.Background()

	channel := &model.ChannelWithKeys{
		Config: &model.Config{
			Name:        "imported-channel-order",
			URL:         "https://api.example.com",
			Priority:    5,
			Enabled:     true,
			ChannelType: "openai",
			ModelEntries: []model.ModelEntry{
				{Model: "z-last"},
				{Model: "a-first"},
			},
		},
		APIKeys: []model.APIKey{
			{KeyIndex: 0, APIKey: "sk-import-key-order", KeyStrategy: model.KeyStrategySequential},
		},
	}

	created, updated, err := store.ImportChannelBatch(ctx, []*model.ChannelWithKeys{channel})
	if err != nil {
		t.Fatalf("import channel batch: %v", err)
	}
	if created != 1 || updated != 0 {
		t.Fatalf("expected created=1 updated=0, got created=%d updated=%d", created, updated)
	}

	configs, err := store.ListConfigs(ctx)
	if err != nil {
		t.Fatalf("list configs: %v", err)
	}
	if len(configs) != 1 {
		t.Fatalf("expected 1 config, got %d", len(configs))
	}

	got := configs[0].ModelEntries
	if len(got) != 2 {
		t.Fatalf("expected 2 model entries, got %+v", got)
	}
	if got[0].Model != "z-last" || got[1].Model != "a-first" {
		t.Fatalf("expected model order [z-last a-first], got %+v", got)
	}
}