package gochannel

import (
	"context"
	"strconv"
	"sync"
	"testing"

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

	"github.com/ThreeDotsLabs/watermill"
	"github.com/ThreeDotsLabs/watermill/message"
)

func TestSubscribe_clean_subscriber_data(t *testing.T) {
	subCount := 100
	pubSub := NewGoChannel(
		Config{OutputChannelBuffer: int64(subCount)},
		watermill.NewStdLogger(false, false),
	)
	topicName := "test_topic"

	for i := 0; i < subCount; i++ {
		ctx, cancel := context.WithCancel(context.Background())
		_, err := pubSub.Subscribe(ctx, topicName+"_index_"+strconv.Itoa(i))
		require.NoError(t, err)
		cancel()
	}

	err := pubSub.Close()
	require.NoError(t, err)

	assert.Len(t, pubSub.subscribers, 0)
	lockCount := 0
	pubSub.subscribersByTopicLock.Range(func(_, _ any) bool {
		lockCount++
		return true
	})
	assert.Equal(t, 0, lockCount)

	assert.NoError(t, pubSub.Close())
}

func TestPublish_clean_lock_data(t *testing.T) {
	messageCount := 100
	pubSub := NewGoChannel(
		Config{OutputChannelBuffer: int64(messageCount)},
		watermill.NewStdLogger(false, false),
	)
	topicName := "test_topic"

	_, err := pubSub.Subscribe(context.Background(), topicName+"_index_"+strconv.Itoa(0))
	require.NoError(t, err)

	for i := 0; i < messageCount; i++ {
		err := pubSub.Publish(topicName+"_index_"+strconv.Itoa(i), message.NewMessage(watermill.NewShortUUID(), nil))
		require.NoError(t, err)
	}

	lockCount := 0
	pubSub.subscribersByTopicLock.Range(func(_, _ any) bool {
		lockCount++
		return true
	})
	assert.Equal(t, 1, lockCount)

	assert.NoError(t, pubSub.Close())
}

// TestPublish_concurrent_close ensures Close does not race with in-flight persistent Publish calls. 
func TestPublish_concurrent_close(t *testing.T) {
	pubSub := NewGoChannel(
		Config{Persistent: true},
		watermill.NewStdLogger(false, false),
	)
	topicName := "test_topic"

	publisherCount := 20
	startPublishing := make(chan struct{})
	wg := sync.WaitGroup{}

	for i := 0; i < publisherCount; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-startPublishing
			for {
				err := pubSub.Publish(topicName, message.NewMessage(watermill.NewShortUUID(), nil))
				if err != nil {
					return
				}
			}
		}()
	}

	close(startPublishing)
	require.NoError(t, pubSub.Close())
	wg.Wait()
}