package store

import "testing"

func TestFilterPlatformModelBillingConfigKeepsOnlyImagePricing(t *testing.T) {
	model := PlatformModel{
		ModelType: StringList{"image_generate"},
		BillingConfig: map[string]any{
			"text":       map[string]any{"basePrice": 0.01},
			"image":      map[string]any{"basePrice": 10},
			"image_edit": map[string]any{"basePrice": 12},
			"video":      map[string]any{"basePrice": 100},
		},
		BillingConfigOverride: map[string]any{
			"image": map[string]any{"basePrice": 8},
			"video": map[string]any{"basePrice": 80},
		},
	}

	filtered := FilterPlatformModelBillingConfig(model)

	assertHasKeys(t, filtered.BillingConfig, "image")
	assertMissingKeys(t, filtered.BillingConfig, "text", "image_edit", "video")
	assertHasKeys(t, filtered.BillingConfigOverride, "image")
	assertMissingKeys(t, filtered.BillingConfigOverride, "video")
}

func TestFilterPlatformModelBillingConfigKeepsImageEditOrImageFallback(t *testing.T) {
	model := PlatformModel{
		ModelType: StringList{"image_edit"},
		BillingConfig: map[string]any{
			"image": map[string]any{"basePrice": 10},
			"video": map[string]any{"basePrice": 100},
		},
	}

	filtered := FilterPlatformModelBillingConfig(model)

	assertHasKeys(t, filtered.BillingConfig, "image")
	assertMissingKeys(t, filtered.BillingConfig, "video")
}

func TestFilterPlatformModelBillingConfigKeepsVideoPricing(t *testing.T) {
	model := PlatformModel{
		ModelType: StringList{"video_generate", "image_to_video"},
		BillingConfig: map[string]any{
			"image":     map[string]any{"basePrice": 10},
			"video":     map[string]any{"basePrice": 100},
			"videoBase": 100,
		},
	}

	filtered := FilterPlatformModelBillingConfig(model)

	assertHasKeys(t, filtered.BillingConfig, "video", "videoBase")
	assertMissingKeys(t, filtered.BillingConfig, "image")
}

func TestFilterPlatformModelBillingConfigKeepsTextFlatPricing(t *testing.T) {
	model := PlatformModel{
		ModelType: StringList{"text_generate"},
		BillingConfig: map[string]any{
			"textInputPer1k":       0.01,
			"textCachedInputPer1k": 0.003,
			"textOutputPer1k":      0.02,
			"text_total": map[string]any{
				"basePrice":       0.01,
				"formulaConfig":   map[string]any{"inputTokenPrice": 0.01, "cachedInputTokenPrice": 0.003, "outputTokenPrice": 0.02},
				"dynamicWeight":   map[string]any{"cached": 0.5},
				"dimensionSchema": map[string]any{"unit": "1k_tokens"},
			},
			"inputTokenPrice":       0.01,
			"cachedInputTokenPrice": 0.003,
			"outputTokenPrice":      0.02,
			"image":                 map[string]any{"basePrice": 10},
		},
	}

	filtered := FilterPlatformModelBillingConfig(model)

	assertHasKeys(t, filtered.BillingConfig, "textInputPer1k", "textCachedInputPer1k", "textOutputPer1k", "text_total", "inputTokenPrice", "cachedInputTokenPrice", "outputTokenPrice")
	assertMissingKeys(t, filtered.BillingConfig, "image")
}

func TestFilterPlatformModelBillingConfigKeepsMusicPricing(t *testing.T) {
	model := PlatformModel{
		ModelType: StringList{"audio_generate"},
		BillingConfig: map[string]any{
			"music":     map[string]any{"basePrice": 6},
			"musicBase": 6,
			"audio":     map[string]any{"basePrice": 1},
			"image":     map[string]any{"basePrice": 10},
		},
	}

	filtered := FilterPlatformModelBillingConfig(model)

	assertHasKeys(t, filtered.BillingConfig, "music", "musicBase")
	assertMissingKeys(t, filtered.BillingConfig, "audio", "image")
}

func TestFilterPlatformModelBillingConfigKeepsSpeechAudioPricing(t *testing.T) {
	model := PlatformModel{
		ModelType: StringList{"text_to_speech"},
		BillingConfig: map[string]any{
			"audio":     map[string]any{"basePrice": 0.5},
			"audioBase": 0.5,
			"music":     map[string]any{"basePrice": 6},
			"video":     map[string]any{"basePrice": 100},
		},
	}

	filtered := FilterPlatformModelBillingConfig(model)

	assertHasKeys(t, filtered.BillingConfig, "audio", "audioBase")
	assertMissingKeys(t, filtered.BillingConfig, "music", "video")
}

func assertHasKeys(t *testing.T, value map[string]any, keys ...string) {
	t.Helper()
	for _, key := range keys {
		if _, ok := value[key]; !ok {
			t.Fatalf("expected key %q in %#v", key, value)
		}
	}
}

func assertMissingKeys(t *testing.T, value map[string]any, keys ...string) {
	t.Helper()
	for _, key := range keys {
		if _, ok := value[key]; ok {
			t.Fatalf("expected key %q to be filtered from %#v", key, value)
		}
	}
}
