package llm

import (
	"math"
	"testing"
)

func TestModelPricing_Cost(t *testing.T) {
	pricing := ModelPricing{
		InputPricePerToken:  0.000003, // $3/MTok
		OutputPricePerToken: 0.000015, // $15/MTok
	}

	usage := &Usage{
		InputTokens:  1000,
		OutputTokens: 500,
		TotalTokens:  1500,
	}

	cost := pricing.Cost(usage)
	expected := 1000*0.000003 + 500*0.000015
	if math.Abs(cost-expected) > 1e-10 {
		t.Errorf("expected cost %f, got %f", expected, cost)
	}
}

func TestModelPricing_Cost_WithCachedTokens(t *testing.T) {
	pricing := ModelPricing{
		InputPricePerToken:       0.000003, // $3/MTok
		OutputPricePerToken:      0.000015, // $15/MTok
		CachedInputPricePerToken: 0.0000015, // $1.50/MTok (50% discount)
	}

	usage := &Usage{
		InputTokens:  1000,
		OutputTokens: 500,
		TotalTokens:  1500,
		Details: map[string]int{
			UsageDetailCachedInputTokens: 400,
		},
	}

	cost := pricing.Cost(usage)
	// 600 regular input tokens + 400 cached tokens + 500 output tokens
	expected := 600*0.000003 + 400*0.0000015 + 500*0.000015
	if math.Abs(cost-expected) > 1e-10 {
		t.Errorf("expected cost %f, got %f", expected, cost)
	}
}

func TestModelPricing_Cost_NilUsage(t *testing.T) {
	pricing := ModelPricing{
		InputPricePerToken:  0.000003,
		OutputPricePerToken: 0.000015,
	}

	cost := pricing.Cost(nil)
	if cost != 0 {
		t.Errorf("expected 0 for nil usage, got %f", cost)
	}
}

func TestModelPricing_Cost_NoCachedPrice(t *testing.T) {
	// When CachedInputPricePerToken is 0, all input tokens use InputPricePerToken
	pricing := ModelPricing{
		InputPricePerToken:  0.000003,
		OutputPricePerToken: 0.000015,
	}

	usage := &Usage{
		InputTokens:  1000,
		OutputTokens: 500,
		TotalTokens:  1500,
		Details: map[string]int{
			UsageDetailCachedInputTokens: 400,
		},
	}

	cost := pricing.Cost(usage)
	expected := 1000*0.000003 + 500*0.000015
	if math.Abs(cost-expected) > 1e-10 {
		t.Errorf("expected cost %f, got %f", expected, cost)
	}
}

func TestPricingRegistry(t *testing.T) {
	registry := NewPricingRegistry()

	registry.Set("gpt-4o", ModelPricing{
		InputPricePerToken:  0.0000025,
		OutputPricePerToken: 0.00001,
	})

	if !registry.Has("gpt-4o") {
		t.Error("expected Has('gpt-4o') to be true")
	}
	if registry.Has("gpt-3.5-turbo") {
		t.Error("expected Has('gpt-3.5-turbo') to be false")
	}

	usage := &Usage{InputTokens: 1000, OutputTokens: 200, TotalTokens: 1200}

	cost := registry.Cost("gpt-4o", usage)
	expected := 1000*0.0000025 + 200*0.00001
	if math.Abs(cost-expected) > 1e-10 {
		t.Errorf("expected cost %f, got %f", expected, cost)
	}

	// Unknown model returns 0
	cost = registry.Cost("unknown-model", usage)
	if cost != 0 {
		t.Errorf("expected 0 for unknown model, got %f", cost)
	}
}

func TestPricingRegistry_Override(t *testing.T) {
	registry := NewPricingRegistry()

	registry.Set("model-a", ModelPricing{InputPricePerToken: 0.001, OutputPricePerToken: 0.002})
	registry.Set("model-a", ModelPricing{InputPricePerToken: 0.003, OutputPricePerToken: 0.004})

	usage := &Usage{InputTokens: 100, OutputTokens: 50, TotalTokens: 150}
	cost := registry.Cost("model-a", usage)
	expected := 100*0.003 + 50*0.004
	if math.Abs(cost-expected) > 1e-10 {
		t.Errorf("expected overridden cost %f, got %f", expected, cost)
	}
}