package helper

import (
	"net/http/httptest"
	"testing"

	"github.com/QuantumNous/new-api/pkg/billingexpr"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relaykit/types"
	"github.com/QuantumNous/new-api/setting/config"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
)

func TestGroupModelPricingPreservesOtherGroupsAndSnapshot(t *testing.T) {
	saved := map[string]string{}
	require.NoError(t, config.GlobalConfig.SaveToDB(func(k, v string) error { saved[k] = v; return nil }))
	t.Cleanup(func() { require.NoError(t, config.GlobalConfig.LoadFromDB(saved)) })
	require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
		"billing_setting.billing_mode":       `{"gpt-5.6-sol":"tiered_expr"}`,
		"billing_setting.billing_expr":       `{"gpt-5.6-sol":"tier(\"standard\", p * 4 + c * 20 + cr * 0.4)"}`,
		"billing_setting.group_billing_expr": `{"gpt-v-pro":{"gpt-5.6-sol":"tier(\"standard\", p * 5 + c * 30 + cr * 0.5)"}}`,
		"group_ratio_setting.group_ratio":    `{"gpt-v-pro":0.2052,"gpt-v-mix":0.18}`,
	}))
	for _, tc := range []struct {
		group    string
		expected int
		cost     float64
	}{
		{"gpt-v-pro", 513, 8500},
		{"gpt-v-mix", 360, 6400},
	} {
		t.Run(tc.group, func(t *testing.T) {
			ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
			ctx.Request = httptest.NewRequest("POST", "/v1/responses", nil)
			ctx.Set("auto_group", tc.group)
			info := &relaycommon.RelayInfo{
				OriginModelName: "gpt-5.6-sol", UserGroup: "default", UsingGroup: "auto",
				BillingRequestInput: &billingexpr.RequestInput{Body: []byte(`{}`)},
			}
			data, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{MaxTokens: 100})
			require.NoError(t, err)
			require.Equal(t, tc.expected, data.QuotaToPreConsume)
			require.Equal(t, tc.group, info.UsingGroup)
			raw, _, err := billingexpr.RunExprWithRequest(info.TieredBillingSnapshot.ExprString,
				billingexpr.TokenParams{P: 1000, C: 100, CR: 1000, Len: 2000}, billingexpr.RequestInput{})
			require.NoError(t, err)
			require.Equal(t, tc.cost, raw)
		})
	}
}