package clickhouse

import (
	"context"
	"log/slog"
	"sync"
	"testing"
	"time"

	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestConnPool_Cap(t *testing.T) {
	tests := []struct {
		name     string
		capacity int
	}{
		{
			name:     "capacity 1",
			capacity: 1,
		},
		{
			name:     "capacity 5",
			capacity: 5,
		},
		{
			name:     "capacity 10",
			capacity: 10,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			pool := newConnPool(time.Hour, tt.capacity)
			defer pool.Close()

			assert.Equal(t, tt.capacity, pool.Cap())
		})
	}
}

func TestConnPool_Len(t *testing.T) {
	pool := newConnPool(time.Hour, 5)
	defer pool.Close()

	assert.Equal(t, 0, pool.Len(), "new pool should have length 0")

	// Add connections
	conn1 := &mockTransport{connectedAt: time.Now(), id: 1}
	pool.Put(conn1)
	assert.Equal(t, 1, pool.Len())

	conn2 := &mockTransport{connectedAt: time.Now(), id: 2}
	pool.Put(conn2)
	assert.Equal(t, 2, pool.Len())

	// Remove connection
	ctx := context.Background()
	conn, err := pool.Get(ctx)
	require.NoError(t, err)
	require.NotNil(t, conn)
	assert.Equal(t, 1, pool.Len())

	conn, err = pool.Get(ctx)
	require.NoError(t, err)
	require.NotNil(t, conn)
	assert.Equal(t, 0, pool.Len())
}

func TestConnPool_GetEmpty(t *testing.T) {
	pool := newConnPool(time.Hour, 5)
	defer pool.Close()

	ctx := context.Background()
	conn, err := pool.Get(ctx)
	require.ErrorIs(t, err, errQueueEmpty)
	assert.Nil(t, conn, "getting from empty pool should return nil")
}

func TestConnPool_PutAndGet(t *testing.T) {
	pool := newConnPool(time.Hour, 5)
	defer pool.Close()

	now := time.Now()
	conn1 := &mockTransport{connectedAt: now, id: 1}
	conn2 := &mockTransport{connectedAt: now.Add(time.Second), id: 2}
	conn3 := &mockTransport{connectedAt: now.Add(2 * time.Second), id: 3}

	// Put connections
	pool.Put(conn1)
	pool.Put(conn2)
	pool.Put(conn3)

	assert.Equal(t, 3, pool.Len())

	ctx := context.Background()

	// Get should return oldest connection first (heap min)
	retrieved, err := pool.Get(ctx)
	require.NoError(t, err)
	require.NotNil(t, retrieved)
	assert.Equal(t, 1, retrieved.connID(), "should get oldest connection first")

	retrieved, err = pool.Get(ctx)
	require.NoError(t, err)
	require.NotNil(t, retrieved)
	assert.Equal(t, 2, retrieved.connID())

	retrieved, err = pool.Get(ctx)
	require.NoError(t, err)
	require.NotNil(t, retrieved)
	assert.Equal(t, 3, retrieved.connID())

	// Pool should be empty now
	retrieved, err = pool.Get(ctx)
	require.ErrorIs(t, err, errQueueEmpty)
	assert.Nil(t, retrieved)
}

func TestConnPool_CapacityLimit(t *testing.T) {
	capacity := 3
	pool := newConnPool(time.Hour, capacity)
	defer pool.Close()

	now := time.Now()

	// Add more connections than capacity
	allConns := make([]*mockTransport, 5)
	for i := 0; i < 5; i++ {
		conn := &mockTransport{
			connectedAt: now.Add(time.Duration(i) * time.Second),
			id:          i + 1,
		}
		allConns[i] = conn
		pool.Put(conn)
	}

	// Pool should not exceed capacity
	assert.Equal(t, capacity, pool.Len())

	// Connections 4 and 5 should have been closed (rejected when pool at capacity)
	assert.True(t, allConns[3].closed, "connection 4 should be closed when pool is full")
	assert.True(t, allConns[4].closed, "connection 5 should be closed when pool is full")
	assert.False(t, allConns[0].closed, "connection 1 should not be closed")
	assert.False(t, allConns[1].closed, "connection 2 should not be closed")
	assert.False(t, allConns[2].closed, "connection 3 should not be closed")

	ctx := context.Background()

	// With FIFO insertion order, the first 3 connections are kept [1, 2, 3]
	// connections (4) and (5) are rejected when pool is at capacity
	for i := 0; i < 3; i++ {
		retrieved, err := pool.Get(ctx)
		require.NoError(t, err)
		require.NotNil(t, retrieved)

		assert.Equal(t, i+1, retrieved.connID(), "should get connections in FIFO order")
	}
}

func TestConnPool_ExpiredConnectionNotReturned(t *testing.T) {
	// Pool with very short lifetime
	lifetime := 100 * time.Millisecond
	pool := newConnPool(lifetime, 5)
	defer pool.Close()

	// Add connection that is not yet expired (but close to expiration)
	oldConn := &mockTransport{
		connectedAt: time.Now().Add(-50 * time.Millisecond),
		id:          1,
	}
	pool.Put(oldConn)

	// Add fresh connection
	freshConn := &mockTransport{
		connectedAt: time.Now(),
		id:          2,
	}
	pool.Put(freshConn)

	// Wait for the old connection to expire
	time.Sleep(60 * time.Millisecond)

	ctx := context.Background()

	// Get should skip expired connection and return fresh one
	retrieved, err := pool.Get(ctx)
	require.NoError(t, err)
	require.NotNil(t, retrieved)
	assert.Equal(t, 2, retrieved.connID(), "should skip expired connection")

	// The expired connection should have been closed during Get()
	assert.True(t, oldConn.closed, "expired connection should be closed")
	assert.False(t, freshConn.closed, "fresh connection should not be closed")

	// Pool should be empty now
	retrieved, err = pool.Get(ctx)
	require.ErrorIs(t, err, errQueueEmpty)
	assert.Nil(t, retrieved)
}

func TestConnPool_PutExpiredConnection(t *testing.T) {
	lifetime := 100 * time.Millisecond
	pool := newConnPool(lifetime, 5)
	defer pool.Close()

	// Try to put already expired connection
	expiredConn := &mockTransport{
		connectedAt: time.Now().Add(-200 * time.Millisecond),
		id:          1,
	}
	pool.Put(expiredConn)

	// Pool should not accept expired connection
	assert.Equal(t, 0, pool.Len())
}

func TestConnPool_PutOlderThanMinimumWithCapacity(t *testing.T) {
	pool := newConnPool(time.Hour, 5)
	defer pool.Close()

	now := time.Now()

	// Add a connection
	conn1 := &mockTransport{connectedAt: now, id: 1}
	pool.Put(conn1)

	// Add an older connection (but inserted second)
	olderConn := &mockTransport{connectedAt: now.Add(-time.Minute), id: 2}
	pool.Put(olderConn)

	// Both connections should be in the pool
	assert.Equal(t, 2, pool.Len())

	ctx := context.Background()
	retrieved, err := pool.Get(ctx)
	require.NoError(t, err)
	require.NotNil(t, retrieved)
	// With FIFO insertion order, we get the first inserted connection
	assert.Equal(t, 1, retrieved.connID(), "should retrieve the first inserted connection")
}

func TestConnPool_GetWithCancelledContext(t *testing.T) {
	pool := newConnPool(time.Hour, 5)
	defer pool.Close()

	// Add a connection
	conn := &mockTransport{connectedAt: time.Now(), id: 1}
	pool.Put(conn)

	// Create cancelled context
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	// Get with cancelled context should return context error
	retrieved, err := pool.Get(ctx)
	assert.Error(t, err)
	assert.Equal(t, context.Canceled, err)
	assert.Nil(t, retrieved)

	// Connection should still be in pool
	assert.Equal(t, 1, pool.Len())
}

func TestConnPool_Close(t *testing.T) {
	pool := newConnPool(time.Hour, 5)

	// Add connections
	for i := 0; i < 3; i++ {
		conn := &mockTransport{
			connectedAt: time.Now(),
			id:          i + 1,
		}
		pool.Put(conn)
	}

	assert.Equal(t, 3, pool.Len())

	// Close the pool
	err := pool.Close()
	assert.NoError(t, err)

	// Verify pool is closed
	assert.True(t, pool.closed())

	// Get should return ErrConnectionClosed from closed pool
	ctx := context.Background()
	conn, err := pool.Get(ctx)
	assert.Error(t, err)
	assert.Equal(t, ErrConnectionClosed, err)
	assert.Nil(t, conn)

	// Put should be ignored on closed pool
	initialLen := pool.Len()
	newConn := &mockTransport{connectedAt: time.Now(), id: 99}
	pool.Put(newConn)
	assert.Equal(t, initialLen, pool.Len(), "closed pool should not accept new connections")

	// Closing again should be safe
	err = pool.Close()
	assert.NoError(t, err)
}

func TestConnPool_CloseWithDrain(t *testing.T) {
	pool := newConnPool(time.Hour, 5)

	// Add connections
	allConns := make([]*mockTransport, 3)
	for i := 0; i < 3; i++ {
		mock := &mockTransport{
			connectedAt: time.Now(),
			id:          i + 1,
		}
		allConns[i] = mock
		pool.Put(mock)
	}

	assert.Equal(t, 3, pool.Len(), "pool should have 3 connections before close")

	// Close the pool
	err := pool.Close()
	assert.NoError(t, err)

	// Verify pool is closed
	assert.True(t, pool.closed())

	// Verify all connections are drained from the pool
	assert.Equal(t, 0, pool.Len(), "pool should be empty after close (all connections drained)")

	// Verify all connections were closed
	for i, conn := range allConns {
		assert.True(t, conn.closed, "connection %d should be closed after pool close", i+1)
	}

	// Verify no connections can be retrieved
	ctx := context.Background()
	conn, err := pool.Get(ctx)
	assert.Error(t, err)
	assert.Equal(t, ErrConnectionClosed, err)
	assert.Nil(t, conn, "get should return nil after pool is closed and drained")
}

func TestConnPool_DrainExpiredConnections(t *testing.T) {
	lifetime := 100 * time.Millisecond
	pool := newConnPool(lifetime, 5)
	defer pool.Close()

	// Add connections that are already old (so they will definitely expire)
	oldTime := time.Now().Add(-50 * time.Millisecond)
	expiredConns := make([]*mockTransport, 3)
	for i := 0; i < 3; i++ {
		conn := &mockTransport{
			connectedAt: oldTime.Add(time.Duration(i) * time.Millisecond),
			id:          i + 1,
		}
		expiredConns[i] = conn
		pool.Put(conn)
	}

	assert.Equal(t, 3, pool.Len())

	// Wait for connections to expire and drain cycle to run
	// The connections will be 50ms + 100ms (sleep) = 150ms old, exceeding the 100ms lifetime
	time.Sleep(lifetime + 50*time.Millisecond)

	// At this point the drain should have run and removed the expired connections
	assert.Equal(t, 0, pool.Len(), "all expired connections should be drained")

	// Verify all expired connections were closed
	for i, conn := range expiredConns {
		assert.True(t, conn.closed, "expired connection %d should be closed", i+1)
	}

	// Add a fresh connection to verify pool still works after drain
	freshConn := &mockTransport{
		connectedAt: time.Now(),
		id:          99,
	}
	pool.Put(freshConn)

	// Fresh connection should be in the pool
	assert.Equal(t, 1, pool.Len(), "fresh connection should be added after drain")
	assert.False(t, freshConn.closed, "fresh connection should not be closed")
}

func TestConnPool_ConcurrentAccess(t *testing.T) {
	pool := newConnPool(time.Hour, 10)
	defer pool.Close()

	ctx := context.Background()
	done := make(chan struct{})

	// Concurrent puts
	go func() {
		for i := 0; i < 20; i++ {
			conn := &mockTransport{
				connectedAt: time.Now(),
				id:          i,
			}
			pool.Put(conn)
			time.Sleep(time.Millisecond)
		}
		close(done)
	}()

	// Concurrent gets
	go func() {
		for i := 0; i < 20; i++ {
			_, _ = pool.Get(ctx)
			time.Sleep(time.Millisecond)
		}
	}()

	// Wait for puts to complete
	<-done

	// Give gets time to complete
	time.Sleep(50 * time.Millisecond)

	// Pool should not exceed capacity
	assert.LessOrEqual(t, pool.Len(), pool.Cap())
}

func TestConnPool_FIFOOrdering(t *testing.T) {
	pool := newConnPool(time.Hour, 10)
	defer pool.Close()

	now := time.Now()

	// Add connections with varying timestamps in non-chronological order
	// to test that FIFO is based on insertion order, not connection age
	connections := []struct {
		time time.Time
		id   int
	}{
		{now.Add(1 * time.Second), 1},
		{now.Add(5 * time.Second), 5},
		{now.Add(2 * time.Second), 2},
		{now.Add(4 * time.Second), 4},
		{now.Add(3 * time.Second), 3},
	}

	// Add all connections
	for _, c := range connections {
		conn := &mockTransport{
			connectedAt: c.time,
			id:          c.id,
		}
		pool.Put(conn)
	}

	expectedCount := 5
	assert.Equal(t, expectedCount, pool.Len())

	ctx := context.Background()

	// Get all connections and verify they come out in insertion order (FIFO)
	for _, id := range []int{1, 5, 2, 4, 3} {
		conn, err := pool.Get(ctx)
		require.NoError(t, err)
		require.NotNil(t, conn, "should get connection %d", id)
		assert.Equal(t, id, conn.connID(),
			"connections should be returned in FIFO insertion order")
	}
}

// mockTransport implements nativeTransport for testing
type mockTransport struct {
	connectedAt   time.Time
	id            int
	released      bool
	closed        bool
	bad           bool
	bufferFreed   bool
	debugMessages []string
	mu            sync.Mutex
}

func (m *mockTransport) serverVersion() (*ServerVersion, error) {
	return nil, nil
}

func (m *mockTransport) query(ctx context.Context, release nativeTransportRelease, query string, args ...any) (*rows, error) {
	return nil, nil
}

func (m *mockTransport) queryRow(ctx context.Context, release nativeTransportRelease, query string, args ...any) *row {
	return nil
}

func (m *mockTransport) prepareBatch(ctx context.Context, release nativeTransportRelease, acquire nativeTransportAcquire, query string, opts driver.PrepareBatchOptions) (driver.Batch, error) {
	return nil, nil
}

func (m *mockTransport) exec(ctx context.Context, query string, args ...any) error {
	return nil
}

func (m *mockTransport) asyncInsert(ctx context.Context, query string, wait bool, args ...any) error {
	return nil
}

func (m *mockTransport) ping(ctx context.Context) error {
	return nil
}

func (m *mockTransport) isBad() bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.bad
}

func (m *mockTransport) connID() int {
	return m.id
}

func (m *mockTransport) connectedAtTime() time.Time {
	return m.connectedAt
}

func (m *mockTransport) isReleased() bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.released
}

func (m *mockTransport) setReleased(released bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.released = released
}

func (m *mockTransport) getLogger() *slog.Logger {
	return newNoopLogger()
}

func (m *mockTransport) freeBuffer() {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.bufferFreed = true
}

func (m *mockTransport) close() error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.closed = true
	return nil
}

// Helper methods for testing
func newMockTransport(id int) *mockTransport {
	return &mockTransport{
		connectedAt: time.Now(),
		id:          id,
	}
}

func (m *mockTransport) setBad(bad bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.bad = bad
}

func (m *mockTransport) isClosed() bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.closed
}

func (m *mockTransport) wasBufferFreed() bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.bufferFreed
}