package probe

import (
	"strings"
	"unicode"
)

type AliasResult struct {
	Raw        string
	Normalized string
	Canonical  string
}

func NormalizeModelID(raw string) string {
	trimmed := strings.TrimSpace(strings.ToLower(raw))
	if trimmed == "" {
		return ""
	}

	if idx := strings.LastIndex(trimmed, "/"); idx >= 0 {
		trimmed = trimmed[idx+1:]
	}

	replacer := strings.NewReplacer("_", "-", " ", "-", ".", ".", "--", "-")
	normalized := replacer.Replace(trimmed)
	for strings.Contains(normalized, "--") {
		normalized = strings.ReplaceAll(normalized, "--", "-")
	}
	return strings.Trim(normalized, "-")
}

func CanonicalModelID(raw string) string {
	return NormalizeModelID(raw)
}

func CanonicalModelFamily(raw string) string {
	normalized := NormalizeModelID(raw)
	switch {
	case strings.HasPrefix(normalized, "kimi-k2."):
		return strings.Replace(normalized, "kimi-k2.", "kimi-2.", 1)
	case strings.HasPrefix(normalized, "kimi-k2-"):
		return strings.Replace(normalized, "kimi-k2-", "kimi-2-", 1)
	default:
		return normalized
	}
}

func BuildAliasTable(rawModels []string) map[string]AliasResult {
	table := make(map[string]AliasResult, len(rawModels)*4)
	for _, rawModel := range rawModels {
		rawModel = strings.TrimSpace(rawModel)
		if rawModel == "" {
			continue
		}

		result := AliasResult{
			Raw:        rawModel,
			Normalized: NormalizeModelID(rawModel),
			Canonical:  CanonicalModelFamily(rawModel),
		}

		keys := []string{
			rawModel,
			result.Normalized,
			result.Canonical,
			CanonicalModelID(rawModel),
			lookupKey(rawModel),
			lookupKey(result.Normalized),
			lookupKey(result.Canonical),
			lookupKey(CanonicalModelID(rawModel)),
		}
		for _, key := range keys {
			if key == "" {
				continue
			}
			if _, exists := table[key]; !exists {
				table[key] = result
			}
		}
	}
	return table
}

func ResolveRequestedModel(requested string, rawModels []string) (resolved string, ok bool) {
	result, ok := BuildAliasTable(rawModels)[lookupKey(requested)]
	if !ok {
		return "", false
	}
	return result.Raw, true
}

func RecommendModels(requested []string, rawModels []string) []string {
	table := BuildAliasTable(rawModels)
	recommended := make([]string, 0, len(requested))
	seen := make(map[string]struct{}, len(requested))

	for _, requestedModel := range requested {
		result, ok := table[lookupKey(requestedModel)]
		if !ok {
			continue
		}
		if _, exists := seen[result.Raw]; exists {
			continue
		}
		seen[result.Raw] = struct{}{}
		recommended = append(recommended, result.Raw)
	}

	return recommended
}

func lookupKey(raw string) string {
	canonical := CanonicalModelFamily(raw)
	if canonical == "" {
		return ""
	}

	var builder strings.Builder
	builder.Grow(len(canonical))
	for _, r := range canonical {
		if unicode.IsLetter(r) || unicode.IsDigit(r) {
			builder.WriteRune(unicode.ToLower(r))
		}
	}
	return builder.String()
}