package store

import "strings"

func FilterPlatformModelBillingConfigs(models []PlatformModel) []PlatformModel {
	filtered := make([]PlatformModel, len(models))
	for i, model := range models {
		filtered[i] = FilterPlatformModelBillingConfig(model)
	}
	return filtered
}

func FilterPlatformModelBillingConfig(model PlatformModel) PlatformModel {
	model.BillingConfig = filterBillingConfigByModelTypes(model.BillingConfig, model.ModelType)
	model.BillingConfigOverride = filterBillingConfigByModelTypes(model.BillingConfigOverride, model.ModelType)
	return model
}

func filterBillingConfigByModelTypes(config map[string]any, modelTypes []string) map[string]any {
	if len(config) == 0 {
		return nil
	}
	resources := billingResourcesForModelTypes(modelTypes)
	if len(resources) == 0 {
		return cloneBillingConfig(config)
	}

	filtered := map[string]any{}
	for key, value := range config {
		if billingConfigKeyAllowed(key, resources) {
			filtered[key] = value
		}
	}

	if resources["image_edit"] && !hasAnyKey(filtered, "image_edit", "editBase") && !hasAnyKey(config, "image_edit", "editBase") {
		copyBillingConfigKey(filtered, config, "image")
		copyBillingConfigKey(filtered, config, "imageBase")
	}

	if len(filtered) == 0 {
		return nil
	}
	return filtered
}

func billingResourcesForModelTypes(modelTypes []string) map[string]bool {
	resources := map[string]bool{}
	for _, modelType := range modelTypes {
		switch normalizeBillingType(modelType) {
		case "chat", "text", "responses", "text_generate", "text_embedding", "embedding", "text_rerank", "rerank",
			"image_analysis", "video_understanding", "audio_understanding", "omni", "tools_call":
			resources["text"] = true
		case "image", "images.generations", "image_generate":
			resources["image"] = true
		case "images.edits", "image_edit":
			resources["image_edit"] = true
		case "image_vectorize", "images.vectorize":
			resources["image_vectorize"] = true
		case "video", "videos.generations", "video_generate", "image_to_video", "text_to_video",
			"video_edit", "omni_video", "video_reference", "video_first_last_frame":
			resources["video"] = true
		case "video_enhance", "videos.upscales", "video_upscale":
			resources["video_enhance"] = true
		case "audio", "text_to_speech", "speech", "voice_clone":
			resources["audio"] = true
		case "music", "music_generate", "audio_generate":
			resources["music"] = true
		case "digital_human", "digital_human_generate":
			resources["digital_human"] = true
		case "model", "model_3d", "text_to_model", "image_to_model", "multiview_to_model", "mesh_edit":
			resources["model"] = true
		default:
			inferBillingResources(modelType, resources)
		}
	}
	return resources
}

func inferBillingResources(modelType string, resources map[string]bool) {
	normalized := normalizeBillingType(modelType)
	switch {
	case strings.Contains(normalized, "digital_human"):
		resources["digital_human"] = true
	case strings.Contains(normalized, "video"):
		resources["video"] = true
	case strings.Contains(normalized, "image"):
		resources["image"] = true
	case strings.Contains(normalized, "audio") || strings.Contains(normalized, "speech"):
		resources["audio"] = true
	case strings.Contains(normalized, "music"):
		resources["music"] = true
	case strings.Contains(normalized, "model") || strings.Contains(normalized, "mesh"):
		resources["model"] = true
	case strings.Contains(normalized, "text") || strings.Contains(normalized, "token"):
		resources["text"] = true
	}
}

func billingConfigKeyAllowed(key string, resources map[string]bool) bool {
	switch normalizeBillingConfigKey(key) {
	case "text", "texttotal", "text_total", "textinputper1k", "textcachedinputper1k", "textoutputper1k", "textinput", "textcachedinput", "textoutput", "text_input", "text_cached_input", "text_output", "inputtokenprice", "cachedinputtokenprice", "inputcachehittokenprice", "outputtokenprice":
		return resources["text"]
	case "image", "imagebase":
		return resources["image"]
	case "image_edit", "imageedit", "editbase":
		return resources["image_edit"]
	case "image_vectorize", "imagevectorize", "vectorizebase":
		return resources["image_vectorize"]
	case "video", "videobase":
		return resources["video"]
	case "video_enhance", "videoenhance", "videoenhancebase":
		return resources["video_enhance"]
	case "audio", "audiobase":
		return resources["audio"]
	case "music", "musicbase":
		return resources["music"]
	case "digital_human", "digitalhuman", "digitalhumanbase":
		return resources["digital_human"]
	case "model", "modelbase":
		return resources["model"]
	default:
		return false
	}
}

func normalizeBillingType(value string) string {
	return strings.ToLower(strings.TrimSpace(value))
}

func normalizeBillingConfigKey(value string) string {
	return strings.ReplaceAll(strings.ReplaceAll(normalizeBillingType(value), "-", "_"), " ", "")
}

func cloneBillingConfig(config map[string]any) map[string]any {
	clone := make(map[string]any, len(config))
	for key, value := range config {
		clone[key] = value
	}
	return clone
}

func hasAnyKey(config map[string]any, keys ...string) bool {
	for _, key := range keys {
		if _, ok := config[key]; ok {
			return true
		}
	}
	return false
}

func copyBillingConfigKey(target map[string]any, source map[string]any, key string) {
	if value, ok := source[key]; ok {
		target[key] = value
	}
}
