package gateway

import (
	"context"
	"io"
	"math"
	"testing"

	"github.com/enterpilot/gomodel/internal/core"
	"github.com/enterpilot/gomodel/internal/usage"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type usageCaptureLogger struct {
	config  usage.Config
	entries []*usage.UsageEntry
}

func (l *usageCaptureLogger) Write(entry *usage.UsageEntry) {
	l.entries = append(l.entries, entry)
}

func (l *usageCaptureLogger) Config() usage.Config { return l.config }
func (l *usageCaptureLogger) Close() error         { return nil }

type pricingCaptureResolver struct {
	model    string
	provider string
	pricing  *core.ModelPricing
}

func (r *pricingCaptureResolver) ResolvePricing(model, provider string) *core.ModelPricing {
	r.model = model
	r.provider = provider
	return r.pricing
}

func TestInferenceOrchestratorLogUsageAssignsUserPathAndProviderName(t *testing.T) {
	logger := &usageCaptureLogger{config: usage.Config{Enabled: true}}
	orchestrator := NewInferenceOrchestrator(InferenceConfig{UsageLogger: logger})
	ctx := core.WithRequestSnapshot(context.Background(), &core.RequestSnapshot{UserPath: "/team/alpha"})
	ctx = core.WithAuthKeyID(ctx, "key-a")

	orchestrator.LogUsage(ctx, nil, "gpt-5-nano", "openai", "primary-openai", func(*core.ModelPricing) *usage.UsageEntry {
		return &usage.UsageEntry{ID: "usage-1"}
	})

	require.Len(t, logger.entries, 1)
	got := logger.entries[0].UserPath
	require.Equal(t, "/team/alpha", got)
	got = logger.entries[0].ProviderName
	require.Equal(t, "primary-openai", got)
	assert.Equal(t, "key-a", logger.entries[0].AuthKeyID)
}

func TestExecuteChatCompletionPricesRequestedModelWhenResponseModelIsVersioned(t *testing.T) {
	zero := 0.0
	perRequest := 0.033333
	logger := &usageCaptureLogger{config: usage.Config{Enabled: true}}
	pricing := &pricingCaptureResolver{pricing: &core.ModelPricing{
		InputPerMtok:  &zero,
		OutputPerMtok: &zero,
		PerRequest:    &perRequest,
	}}
	provider := &providerTypeResolverStub{
		chatResponse: &core.ChatResponse{
			ID:       "chatcmpl-test",
			Model:    "gpt-4o-mini-2024-07-18",
			Provider: "openai",
			Choices:  []core.Choice{{FinishReason: "stop"}},
			Usage: core.Usage{
				PromptTokens:     12,
				CompletionTokens: 1,
				TotalTokens:      13,
			},
		},
	}
	orchestrator := NewInferenceOrchestrator(InferenceConfig{
		Provider:        provider,
		UsageLogger:     logger,
		PricingResolver: pricing,
	})
	workflow := &core.Workflow{
		ProviderType: "openai",
		Resolution: &core.RequestModelResolution{
			Requested:        core.NewRequestedModelSelector("openai/gpt-4o-mini", ""),
			ResolvedSelector: core.ModelSelector{Provider: "openai", Model: "gpt-4o-mini"},
			ProviderType:     "openai",
			ProviderName:     "openai",
		},
	}

	_, err := orchestrator.ExecuteChatCompletion(
		context.Background(),
		workflow,
		&core.ChatRequest{Model: "openai/gpt-4o-mini"},
		"req-usage-pricing",
		"/v1/chat/completions",
	)
	require.NoError(t, err)
	require.Equal(t, "gpt-4o-mini", pricing.model)
	require.Equal(t, "openai", pricing.provider)
	require.Len(t, logger.entries, 1)

	entry := logger.entries[0]
	require.Equal(t, "gpt-4o-mini-2024-07-18", entry.Model)
	assert.Equal(t, "gpt-4o-mini", entry.RoutedModel, "markups select on the routed model")
	require.NotNil(t, entry.TotalCost)
	require.LessOrEqual(t, math.Abs(*entry.TotalCost-perRequest), 0.0000001, "total cost = %v, want %f", entry.TotalCost, perRequest)
}

// The response's provider field names the configured instance that served the
// request, while the execution metadata keeps the provider type that drives
// routing and audit filters.
func TestExecuteChatCompletionNamesTheProviderInstance(t *testing.T) {
	tests := []struct {
		name         string
		providerName string
		wantProvider string
	}{
		{name: "configured instance name", providerName: "mockds", wantProvider: "mockds"},
		{name: "no configured instance name", wantProvider: "deepseek"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := &providerTypeResolverStub{
				chatResponse: &core.ChatResponse{
					ID:       "chatcmpl-test",
					Model:    "deepseek-chat",
					Provider: "deepseek",
					Choices:  []core.Choice{{FinishReason: "stop"}},
				},
			}
			orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: provider})
			workflow := &core.Workflow{
				ProviderType: "deepseek",
				Resolution: &core.RequestModelResolution{
					Requested:        core.NewRequestedModelSelector("deepseek/deepseek-chat", ""),
					ResolvedSelector: core.ModelSelector{Provider: "deepseek", Model: "deepseek-chat"},
					ProviderType:     "deepseek",
					ProviderName:     tt.providerName,
				},
			}

			result, err := orchestrator.ExecuteChatCompletion(
				context.Background(),
				workflow,
				&core.ChatRequest{Model: "deepseek/deepseek-chat"},
				"req-provider-name",
				"/v1/chat/completions",
			)
			require.NoError(t, err)
			require.NotNil(t, result.Response)
			assert.Equal(t, tt.wantProvider, result.Response.Provider, "response provider")
			assert.Equal(t, "deepseek", result.Meta.ProviderType, "meta provider type must stay the type")
			assert.Equal(t, tt.providerName, result.Meta.ProviderName, "meta provider name")
		})
	}
}

func TestExecuteResponsesNamesTheProviderInstance(t *testing.T) {
	tests := []struct {
		name         string
		providerName string
		wantProvider string
	}{
		{name: "configured instance name", providerName: "mockds", wantProvider: "mockds"},
		{name: "no configured instance name", wantProvider: "deepseek"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := &providerTypeResolverStub{
				responsesResponse: &core.ResponsesResponse{
					ID:       "resp-test",
					Object:   "response",
					Model:    "deepseek-chat",
					Provider: "deepseek",
					Status:   "completed",
				},
			}
			orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: provider})
			workflow := &core.Workflow{
				ProviderType: "deepseek",
				Resolution: &core.RequestModelResolution{
					Requested:        core.NewRequestedModelSelector("deepseek/deepseek-chat", ""),
					ResolvedSelector: core.ModelSelector{Provider: "deepseek", Model: "deepseek-chat"},
					ProviderType:     "deepseek",
					ProviderName:     tt.providerName,
				},
			}

			result, err := orchestrator.ExecuteResponses(
				context.Background(),
				workflow,
				&core.ResponsesRequest{Model: "deepseek/deepseek-chat"},
				"req-provider-name",
				"/v1/responses",
			)
			require.NoError(t, err)
			require.NotNil(t, result.Response)
			assert.Equal(t, tt.wantProvider, result.Response.Provider, "response provider")
			assert.Equal(t, "deepseek", result.Meta.ProviderType, "meta provider type must stay the type")
			assert.Equal(t, tt.providerName, result.Meta.ProviderName, "meta provider name")
		})
	}
}

func TestInferenceOrchestratorLogUsageSkipsWhenWorkflowDisablesUsage(t *testing.T) {
	logger := &usageCaptureLogger{config: usage.Config{Enabled: true}}
	orchestrator := NewInferenceOrchestrator(InferenceConfig{UsageLogger: logger})

	orchestrator.LogUsage(context.Background(), &core.Workflow{
		Policy: &core.ResolvedWorkflowPolicy{
			VersionID: "workflow-usage-off",
			Features: core.WorkflowFeatures{
				Cache:      true,
				Audit:      true,
				Usage:      false,
				Guardrails: true,
			},
		},
	}, "gpt-5-nano", "openai", "primary-openai", func(*core.ModelPricing) *usage.UsageEntry {
		return &usage.UsageEntry{ID: "usage-1"}
	})

	require.Empty(t, logger.entries)
}

func TestInferenceOrchestratorWithCacheRequestContextClearsInheritedGuardrailsHash(t *testing.T) {
	orchestrator := NewInferenceOrchestrator(InferenceConfig{GuardrailsHash: "service-default"})
	ctx := core.WithGuardrailsHash(context.Background(), "caller-hash")
	workflow := &core.Workflow{
		Policy: &core.ResolvedWorkflowPolicy{
			VersionID:      "workflow-1",
			GuardrailsHash: "",
			Features: core.WorkflowFeatures{
				Cache:      true,
				Audit:      true,
				Usage:      true,
				Guardrails: false,
				Failover:   true,
			},
		},
	}

	got := orchestrator.WithCacheRequestContext(ctx, workflow)
	hash := core.GetGuardrailsHash(got)
	require.Empty(t, hash)
}

func TestInferenceOrchestratorProviderTypeForSelectorPrefersExplicitProvider(t *testing.T) {
	orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: &providerTypeResolverStub{}})

	got := orchestrator.ProviderTypeForSelector(core.ModelSelector{Provider: "azure", Model: "gpt-4o"}, "openai")
	require.Equal(t, "azure", got)
}

func TestInferenceOrchestratorProviderTypeForSelectorCanonicalizesProviderNameSelectors(t *testing.T) {
	orchestrator := NewInferenceOrchestrator(InferenceConfig{
		Provider: &providerTypeResolverStub{
			providerTypes: map[string]string{
				"openai_test/gpt-4o": "openai",
			},
		},
	})

	got := orchestrator.ProviderTypeForSelector(core.ModelSelector{Provider: "openai_test", Model: "gpt-4o"}, "anthropic")
	require.Equal(t, "openai", got)
}

// namedProviderStub knows configured provider names and their types, but no
// catalog entry for the models under test.
type namedProviderStub struct {
	providerTypeResolverStub
	typesByName map[string]string
}

func (p *namedProviderStub) GetProviderTypeForName(name string) string { return p.typesByName[name] }

// A failover target the catalog does not list, such as a pinned jev version
// on a provider named kev, is routed to the provider its selector names, with
// that provider's type, never to the primary's provider.
func TestUnlistedSelectorResolvesToTheProviderItNames(t *testing.T) {
	provider := &namedProviderStub{typesByName: map[string]string{"kev": "jev", "jev-down": "jev"}}
	orchestrator := NewInferenceOrchestrator(InferenceConfig{Provider: provider})

	pinned := core.ModelSelector{Provider: "kev", Model: "kev-4b-2026-09"}
	assert.Equal(t, "jev", orchestrator.ProviderTypeForSelector(pinned, "openai"))
	assert.Equal(t, "kev", ResolvedProviderName(provider, pinned, "jev-down"))

	unknown := core.ModelSelector{Provider: "nope", Model: "x"}
	assert.Equal(t, "jev-down", ResolvedProviderName(provider, unknown, "jev-down"))
}

func TestQualifyModelWithProviderPrefixesSlashModelIDs(t *testing.T) {
	got := QualifyModelWithProvider("openai/gpt-4o-mini", "openrouter")
	require.Equal(t, "openrouter/openai/gpt-4o-mini", got)
}

func TestQualifyModelWithProviderKeepsAlreadyQualifiedModelIDs(t *testing.T) {
	got := QualifyModelWithProvider("openrouter/openai/gpt-4o-mini", "openrouter")
	require.Equal(t, "openrouter/openai/gpt-4o-mini", got)
}

type providerTypeResolverStub struct {
	providerTypes     map[string]string
	chatResponse      *core.ChatResponse
	responsesResponse *core.ResponsesResponse
}

func (p *providerTypeResolverStub) ChatCompletion(context.Context, *core.ChatRequest) (*core.ChatResponse, error) {
	return p.chatResponse, nil
}

func (p *providerTypeResolverStub) StreamChatCompletion(context.Context, *core.ChatRequest) (io.ReadCloser, error) {
	return nil, nil
}

func (p *providerTypeResolverStub) ListModels(context.Context) (*core.ModelsResponse, error) {
	return nil, nil
}

func (p *providerTypeResolverStub) Responses(context.Context, *core.ResponsesRequest) (*core.ResponsesResponse, error) {
	return p.responsesResponse, nil
}

func (p *providerTypeResolverStub) StreamResponses(context.Context, *core.ResponsesRequest) (io.ReadCloser, error) {
	return nil, nil
}

func (p *providerTypeResolverStub) Embeddings(context.Context, *core.EmbeddingRequest) (*core.EmbeddingResponse, error) {
	return nil, nil
}

func (p *providerTypeResolverStub) Supports(string) bool { return true }

func (p *providerTypeResolverStub) GetProviderType(model string) string {
	return p.providerTypes[model]
}