package controller

import (
	"context"
	"errors"
	hostdto "github.com/QuantumNous/new-api/dto"
	"github.com/QuantumNous/new-api/model"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relaykit/dto"
	"github.com/QuantumNous/new-api/relaykit/types"
	"github.com/QuantumNous/new-api/service"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	"net/http/httptest"
	"testing"
)

func TestRelayAttemptSwitchSafety(t *testing.T) {
	for _, tc := range []struct {
		name, reason, previous, conversation string
		count                                int
		pin, cancel, written, disabled, want bool
	}{
		{name: "explicit mismatch", reason: "model_mismatch", count: 1, want: true},
		{name: "token disabled", reason: "model_mismatch", count: 1, disabled: true},
		{name: "third channel allowed", reason: "model_mismatch", count: 2, want: true},
		{name: "third channel final", reason: "model_mismatch", count: 3},
		{name: "missing model", reason: "missing_model", count: 1},
		{name: "confirmation timeout", reason: "confirmation_timeout", count: 1},
		{name: "pinned", reason: "model_mismatch", count: 1, pin: true},
		{name: "cancelled", reason: "model_mismatch", count: 1, cancel: true},
		{name: "already delivered", reason: "model_mismatch", count: 1, written: true},
		{name: "upstream response id", reason: "model_mismatch", count: 1, previous: "resp_private"},
		{name: "upstream conversation", reason: "model_mismatch", count: 1, conversation: `"conv_private"`},
	} {
		t.Run(tc.name, func(t *testing.T) {
			c, _ := gin.CreateTestContext(httptest.NewRecorder())
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			c.Request = httptest.NewRequest("POST", "/v1/responses", nil).WithContext(ctx)
			c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "a"}}})
			if tc.pin {
				service.GetChannelConstraints(c).AddPin(hostdto.ChannelPin{ChannelId: 1, Source: hostdto.PinSourceToken, Rank: hostdto.PinRankToken, RetryMode: hostdto.PinRetrySingleAttempt})
			}
			if tc.cancel {
				cancel()
			}
			if tc.written {
				_, _ = c.Writer.WriteString("already sent")
			}
			info := &relaycommon.RelayInfo{DurableAttemptBilling: true, UpstreamRequests: tc.count, Request: &dto.OpenAIResponsesRequest{PreviousResponseID: tc.previous, Conversation: []byte(tc.conversation)}, ResponseGuard: &relaycommon.ResponseModelGuard{Attempts: []*relaycommon.ResponseGuardAttempt{{Reason: tc.reason}}}}
			err := types.NewError(errors.New("guard rejected"), relaycommon.ResponseGuardErrorCode)
			info.AstraAutoSwitch = !tc.disabled
			require.Equal(t, tc.want, canSwitchRejectedAstra(c, info, err))
		})
	}
}

func TestRelayAttemptGuardResult(t *testing.T) {
	require.Equal(t, "mismatch_rejected", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "model_mismatch", Number: 1}))
	require.Equal(t, "switched_success", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "accepted", Number: 2}))
	require.Equal(t, "accepted", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "accepted", Number: 1}))
	require.Equal(t, "accepted", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "unchecked", Number: 1, RequestedModel: "gpt-6-astra", ReturnedModel: "gpt-6-astra"}))
	require.Equal(t, "unchecked", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "unchecked", Number: 1, RequestedModel: "gpt-6-astra"}))
	require.Equal(t, "transport_error", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "transport_error", Number: 1}))
}