package controller

import (
	"errors"
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/dto"
	"github.com/QuantumNous/new-api/relaykit/types"
	"github.com/QuantumNous/new-api/service"
	"github.com/stretchr/testify/require"
)

func TestFixedBackupRetryOverridesAffinityOnly(t *testing.T) {
	configureTokenAutoGroupsTest(t, "2", "[]")
	for _, tc := range []struct {
		name string
		group string
		backups []string
		cross bool
		status int
		skip bool
		pinned bool
		want bool
	}{
		{"503 with backup", "default", []string{"vip"}, true, 503, false, false, true},
		{"no backup", "default", nil, true, 503, false, false, false},
		{"same group", "default", []string{"default"}, true, 503, false, false, false},
		{"revoked backup", "default", []string{"unavailable"}, true, 503, false, false, false},
		{"retry disabled", "default", []string{"vip"}, false, 503, false, false, false},
		{"auto unchanged", "auto", []string{"vip"}, true, 503, false, false, false},
		{"invalid request", "default", []string{"vip"}, true, 400, true, false, false},
		{"skip retry", "default", []string{"vip"}, true, 503, true, false, false},
		{"explicit channel pin", "default", []string{"vip"}, true, 503, false, true, false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			c := newPinRetryContext()
			c.Set("channel_affinity_skip_retry_on_failure", true)
			common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
			common.SetContextKey(c, constant.ContextKeyTokenGroup, tc.group)
			common.SetContextKey(c, constant.ContextKeyTokenAutoGroups, tc.backups)
			common.SetContextKey(c, constant.ContextKeyTokenCrossGroupRetry, tc.cross)
			if tc.pinned {
				service.GetChannelConstraints(c).AddPin(dto.ChannelPin{
					ChannelId: 13, Source: dto.PinSourceToken,
					Rank: dto.PinRankToken, RetryMode: dto.PinRetrySingleAttempt,
				})
			}
			err := types.NewOpenAIError(errors.New("upstream unavailable"), types.ErrorCodeBadResponseStatusCode, tc.status)
			if tc.skip {
				err = types.NewOpenAIError(errors.New("do not replay"), types.ErrorCodeBadResponseStatusCode, tc.status, types.ErrOptionWithSkipRetry())
			}
			require.Equal(t, tc.want, shouldRetry(c, err, 1))
			require.False(t, shouldRetry(c, err, 0), "exhausted retry budget stays bounded")
		})
	}
}