package modelcatalog

import (
	"testing"

	"github.com/maximhq/bifrost/core/schemas"
	"github.com/maximhq/bifrost/framework/configstore"
	configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables"
	"github.com/stretchr/testify/assert"
)

// newTestCatalog creates a minimal ModelCatalog for testing within the package.
func newTestCatalog(modelPool map[schemas.ModelProvider][]string, baseModelIndex map[string]string) *ModelCatalog {
	if modelPool == nil {
		modelPool = make(map[schemas.ModelProvider][]string)
	}
	if baseModelIndex == nil {
		baseModelIndex = make(map[string]string)
	}
	return &ModelCatalog{
		modelPool:      modelPool,
		baseModelIndex: baseModelIndex,
		pricingData:    make(map[string]configstoreTables.TableModelPricing),
	}
}

// --- GetBaseModelName tests ---

func TestGetBaseModelName_Simple(t *testing.T) {
	mc := newTestCatalog(nil, nil)
	// No catalog data, no prefix — returns as-is (no date suffix to strip either)
	assert.Equal(t, "gpt-4o", mc.GetBaseModelName("gpt-4o"))
}

func TestGetBaseModelName_Prefixed(t *testing.T) {
	mc := newTestCatalog(nil, nil)
	// Provider prefix stripped, no catalog — algorithmic fallback returns base
	assert.Equal(t, "gpt-4o", mc.GetBaseModelName("openai/gpt-4o"))
}

func TestGetBaseModelName_PrefixedAnthropic(t *testing.T) {
	mc := newTestCatalog(nil, nil)
	assert.Equal(t, "claude-3-5-sonnet", mc.GetBaseModelName("anthropic/claude-3-5-sonnet"))
}

func TestGetBaseModelName_FromCatalog(t *testing.T) {
	// Model has a pre-computed base_model in the catalog
	mc := newTestCatalog(nil, map[string]string{
		"gpt-4o":            "gpt-4o",
		"gpt-4o-2024-08-06": "gpt-4o",
	})
	assert.Equal(t, "gpt-4o", mc.GetBaseModelName("gpt-4o"))
	assert.Equal(t, "gpt-4o", mc.GetBaseModelName("gpt-4o-2024-08-06"))
}

func TestGetBaseModelName_ProviderPrefixWithCatalog(t *testing.T) {
	// Model has provider prefix — strip prefix, then find in catalog
	mc := newTestCatalog(nil, map[string]string{
		"gpt-4o": "gpt-4o",
	})
	assert.Equal(t, "gpt-4o", mc.GetBaseModelName("openai/gpt-4o"))
}

func TestGetBaseModelName_FallbackAlgorithmic(t *testing.T) {
	// Model NOT in catalog — falls back to schemas.BaseModelName (date stripping)
	mc := newTestCatalog(nil, nil)
	// Anthropic-style date suffix
	assert.Equal(t, "claude-sonnet-4", mc.GetBaseModelName("claude-sonnet-4-20250514"))
	// OpenAI-style date suffix
	assert.Equal(t, "gpt-4o", mc.GetBaseModelName("gpt-4o-2024-08-06"))
}

func TestGetBaseModelName_FallbackAlgorithmicWithPrefix(t *testing.T) {
	// Provider prefix + not in catalog — strip prefix, then algorithmic fallback
	mc := newTestCatalog(nil, nil)
	assert.Equal(t, "claude-sonnet-4", mc.GetBaseModelName("anthropic/claude-sonnet-4-20250514"))
}

func TestGetBaseModelName_UnknownModel(t *testing.T) {
	mc := newTestCatalog(nil, nil)
	assert.Equal(t, "some-random-model", mc.GetBaseModelName("some-random-model"))
}

func TestGetBaseModelName_CatalogTakesPrecedence(t *testing.T) {
	// If catalog says the base_model is X, use it even if algorithmic would give Y
	mc := newTestCatalog(nil, map[string]string{
		"my-custom-model-20250101": "my-custom-model-20250101", // catalog says keep the date
	})
	assert.Equal(t, "my-custom-model-20250101", mc.GetBaseModelName("my-custom-model-20250101"))
}

// --- IsSameModel tests ---

func TestIsSameModel_DirectMatch(t *testing.T) {
	mc := newTestCatalog(nil, nil)
	assert.True(t, mc.IsSameModel("gpt-4o", "gpt-4o"))
}

func TestIsSameModel_ProviderPrefix(t *testing.T) {
	mc := newTestCatalog(nil, nil)
	assert.True(t, mc.IsSameModel("openai/gpt-4o", "gpt-4o"))
	assert.True(t, mc.IsSameModel("gpt-4o", "openai/gpt-4o"))
}

func TestIsSameModel_BothPrefixed(t *testing.T) {
	mc := newTestCatalog(nil, nil)
	assert.True(t, mc.IsSameModel("openai/gpt-4o", "openai/gpt-4o"))
}

func TestIsSameModel_DifferentProvidersSameBase(t *testing.T) {
	mc := newTestCatalog(nil, nil)
	// Both have the same base model after stripping different provider prefixes
	assert.True(t, mc.IsSameModel("openai/gpt-4o", "azure/gpt-4o"))
}

func TestIsSameModel_DifferentModels(t *testing.T) {
	mc := newTestCatalog(nil, nil)
	assert.False(t, mc.IsSameModel("gpt-4o", "claude-3-5-sonnet"))
}

func TestIsSameModel_DifferentModelsBothPrefixed(t *testing.T) {
	mc := newTestCatalog(nil, nil)
	assert.False(t, mc.IsSameModel("openai/gpt-4o", "anthropic/claude-3-5-sonnet"))
}

func TestIsSameModel_CatalogBacked(t *testing.T) {
	// Two model strings that look different but the catalog says they have the same base_model
	mc := newTestCatalog(nil, map[string]string{
		"claude-3-5-sonnet":          "claude-3-5-sonnet",
		"claude-3-5-sonnet-20241022": "claude-3-5-sonnet",
	})
	assert.True(t, mc.IsSameModel("claude-3-5-sonnet", "claude-3-5-sonnet-20241022"))
	assert.True(t, mc.IsSameModel("claude-3-5-sonnet-20241022", "claude-3-5-sonnet"))
}

func TestIsSameModel_AlgorithmicFallback(t *testing.T) {
	// Models not in catalog — use algorithmic date stripping
	mc := newTestCatalog(nil, nil)
	assert.True(t, mc.IsSameModel("custom-model-20250101", "custom-model"))
}

func TestIsSameModel_EmptyStrings(t *testing.T) {
	mc := newTestCatalog(nil, nil)
	assert.True(t, mc.IsSameModel("", ""))
	assert.False(t, mc.IsSameModel("gpt-4o", ""))
	assert.False(t, mc.IsSameModel("", "gpt-4o"))
}

func TestIsModelAllowedForProvider_PrefixedAllowedModelInCatalog(t *testing.T) {
	mc := newTestCatalog(
		map[schemas.ModelProvider][]string{
			schemas.OpenRouter: {"openai/gpt-4o"},
		},
		nil,
	)

	providerConfig := configstore.ProviderConfig{}

	assert.True(t, mc.IsModelAllowedForProvider(schemas.OpenRouter, "gpt-4o", &providerConfig, []string{"openai/gpt-4o"}))
}

func TestIsModelAllowedForProvider_CustomProviderListModelsDisabled(t *testing.T) {
	mc := newTestCatalog(nil, nil)

	// Custom provider with list-models disabled + ["*"] → should return true
	providerConfig := configstore.ProviderConfig{
		CustomProviderConfig: &schemas.CustomProviderConfig{
			AllowedRequests: &schemas.AllowedRequests{
				ListModels: false,
			},
		},
	}
	assert.True(t, mc.IsModelAllowedForProvider("custom-provider", "any-model", &providerConfig, []string{"*"}))
}

func TestIsModelAllowedForProvider_CustomProviderListModelsEnabled(t *testing.T) {
	mc := newTestCatalog(
		map[schemas.ModelProvider][]string{
			"custom-provider": {"model-a"},
		},
		nil,
	)

	// Custom provider with list-models enabled + ["*"] → should go through catalog
	providerConfig := configstore.ProviderConfig{
		CustomProviderConfig: &schemas.CustomProviderConfig{
			AllowedRequests: &schemas.AllowedRequests{
				ListModels: true,
			},
		},
	}
	// model-a is in catalog → allowed
	assert.True(t, mc.IsModelAllowedForProvider("custom-provider", "model-a", &providerConfig, []string{"*"}))
	// model-b is NOT in catalog → denied
	assert.False(t, mc.IsModelAllowedForProvider("custom-provider", "model-b", &providerConfig, []string{"*"}))
}

func TestIsModelAllowedForProvider_NilProviderConfig(t *testing.T) {
	mc := newTestCatalog(
		map[schemas.ModelProvider][]string{
			"some-provider": {"model-x"},
		},
		nil,
	)

	// nil providerConfig + ["*"] → should go through catalog (not bypass)
	assert.True(t, mc.IsModelAllowedForProvider("some-provider", "model-x", nil, []string{"*"}))
	assert.False(t, mc.IsModelAllowedForProvider("some-provider", "model-y", nil, []string{"*"}))
}