package controller

import (
	"net/http"
	"testing"

	"github.com/QuantumNous/new-api/model"

	"github.com/stretchr/testify/require"
)

func restoreRelayConcurrencyStub(t *testing.T) {
	t.Helper()

	previousAcquire := acquireChannelConcurrencyForRelay
	t.Cleanup(func() {
		acquireChannelConcurrencyForRelay = previousAcquire
	})
}

func TestAcquireRetryChannelConcurrencySkipsInitialAttempt(t *testing.T) {
	restoreRelayConcurrencyStub(t)

	acquireChannelConcurrencyForRelay = func(channelId, limit int) (bool, func()) {
		t.Fatalf("expected initial attempt to use distributor-owned acquisition")
		return false, func() {}
	}

	release, err := acquireRetryChannelConcurrency(&model.Channel{
		Id:               1,
		ConcurrencyLimit: 1,
	}, 0)

	require.Nil(t, err)
	require.NotNil(t, release)
	release()
}

func TestAcquireRetryChannelConcurrencyRejectsExceededRetryChannel(t *testing.T) {
	restoreRelayConcurrencyStub(t)

	acquireChannelConcurrencyForRelay = func(channelId, limit int) (bool, func()) {
		require.Equal(t, 2, channelId)
		require.Equal(t, 1, limit)
		return false, func() {}
	}

	release, err := acquireRetryChannelConcurrency(&model.Channel{
		Id:               2,
		ConcurrencyLimit: 1,
	}, 1)

	require.Nil(t, release)
	require.NotNil(t, err)
	require.Equal(t, http.StatusTooManyRequests, err.StatusCode)
}

func TestAcquireRetryChannelConcurrencyReleasesRetryChannel(t *testing.T) {
	restoreRelayConcurrencyStub(t)

	released := false
	acquireChannelConcurrencyForRelay = func(channelId, limit int) (bool, func()) {
		return true, func() {
			released = true
		}
	}

	release, err := acquireRetryChannelConcurrency(&model.Channel{
		Id:               3,
		ConcurrencyLimit: 1,
	}, 1)

	require.Nil(t, err)
	require.NotNil(t, release)
	release()
	require.True(t, released)
}