package cql

import (
	"net"
	"testing"
	"time"

	"go.uber.org/zap"
)

func TestMaxInFlightFramesConfig(t *testing.T) {
	server := NewServer("127.0.0.1:0", nil, zap.NewNop(), &ServerOptions{MaxInFlightFrames: 128})
	client, srv := net.Pipe()
	defer client.Close()
	defer srv.Close()

	conn := NewConn(srv, server)
	if conn.maxInFlightFrames != 128 {
		t.Fatalf("expected maxInFlightFrames 128, got %d", conn.maxInFlightFrames)
	}
}

func TestMaxInFlightFramesClamp(t *testing.T) {
	server := NewServer("127.0.0.1:0", nil, zap.NewNop(), &ServerOptions{MaxInFlightFrames: int(MaxStreamIDv3p) + 1})
	client, srv := net.Pipe()
	defer client.Close()
	defer srv.Close()

	conn := NewConn(srv, server)
	if conn.maxInFlightFrames != int(MaxStreamIDv3p) {
		t.Fatalf("expected maxInFlightFrames %d, got %d", int(MaxStreamIDv3p), conn.maxInFlightFrames)
	}
}

func TestMaxInFlightFramesMinimum(t *testing.T) {
	server := NewServer("127.0.0.1:0", nil, zap.NewNop(), &ServerOptions{MaxInFlightFrames: -5})
	client, srv := net.Pipe()
	defer client.Close()
	defer srv.Close()

	conn := NewConn(srv, server)
	if conn.maxInFlightFrames != 1 {
		t.Fatalf("expected maxInFlightFrames 1, got %d", conn.maxInFlightFrames)
	}
}

func TestBatchFramesLeaveHeadroomForNonBatch(t *testing.T) {
	server := NewServer("127.0.0.1:0", nil, zap.NewNop(), &ServerOptions{MaxInFlightFrames: 4})
	client, srv := net.Pipe()
	defer client.Close()
	defer srv.Close()

	conn := NewConn(srv, server)
	if conn.maxBatchInFlight != 3 {
		t.Fatalf("expected maxBatchInFlight 3, got %d", conn.maxBatchInFlight)
	}

	releases := make([]func(), 0, conn.maxBatchInFlight)
	for i := 0; i < conn.maxBatchInFlight; i++ {
		releases = append(releases, conn.acquireFramePermit(&Frame{Opcode: OpcodeBatch}, nil))
	}
	defer func() {
		for _, release := range releases {
			release()
		}
	}()

	selectDone := make(chan struct{})
	go func() {
		release := conn.acquireFramePermit(&Frame{Opcode: OpcodeQuery}, nil)
		release()
		close(selectDone)
	}()

	select {
	case <-selectDone:
	case <-time.After(200 * time.Millisecond):
		t.Fatal("non-batch frame should acquire reserved in-flight headroom")
	}

	batchDone := make(chan struct{})
	go func() {
		release := conn.acquireFramePermit(&Frame{Opcode: OpcodeBatch}, nil)
		release()
		close(batchDone)
	}()

	select {
	case <-batchDone:
		t.Fatal("batch frame should block when batch headroom is exhausted")
	case <-time.After(50 * time.Millisecond):
	}

	releases[0]()
	releases = releases[1:]

	select {
	case <-batchDone:
	case <-time.After(200 * time.Millisecond):
		t.Fatal("batch frame should proceed after a batch permit is released")
	}
}

func TestQueuedBatchDoesNotStealNonBatchHeadroom(t *testing.T) {
	server := NewServer("127.0.0.1:0", nil, zap.NewNop(), &ServerOptions{MaxInFlightFrames: 2})
	client, srv := net.Pipe()
	defer client.Close()
	defer srv.Close()

	conn := NewConn(srv, server)
	if conn.maxBatchInFlight != 1 {
		t.Fatalf("expected maxBatchInFlight 1, got %d", conn.maxBatchInFlight)
	}

	first := conn.acquireFramePermit(&Frame{Opcode: OpcodeQuery}, nil)
	second := conn.acquireFramePermit(&Frame{Opcode: OpcodeQuery}, nil)
	defer first()
	defer second()

	batchAcquired := make(chan struct{})
	go func() {
		release := conn.acquireFramePermit(&Frame{Opcode: OpcodeBatch}, nil)
		close(batchAcquired)
		release()
	}()

	selectAcquired := make(chan struct{})
	releaseSelect := make(chan struct{})
	go func() {
		release := conn.acquireFramePermit(&Frame{Opcode: OpcodeQuery}, nil)
		close(selectAcquired)
		<-releaseSelect
		release()
	}()

	// Wait until the non-batch goroutine has registered its reservation
	// (waitingNonBatch incremented under framePermitMu) before freeing a slot.
	// Otherwise the batch goroutine may re-evaluate the permit predicate while
	// waitingNonBatch is still 0 and steal the reserved headroom, making this
	// test flaky.
	waitDeadline := time.Now().Add(2 * time.Second)
	for {
		conn.framePermitMu.Lock()
		waiting := conn.waitingNonBatch
		conn.framePermitMu.Unlock()
		if waiting > 0 {
			break
		}
		if time.Now().After(waitDeadline) {
			t.Fatal("non-batch goroutine did not register as waiting")
		}
		time.Sleep(time.Millisecond)
	}

	second()

	select {
	case <-selectAcquired:
	case <-time.After(200 * time.Millisecond):
		t.Fatal("non-batch frame should acquire reserved headroom before queued batch")
	}

	select {
	case <-batchAcquired:
		t.Fatal("queued batch should not steal reserved headroom while non-batch is waiting")
	case <-time.After(50 * time.Millisecond):
	}

	close(releaseSelect)

	select {
	case <-batchAcquired:
	case <-time.After(200 * time.Millisecond):
		t.Fatal("batch frame should proceed after non-batch headroom is released")
	}
}

func TestFramePermitReleaseIsIdempotent(t *testing.T) {
	server := NewServer("127.0.0.1:0", nil, zap.NewNop(), &ServerOptions{MaxInFlightFrames: 2})
	client, srv := net.Pipe()
	defer client.Close()
	defer srv.Close()

	conn := NewConn(srv, server)
	release := conn.acquireFramePermit(&Frame{Opcode: OpcodeBatch}, nil)
	release()
	release()

	done := make(chan struct{})
	go func() {
		nextRelease := conn.acquireFramePermit(&Frame{Opcode: OpcodeBatch}, nil)
		nextRelease()
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(200 * time.Millisecond):
		t.Fatal("double release should not corrupt frame permit accounting")
	}
}