internal/api/modules/amp/model_mapping_test.go

package amp

import (
	"testing"

	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
)

func TestNewModelMapper(t *testing.T) {
	mappings := []config.AmpModelMapping{
		{From: "claude-opus-4.5", To: "claude-sonnet-4"},
		{From: "gpt-5", To: "gemini-2.5-pro"},
	}

	mapper := NewModelMapper(mappings)
	if mapper == nil {
		t.Fatal("Expected non-nil mapper")
	}

	result := mapper.GetMappings()
	if len(result) != 2 {
		t.Errorf("Expected 2 mappings, got %d", len(result))
	}
}

func TestNewModelMapper_Empty(t *testing.T) {
	mapper := NewModelMapper(nil)
	if mapper == nil {
		t.Fatal("Expected non-nil mapper")
	}

	result := mapper.GetMappings()
	if len(result) != 0 {
		t.Errorf("Expected 0 mappings, got %d", len(result))
	}
}

func TestModelMapper_MapModel_NoProvider(t *testing.T) {
	mappings := []config.AmpModelMapping{
		{From: "claude-opus-4.5", To: "claude-sonnet-4"},
	}

	mapper := NewModelMapper(mappings)

	// Without a registered provider for the target, mapping should return empty
	result := mapper.MapModel("claude-opus-4.5")
	if result != "" {
		t.Errorf("Expected empty result when target has no provider, got %s", result)
	}
}

func TestModelMapper_MapModel_WithProvider(t *testing.T) {
	// Register a mock provider for the target model
	reg := registry.GetGlobalRegistry()
	reg.RegisterClient("test-client", "claude", []*registry.ModelInfo{
		{ID: "claude-sonnet-4", OwnedBy: "anthropic", Type: "claude"},
	})
	defer reg.UnregisterClient("test-client")

	mappings := []config.AmpModelMapping{
		{From: "claude-opus-4.5", To: "claude-sonnet-4"},
	}

	mapper := NewModelMapper(mappings)

	// With a registered provider, mapping should work
	result := mapper.MapModel("claude-opus-4.5")
	if result != "claude-sonnet-4" {
		t.Errorf("Expected claude-sonnet-4, got %s", result)
	}
}

func TestModelMapper_MapModel_TargetWithThinkingSuffix(t *testing.T) {
	reg := registry.GetGlobalRegistry()
	reg.RegisterClient("test-client-thinking", "codex", []*registry.ModelInfo{
		{ID: "gpt-5.2", OwnedBy: "openai", Type: "codex"},
	})
	defer reg.UnregisterClient("test-client-thinking")

	mappings := []config.AmpModelMapping{
		{From: "gpt-5.2-alias", To: "gpt-5.2(xhigh)"},
	}

	mapper := NewModelMapper(mappings)

	result := mapper.MapModel("gpt-5.2-alias")
	if result != "gpt-5.2(xhigh)" {
		t.Errorf("Expected gpt-5.2(xhigh), got %s", result)
	}
}

func TestModelMapper_MapModel_CaseInsensitive(t *testing.T) {
	reg := registry.GetGlobalRegistry()
	reg.RegisterClient("test-client2", "claude", []*registry.ModelInfo{
		{ID: "claude-sonnet-4", OwnedBy: "anthropic", Type: "claude"},
	})
	defer reg.UnregisterClient("test-client2")

	mappings := []config.AmpModelMapping{
		{From: "Claude-Opus-4.5", To: "claude-sonnet-4"},
	}

	mapper := NewModelMapper(mappings)

	// Should match case-insensitively
	result := mapper.MapModel("claude-opus-4.5")
	if result != "claude-sonnet-4" {
		t.Errorf("Expected claude-sonnet-4, got %s", result)
	}
}

func TestModelMapper_MapModel_NotFound(t *testing.T) {
	mappings := []config.AmpModelMapping{
		{From: "claude-opus-4.5", To: "claude-sonnet-4"},
	}

	mapper := NewModelMapper(mappings)

	// Unknown model should return empty
	result := mapper.MapModel("unknown-model")
	if result != "" {
		t.Errorf("Expected empty for unknown model, got %s", result)
	}
}

func TestModelMapper_MapModel_EmptyInput(t *testing.T) {
	mappings := []config.AmpModelMapping{
		{From: "claude-opus-4.5", To: "claude-sonnet-4"},
	}

	mapper := NewModelMapper(mappings)

	result := mapper.MapModel("")
	if result != "" {
		t.Errorf("Expected empty for empty input, got %s", result)
	}
}

func TestModelMapper_UpdateMappings(t *testing.T) {
	mapper := NewModelMapper(nil)

	// Initially empty
	if len(mapper.GetMappings()) != 0 {
		t.Error("Expected 0 initial mappings")
	}

	// Update with new mappings
	mapper.UpdateMappings([]config.AmpModelMapping{
		{From: "model-a", To: "model-b"},
		{From: "model-c", To: "model-d"},
	})

	result := mapper.GetMappings()
	if len(result) != 2 {
		t.Errorf("Expected 2 mappings after update, got %d", len(result))
	}

	// Update again should replace, not append
	mapper.UpdateMappings([]config.AmpModelMapping{
		{From: "model-x", To: "model-y"},
	})

	result = mapper.GetMappings()
	if len(result) != 1 {
		t.Errorf("Expected 1 mapping after second update, got %d", len(result))
	}
}

func TestModelMapper_UpdateMappings_SkipsInvalid(t *testing.T) {
	mapper := NewModelMapper(nil)

	mapper.UpdateMappings([]config.AmpModelMapping{
		{From: "", To: "model-b"},        // Invalid: empty from
		{From: "model-a", To: ""},        // Invalid: empty to
		{From: "  ", To: "model-b"},      // Invalid: whitespace from
		{From: "model-c", To: "model-d"}, // Valid
	})

	result := mapper.GetMappings()
	if len(result) != 1 {
		t.Errorf("Expected 1 valid mapping, got %d", len(result))
	}
}

func TestModelMapper_GetMappings_ReturnsCopy(t *testing.T) {
	mappings := []config.AmpModelMapping{
		{From: "model-a", To: "model-b"},
	}

	mapper := NewModelMapper(mappings)

	// Get mappings and modify the returned map
	result := mapper.GetMappings()
	result["new-key"] = "new-value"

	// Original should be unchanged
	original := mapper.GetMappings()
	if len(original) != 1 {
		t.Errorf("Expected original to have 1 mapping, got %d", len(original))
	}
	if _, exists := original["new-key"]; exists {
		t.Error("Original map was modified")
	}
}

func TestModelMapper_Regex_MatchBaseWithoutParens(t *testing.T) {
	reg := registry.GetGlobalRegistry()
	reg.RegisterClient("test-client-regex-1", "gemini", []*registry.ModelInfo{
		{ID: "gemini-2.5-pro", OwnedBy: "google", Type: "gemini"},
	})
	defer reg.UnregisterClient("test-client-regex-1")

	mappings := []config.AmpModelMapping{
		{From: "^gpt-5$", To: "gemini-2.5-pro", Regex: true},
	}

	mapper := NewModelMapper(mappings)

	// Incoming model has reasoning suffix, regex matches base, suffix is preserved
	result := mapper.MapModel("gpt-5(high)")
	if result != "gemini-2.5-pro(high)" {
		t.Errorf("Expected gemini-2.5-pro(high), got %s", result)
	}
}

func TestModelMapper_Regex_ExactPrecedence(t *testing.T) {
	reg := registry.GetGlobalRegistry()
	reg.RegisterClient("test-client-regex-2", "claude", []*registry.ModelInfo{
		{ID: "claude-sonnet-4", OwnedBy: "anthropic", Type: "claude"},
	})
	reg.RegisterClient("test-client-regex-3", "gemini", []*registry.ModelInfo{
		{ID: "gemini-2.5-pro", OwnedBy: "google", Type: "gemini"},
	})
	defer reg.UnregisterClient("test-client-regex-2")
	defer reg.UnregisterClient("test-client-regex-3")

	mappings := []config.AmpModelMapping{
		{From: "gpt-5", To: "claude-sonnet-4"},                 // exact
		{From: "^gpt-5.*$", To: "gemini-2.5-pro", Regex: true}, // regex
	}

	mapper := NewModelMapper(mappings)

	// Exact match should win over regex
	result := mapper.MapModel("gpt-5")
	if result != "claude-sonnet-4" {
		t.Errorf("Expected claude-sonnet-4, got %s", result)
	}
}

func TestModelMapper_Regex_InvalidPattern_Skipped(t *testing.T) {
	// Invalid regex should be skipped and not cause panic
	mappings := []config.AmpModelMapping{
		{From: "(", To: "target", Regex: true},
	}

	mapper := NewModelMapper(mappings)

	result := mapper.MapModel("anything")
	if result != "" {
		t.Errorf("Expected empty result due to invalid regex, got %s", result)
	}
}

func TestModelMapper_Regex_CaseInsensitive(t *testing.T) {
	reg := registry.GetGlobalRegistry()
	reg.RegisterClient("test-client-regex-4", "claude", []*registry.ModelInfo{
		{ID: "claude-sonnet-4", OwnedBy: "anthropic", Type: "claude"},
	})
	defer reg.UnregisterClient("test-client-regex-4")

	mappings := []config.AmpModelMapping{
		{From: "^CLAUDE-OPUS-.*$", To: "claude-sonnet-4", Regex: true},
	}

	mapper := NewModelMapper(mappings)

	result := mapper.MapModel("claude-opus-4.5")
	if result != "claude-sonnet-4" {
		t.Errorf("Expected claude-sonnet-4, got %s", result)
	}
}

func TestModelMapper_SuffixPreservation(t *testing.T) {
	reg := registry.GetGlobalRegistry()

	// Register test models
	reg.RegisterClient("test-client-suffix", "gemini", []*registry.ModelInfo{
		{ID: "gemini-2.5-pro", OwnedBy: "google", Type: "gemini"},
	})
	reg.RegisterClient("test-client-suffix-2", "claude", []*registry.ModelInfo{
		{ID: "claude-sonnet-4", OwnedBy: "anthropic", Type: "claude"},
	})
	defer reg.UnregisterClient("test-client-suffix")
	defer reg.UnregisterClient("test-client-suffix-2")

	tests := []struct {
		name     string
		mappings []config.AmpModelMapping
		input    string
		want     string
	}{
		{
			name:     "numeric suffix preserved",
			mappings: []config.AmpModelMapping{{From: "g25p", To: "gemini-2.5-pro"}},
			input:    "g25p(8192)",
			want:     "gemini-2.5-pro(8192)",
		},
		{
			name:     "level suffix preserved",
			mappings: []config.AmpModelMapping{{From: "g25p", To: "gemini-2.5-pro"}},
			input:    "g25p(high)",
			want:     "gemini-2.5-pro(high)",
		},
		{
			name:     "no suffix unchanged",
			mappings: []config.AmpModelMapping{{From: "g25p", To: "gemini-2.5-pro"}},
			input:    "g25p",
			want:     "gemini-2.5-pro",
		},
		{
			name:     "config suffix takes priority",
			mappings: []config.AmpModelMapping{{From: "alias", To: "gemini-2.5-pro(medium)"}},
			input:    "alias(high)",
			want:     "gemini-2.5-pro(medium)",
		},
		{
			name:     "regex with suffix preserved",
			mappings: []config.AmpModelMapping{{From: "^g25.*", To: "gemini-2.5-pro", Regex: true}},
			input:    "g25p(8192)",
			want:     "gemini-2.5-pro(8192)",
		},
		{
			name:     "auto suffix preserved",
			mappings: []config.AmpModelMapping{{From: "g25p", To: "gemini-2.5-pro"}},
			input:    "g25p(auto)",
			want:     "gemini-2.5-pro(auto)",
		},
		{
			name:     "none suffix preserved",
			mappings: []config.AmpModelMapping{{From: "g25p", To: "gemini-2.5-pro"}},
			input:    "g25p(none)",
			want:     "gemini-2.5-pro(none)",
		},
		{
			name:     "case insensitive base lookup with suffix",
			mappings: []config.AmpModelMapping{{From: "G25P", To: "gemini-2.5-pro"}},
			input:    "g25p(high)",
			want:     "gemini-2.5-pro(high)",
		},
		{
			name:     "empty suffix filtered out",
			mappings: []config.AmpModelMapping{{From: "g25p", To: "gemini-2.5-pro"}},
			input:    "g25p()",
			want:     "gemini-2.5-pro",
		},
		{
			name:     "incomplete suffix treated as no suffix",
			mappings: []config.AmpModelMapping{{From: "g25p(high", To: "gemini-2.5-pro"}},
			input:    "g25p(high",
			want:     "gemini-2.5-pro",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mapper := NewModelMapper(tt.mappings)
			got := mapper.MapModel(tt.input)
			if got != tt.want {
				t.Errorf("MapModel(%q) = %q, want %q", tt.input, got, tt.want)
			}
		})
	}
}