package perfmetrics

import (
	"context"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/common"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relaykit/types"
	"github.com/alicebob/miniredis/v2"
	"github.com/go-redis/redis/v8"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestLiveUserResults(t *testing.T) {
	oldRedis := common.RedisEnabled
	common.RedisEnabled = false
	t.Cleanup(func() { common.RedisEnabled = oldRedis })
	canceled, cancel := context.WithCancel(context.Background())
	cancel()
	for _, tc := range []struct {
		name    string
		test    bool
		token   int
		ctx     context.Context
		err     *types.NewAPIError
		want    bool
		success bool
	}{
		{"success", false, 1, context.Background(), nil, true, true},
		{"upstream 503", false, 1, context.Background(), types.InitOpenAIError("server_error", 503), true, false},
		{"probe", true, 1, context.Background(), types.InitOpenAIError("server_error", 503), false, false},
		{"internal", false, 0, context.Background(), nil, false, false},
		{"canceled", false, 1, canceled, types.InitOpenAIError("server_error", 503), false, false},
		{"bad prompt", false, 1, context.Background(), types.InitOpenAIError("context_length_exceeded", 400), false, false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Cleanup(func() { liveResults.Delete(tc.name) })
			RecordRelayResult(tc.ctx, &relaycommon.RelayInfo{UsingGroup: tc.name, TokenId: tc.token, IsChannelTest: tc.test, IsStream: true, StartTime: time.Now()}, tc.err)
			result := LatestUserResults(context.Background(), []string{tc.name})[tc.name]
			assert.Equal(t, tc.want, result.At > 0)
			if tc.want {
				assert.Equal(t, tc.success, result.Success)
			}
		})
	}
}

func TestUserResultsAreOrderedAndSurviveRestart(t *testing.T) {
	server := miniredis.RunT(t)
	client := redis.NewClient(&redis.Options{Addr: server.Addr()})
	oldRedis, oldClient := common.RedisEnabled, common.RDB
	common.RedisEnabled, common.RDB = true, client
	t.Cleanup(func() {
		common.RedisEnabled, common.RDB = oldRedis, oldClient
		liveResults.Delete("group")
		liveResults.Delete("other")
		client.Close()
	})
	ctx := context.Background()
	require.NoError(t, client.Set(ctx, "status:user-result:group", "9999999999999:false", time.Minute).Err())
	assert.Empty(t, LatestUserResults(ctx, []string{"group"}))
	check := func(want UserRequestResult) {
		assert.Equal(t, want, LatestUserResults(ctx, []string{"group"})["group"])
		liveResults.Delete("group")
		assert.Equal(t, want, LatestUserResults(ctx, []string{"group"})["group"])
	}
	recordUserResult("other", UserRequestResult{At: 2000, Success: true})
	recordUserResult("group", UserRequestResult{At: 2000, Success: true})
	check(UserRequestResult{At: 2000, Success: true})
	recordUserResult("group", UserRequestResult{At: 2001, Success: false})
	check(UserRequestResult{At: 2001, Success: false})
	recordUserResult("group", UserRequestResult{At: 2000, Success: true})
	check(UserRequestResult{At: 2001, Success: false})
	recordUserResult("group", UserRequestResult{At: 2002, Success: true})
	recordUserResult("group", UserRequestResult{At: 2001, Success: false})
	check(UserRequestResult{At: 2002, Success: true})
	assert.True(t, LatestUserResults(ctx, []string{"other"})["other"].Success)
	ttl, err := client.TTL(ctx, "status:stream-user-result:group").Result()
	require.NoError(t, err)
	assert.Equal(t, 15*time.Minute, ttl)
	server.FastForward(15 * time.Minute)
	assert.Zero(t, LatestUserResults(ctx, []string{"group"})["group"].At)
}

func TestNonStreamRequestsCannotChangeHealth(t *testing.T) {
	oldRedis := common.RedisEnabled
	common.RedisEnabled = false
	t.Cleanup(func() {
		common.RedisEnabled = oldRedis
		liveResults.Delete("stream-only")
		hotBuckets.Clear()
	})
	ctx := context.Background()
	info := &relaycommon.RelayInfo{UsingGroup: "stream-only", OriginModelName: "test-model", TokenId: 1, IsStream: true, StartTime: time.Now()}
	RecordRelayResult(ctx, info, nil)
	want := LatestUserResults(ctx, []string{info.UsingGroup})[info.UsingGroup]
	require.True(t, want.Success)
	require.Positive(t, want.At)
	info.IsStream = false
	for _, apiErr := range []*types.NewAPIError{nil, types.InitOpenAIError("server_error", 500), types.InitOpenAIError("server_error", 502)} {
		RecordRelayResult(ctx, info, apiErr)
		assert.Equal(t, want, LatestUserResults(ctx, []string{info.UsingGroup})[info.UsingGroup])
	}
	info.IsStream = true
	RecordRelayResult(ctx, info, types.InitOpenAIError("server_error", 502))
	assert.False(t, LatestUserResults(ctx, []string{info.UsingGroup})[info.UsingGroup].Success)
}