package pubsubclient

import (
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

func TestClientWhileConnectedMakesConn(t *testing.T) {
	// When c.whileConnected is called and the client is
	// uninitialized, the c.newConn function should get called.
	called := false
	newConn := func() (pubsubConn, error) {
		called = true
		return newMockPubsubConn(), nil
	}

	c := newClient(newConn)
	defer func() {
		assert.NoError(t, c.Close())
	}()

	err := c.whileConnected(func(pubsubConn) error { return nil })
	assert.NoError(t, err)
	assert.True(t, called)
	assert.Equal(t, c.nReconnectFailures, 0)
}

func TestClientWhileConnectedMakesConnOnce(t *testing.T) {
	// When c.whileConnected is called and the client is
	// uninitialized, the c.newConn function should get called by just
	// one goroutine.
	nCalls := 0
	newConn := func() (pubsubConn, error) {
		nCalls++
		return newMockPubsubConn(), nil
	}

	c := newClient(newConn)
	defer func() {
		assert.NoError(t, c.Close())
	}()

	const nGoroutine = 10000
	results := make(chan int, nGoroutine)
	var wg sync.WaitGroup
	for i := 0; i < nGoroutine; i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			_ = c.whileConnected(func(pubsubConn) error {
				results <- id
				return nil
			})
		}(i)
	}
	wg.Wait()
	close(results)
	assert.Equal(t, 1, nCalls)

	have := make(map[int]struct{})
	for i := range results {
		have[i] = struct{}{}
	}
	assert.Len(t, have, nGoroutine)
	assert.Equal(t, c.nReconnectFailures, 0)
}

func TestClientWhileConnectedDisconnectsOnFailure(t *testing.T) {
	// When c.whileConnected is called and the newConn function returns
	// an error, the client's state should go to disconnected.
	newConn := func() (pubsubConn, error) {
		return nil, errors.New("failed")
	}

	c := newClient(newConn)
	defer assert.Error(t, c.Close())

	err := c.whileConnected(func(pubsubConn) error { return nil })
	assert.Error(t, err)
	assert.Equal(t, ErrDisconnected, err)
}

func TestClientReconnectsOnError(t *testing.T) {
	// When the conn sends an error to its error channel, the connection
	// should get rebuilt.
	mockConn := newMockPubsubConn()
	connCount := new(int64)
	newConn := func() (pubsubConn, error) {
		atomic.AddInt64(connCount, 1)
		return mockConn, nil
	}
	c := newClient(newConn)
	c.connect()
	require.EqualValues(t, 1, atomic.LoadInt64(connCount))

	mockConn.errch <- errors.New("everything is broken")
	time.Sleep(100 * time.Millisecond)
	require.EqualValues(t, 2, atomic.LoadInt64(connCount))
}

func TestClientReconnectBackoff(t *testing.T) {
	// When the client has trouble connecting multiple times, it should
	// keep retrying until it succeeds.
	mockConn := newMockPubsubConn()
	connCount := new(int64)
	newConn := func() (pubsubConn, error) {
		atomic.AddInt64(connCount, 1)
		if atomic.LoadInt64(connCount) > 3 {
			return mockConn, nil
		}
		return nil, errors.New("derp")
	}
	c := newClient(newConn)
	c.reconnectBackoffUnit = time.Millisecond

	err := c.Subscribe("topic", "")
	require.EqualValues(t, 1, atomic.LoadInt64(connCount))
	assert.Equal(t, err, ErrDisconnected)

	// Try again right away: we should retry and fail, and this time it
	// should have taken 1x the backoff unit
	start := time.Now()
	err = c.Subscribe("topic", "")
	dur := time.Since(start)

	require.EqualValues(t, 2, atomic.LoadInt64(connCount))
	assert.Equal(t, err, ErrDisconnected)
	assertBackoffWithin(t, 1*c.reconnectBackoffUnit, dur)

	// Again, try and fail; now it should be 2x the backoff
	start = time.Now()
	err = c.Subscribe("topic", "")
	dur = time.Since(start)

	require.EqualValues(t, 3, atomic.LoadInt64(connCount))
	assert.Equal(t, err, ErrDisconnected)
	assertBackoffWithin(t, 2*c.reconnectBackoffUnit, dur)

	// Take another shot.. this time it should work after waiting off for 4x the backoff
	start = time.Now()
	err = c.Subscribe("topic", "")
	dur = time.Since(start)

	require.EqualValues(t, 4, atomic.LoadInt64(connCount))
	assert.NoError(t, err)
	assertBackoffWithin(t, 4*c.reconnectBackoffUnit, dur)

	// The next call should be very fast and not error and not reconnect
	start = time.Now()
	err = c.Subscribe("topic", "")
	dur = time.Since(start)

	require.EqualValues(t, 4, atomic.LoadInt64(connCount))
	assert.NoError(t, err)
	if dur >= 4*c.reconnectBackoffUnit {
		assert.Fail(t, "took too long", "should be fast when not backing off, but it took %s", dur)
	}
}

func assertBackoffWithin(t *testing.T, want, have time.Duration) {
	if have < want {
		assert.Fail(t, "did not back off long enough", "expected to take at least %s, but it took %s", want, have)
	}
	if have > 4*want {
		assert.Fail(t, "backed off for too long", "expected to take at about %s, but it took %s, which is much higher", want, have)
	}
}

func TestClientResubscribesOnError(t *testing.T) {
	// When the conn sends an error to its error channel and the client
	// reconnects, it should re-establish its subscriptions.
	firstConn := newMockPubsubConn()
	secondConn := newMockPubsubConn()
	connCount := new(int64)
	newConn := func() (pubsubConn, error) {
		atomic.AddInt64(connCount, 1)
		if atomic.LoadInt64(connCount) == 1 {
			return firstConn, nil
		}
		return secondConn, nil
	}
	c := newClient(newConn)
	c.connStateCond.L.Lock()
	c.connect()
	c.connStateCond.L.Unlock()

	require.EqualValues(t, 1, atomic.LoadInt64(connCount))

	c.subscriptions.listen(
		&listenMsg{
			msg:   msg{listen},
			Nonce: "req1",
			Data: topicList{
				Topics:    []string{"topic1"},
				AuthToken: "token1",
			},
		},
	)
	c.subscriptions.listen(
		&listenMsg{
			msg:   msg{listen},
			Nonce: "req2",
			Data: topicList{
				Topics:    []string{"topic2"},
				AuthToken: "token2",
			},
		},
	)

	var l sync.Mutex
	var recvd []*listenMsg
	secondConn.sendStub = func(msg typedMessage) error {
		l.Lock()
		recvd = append(recvd, msg.(*listenMsg))
		l.Unlock()
		return nil
	}

	firstConn.errch <- errors.New("everything is broken")
	time.Sleep(100 * time.Millisecond)

	require.EqualValues(t, 2, atomic.LoadInt64(connCount))

	assert.NoError(t, c.Close())

	l.Lock()
	have := c.subscriptions.replay()
	assert.Len(t, recvd, len(have))
	l.Unlock()
}

func TestClientWhileConnectedDoesNotRunFunctionWhileDisconnected(t *testing.T) {
	// When c.whileConnected is called and the newConn function returns
	// an error, the client's state should go to disconnected.
	c := newClient(nil)

	c.connState = disconnected

	called := false
	_ = c.whileConnected(func(pubsubConn) error {
		called = true
		return nil
	})
	assert.False(t, called)
}

func TestClientSubscribe(t *testing.T) {
	mockConn := newMockPubsubConn()
	newConn := func() (pubsubConn, error) {
		return mockConn, nil
	}
	c := newClient(newConn)
	defer func() {
		assert.NoError(t, c.Close())
	}()

	err := c.Subscribe("topic", "token")
	assert.NoError(t, err)
	require.Equal(t, 1, len(mockConn.sendCallHistory()))

	call := mockConn.sendCallHistory()[0]
	require.IsType(t, &listenMsg{}, call)

	listen := call.(*listenMsg)

	assert.Equal(t, "token", listen.Data.AuthToken)
	require.Len(t, listen.Data.Topics, 1)
	assert.Equal(t, "topic", listen.Data.Topics[0])
}

func TestClientSubscribeError(t *testing.T) {
	mockConn := newMockPubsubConn()
	sendErr := errors.New("failure")
	mockConn.sendStub = func(typedMessage) error {
		return sendErr
	}
	newConn := func() (pubsubConn, error) {
		return mockConn, nil
	}
	c := newClient(newConn)
	defer func() {
		assert.NoError(t, c.Close())
	}()

	err := c.Subscribe("topic", "token")
	require.Error(t, err)
	assert.Equal(t, sendErr, err)
}

func TestClientSubscribeWhileConnected(t *testing.T) {
	mockConn := newMockPubsubConn()
	connectCount := 0
	newConn := func() (pubsubConn, error) {
		connectCount++
		return mockConn, nil
	}
	c := newClient(newConn)
	defer func() {
		assert.NoError(t, c.Close())
	}()

	c.connect()
	require.Equal(t, 1, connectCount)

	err := c.Subscribe("topic", "token")
	require.Equal(t, 1, connectCount)
	assert.NoError(t, err)
}

func TestClientSubscribeWhileDisconnected(t *testing.T) {
	mockConn := newMockPubsubConn()
	newConn := func() (pubsubConn, error) {
		return mockConn, nil
	}
	c := newClient(newConn)
	c.connState = disconnected

	err := c.Subscribe("topic", "token")
	assert.Error(t, err)
	assert.Equal(t, ErrDisconnected, err)
}

func TestClientUnsubscribe(t *testing.T) {
	mockConn := newMockPubsubConn()
	newConn := func() (pubsubConn, error) {
		return mockConn, nil
	}
	c := newClient(newConn)
	defer func() {
		assert.NoError(t, c.Close())
	}()

	err := c.Unsubscribe("topic")
	assert.NoError(t, err)
	require.Equal(t, 1, len(mockConn.sendCallHistory()))

	call := mockConn.sendCallHistory()[0]
	require.IsType(t, &unlistenMsg{}, call)

	unlisten := call.(*unlistenMsg)

	require.Len(t, unlisten.Data.Topics, 1)
	assert.Equal(t, "topic", unlisten.Data.Topics[0])
}

func TestClientUnsubscribeError(t *testing.T) {
	mockConn := newMockPubsubConn()
	sendErr := errors.New("failure")
	mockConn.sendStub = func(typedMessage) error {
		return sendErr
	}
	newConn := func() (pubsubConn, error) {
		return mockConn, nil
	}
	c := newClient(newConn)
	defer func() {
		assert.NoError(t, c.Close())
	}()

	err := c.Unsubscribe("topic")
	require.Error(t, err)
	assert.Equal(t, sendErr, err)
}

func TestClientUnsubscribeWhileDisconnected(t *testing.T) {
	mockConn := newMockPubsubConn()
	newConn := func() (pubsubConn, error) {
		return mockConn, nil
	}
	c := newClient(newConn)

	c.connState = disconnected

	err := c.Unsubscribe("topic")
	assert.Error(t, err)
	assert.Equal(t, ErrDisconnected, err)
}

func TestClientUnsubscribeWhileConnecting(t *testing.T) {
	mockConn := newMockPubsubConn()
	newConn := func() (pubsubConn, error) {
		return mockConn, nil
	}
	c := newClient(newConn)

	c.connState = connecting

	result := make(chan error)
	go func() {
		result <- c.Unsubscribe("topic")
	}()

	select {
	case <-result:
		require.Fail(t, "should block while connecting")
	case <-time.After(50 * time.Millisecond):
	}

	c.connStateCond.L.Lock()
	c.connect()
	c.connStateCond.L.Unlock()

	select {
	case err := <-result:
		assert.NoError(t, err)
	case <-time.After(50 * time.Millisecond):
		require.Fail(t, "should be unblocked after conn is established")
	}
}

func TestClientCloseWhileReading(t *testing.T) {
	// A call to client.Close() should not be blocked on any pending
	// client.NextMessage() calls.
	mockConn := newMockPubsubConn()
	newConn := func() (pubsubConn, error) {
		return mockConn, nil
	}
	c := newClient(newConn)

	for i := 0; i < 10; i++ {
		go func() {
			_, _ = c.NextMessage()
		}()
	}

	closeDone := make(chan struct{})
	go func() {
		_ = c.Close()
		close(closeDone)
	}()

	select {
	case <-closeDone:
	case <-time.After(time.Second):
		t.Error("timed out waiting for client.Close() to return")
	}
}

func TestPreconnectSleepNoOverflow(t *testing.T) {
	mockConn := newMockPubsubConn()
	newConn := func() (pubsubConn, error) {
		return mockConn, nil
	}
	c := newClient(newConn)

	c.nReconnectFailures = 1000
	c.maxReconnectSleep = 10 * time.Second

	var have time.Duration
	require.NotPanics(t, func() {
		have = c.preconnectSleep()
	})
	assert.Equal(t, c.maxReconnectSleep, have)
}

func TestPreconnectSleep(t *testing.T) {
	mockConn := newMockPubsubConn()
	newConn := func() (pubsubConn, error) {
		return mockConn, nil
	}
	c := newClient(newConn)

	c.nReconnectFailures = 3
	c.maxReconnectSleep = 10 * time.Second
	c.reconnectBackoffUnit = time.Second

	var have time.Duration
	require.NotPanics(t, func() {
		have = c.preconnectSleep()
	})
	assert.Equal(t, 4*time.Second, have)
}
