package cql

import (
	"context"
	"net"
	"testing"
	"time"

	"go.uber.org/zap"
)

func newPipeConn(t *testing.T) (*Conn, net.Conn) {
	t.Helper()
	server := NewServer("127.0.0.1:0", nil, zap.NewNop(), nil)
	client, srv := net.Pipe()
	t.Cleanup(func() { _ = client.Close() })
	conn := NewConn(srv, server)
	return conn, client
}

func TestIsDrainingDefaultsFalse(t *testing.T) {
	conn, _ := newPipeConn(t)
	if conn.isDraining() {
		t.Fatal("new connection should not be draining")
	}
}

func TestBeginDrainSetsDrainingFlag(t *testing.T) {
	conn, _ := newPipeConn(t)
	conn.beginDrain()
	if !conn.isDraining() {
		t.Fatal("beginDrain should mark connection draining")
	}
}

func TestBeginDrainUnblocksBlockedRead(t *testing.T) {
	conn, _ := newPipeConn(t)
	readErr := make(chan error, 1)
	go func() {
		buf := make([]byte, 1)
		_, err := conn.conn.Read(buf)
		readErr <- err
	}()
	// Give the reader time to block.
	time.Sleep(20 * time.Millisecond)
	conn.beginDrain()
	select {
	case err := <-readErr:
		if err == nil {
			t.Fatal("expected blocked read to fail after beginDrain")
		}
	case <-time.After(time.Second):
		t.Fatal("beginDrain did not unblock the read within 1s")
	}
}

func TestDrainInFlightReturnsWhenIdle(t *testing.T) {
	conn, _ := newPipeConn(t)
	done := make(chan struct{})
	go func() {
		conn.drainInFlight(time.Second)
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(500 * time.Millisecond):
		t.Fatal("drainInFlight should return immediately with no in-flight frames")
	}
}

func TestDrainInFlightWaitsForInFlight(t *testing.T) {
	conn, _ := newPipeConn(t)
	release := conn.acquireFramePermit(&Frame{Opcode: OpcodeQuery}, nil)

	done := make(chan struct{})
	go func() {
		conn.drainInFlight(2 * time.Second)
		close(done)
	}()

	select {
	case <-done:
		t.Fatal("drainInFlight returned while a frame was still in flight")
	case <-time.After(100 * time.Millisecond):
	}

	release()

	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("drainInFlight did not return after in-flight frame completed")
	}
}

func TestDrainInFlightReturnsOnClose(t *testing.T) {
	conn, _ := newPipeConn(t)
	_ = conn.acquireFramePermit(&Frame{Opcode: OpcodeQuery}, nil) // never released

	done := make(chan struct{})
	go func() {
		conn.drainInFlight(5 * time.Second)
		close(done)
	}()

	select {
	case <-done:
		t.Fatal("drainInFlight returned before close while frame in flight")
	case <-time.After(100 * time.Millisecond):
	}

	_ = conn.Close()

	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("drainInFlight should return once the connection is closed")
	}
}

func TestDrainInFlightHonorsTimeout(t *testing.T) {
	conn, _ := newPipeConn(t)
	_ = conn.acquireFramePermit(&Frame{Opcode: OpcodeQuery}, nil) // never released

	start := time.Now()
	conn.drainInFlight(150 * time.Millisecond)
	if elapsed := time.Since(start); elapsed < 150*time.Millisecond {
		t.Fatalf("drainInFlight returned before timeout: %s", elapsed)
	}
}

// blockingHandler answers STARTUP immediately and blocks each QUERY until the
// test releases it, so a request can be held in flight while drain is triggered.
type blockingHandler struct {
	inFlight chan struct{} // a QUERY entering HandleMessage signals here
	release  chan struct{} // closed by the test to let blocked QUERYs return
}

func (h *blockingHandler) HandleMessage(_ context.Context, f *Frame) (*Frame, error) {
	switch f.Opcode {
	case OpcodeStartup:
		return &Frame{Opcode: OpcodeReady}, nil
	case OpcodeQuery:
		h.inFlight <- struct{}{}
		<-h.release
		return &Frame{Opcode: OpcodeResult, Body: []byte{0, 0, 0, 1}}, nil // void result
	default:
		return &Frame{Opcode: OpcodeReady}, nil
	}
}

func readFrameWithTimeout(t *testing.T, c net.Conn) (*Frame, error) {
	t.Helper()
	_ = c.SetReadDeadline(time.Now().Add(2 * time.Second))
	return ReadFrame(c)
}

func handshake(t *testing.T, client net.Conn) {
	t.Helper()
	startup := BuildStartupFrame("3.4.0", "", "", "")
	startup.Version, startup.Stream = ProtoVersion4, 1
	if err := startup.WriteFrame(client); err != nil {
		t.Fatalf("write startup: %v", err)
	}
	r, err := readFrameWithTimeout(t, client)
	if err != nil || r.Opcode != OpcodeReady {
		t.Fatalf("want READY for startup, got %v err=%v", r, err)
	}
}

func writeQuery(t *testing.T, client net.Conn, stream int16) {
	t.Helper()
	q := &Frame{Version: ProtoVersion4, Opcode: OpcodeQuery, Stream: stream, Body: []byte{0, 0, 0, 0}}
	if err := q.WriteFrame(client); err != nil {
		t.Fatalf("write query stream=%d: %v", stream, err)
	}
}

func TestServeDeliversInFlightResponseThenClosesOnDrain(t *testing.T) {
	h := &blockingHandler{inFlight: make(chan struct{}, 4), release: make(chan struct{})}
	server := NewServer("127.0.0.1:0", h, zap.NewNop(), &ServerOptions{Authenticator: &AllowAllAuthenticator{}})
	clientConn, serverConn := net.Pipe()
	defer clientConn.Close()
	defer serverConn.Close()
	conn := NewConn(serverConn, server)
	errCh := make(chan error, 1)
	go func() { errCh <- conn.Serve() }()

	handshake(t, clientConn)
	writeQuery(t, clientConn, 5)
	<-h.inFlight // request is now in flight

	conn.beginDrain() // drain while the request is still being handled
	close(h.release)  // let the in-flight request finish

	// Assertion A: the in-flight request's response is delivered (not torn).
	resp, err := readFrameWithTimeout(t, clientConn)
	if err != nil {
		t.Fatalf("in-flight response should be delivered, got err: %v", err)
	}
	if resp.Opcode != OpcodeResult || resp.Stream != 5 {
		t.Fatalf("want RESULT stream=5, got opcode=%s stream=%d", MessageType(resp.Opcode).String(), resp.Stream)
	}

	// Assertion B: the connection closes cleanly only after the response.
	if _, err := readFrameWithTimeout(t, clientConn); err == nil {
		t.Fatal("connection should close after the in-flight response is delivered")
	}

	// Assertion C: Serve exits without error (clean drain, not a torn read).
	select {
	case err := <-errCh:
		if err != nil {
			t.Fatalf("serve returned error during graceful drain: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("serve did not exit after drain")
	}
}

func TestServeDeliversAllInFlightResponsesOnDrain(t *testing.T) {
	h := &blockingHandler{inFlight: make(chan struct{}, 8), release: make(chan struct{})}
	server := NewServer("127.0.0.1:0", h, zap.NewNop(), &ServerOptions{Authenticator: &AllowAllAuthenticator{}})
	clientConn, serverConn := net.Pipe()
	defer clientConn.Close()
	defer serverConn.Close()
	conn := NewConn(serverConn, server)
	errCh := make(chan error, 1)
	go func() { errCh <- conn.Serve() }()

	handshake(t, clientConn)

	streams := []int16{5, 6, 7}
	for _, s := range streams {
		writeQuery(t, clientConn, s)
		<-h.inFlight
	}

	conn.beginDrain()
	close(h.release)

	// All three multiplexed in-flight requests must get their responses.
	got := map[int16]bool{}
	for range streams {
		resp, err := readFrameWithTimeout(t, clientConn)
		if err != nil {
			t.Fatalf("expected %d responses, error after %d: %v", len(streams), len(got), err)
		}
		if resp.Opcode != OpcodeResult {
			t.Fatalf("want RESULT, got %s", MessageType(resp.Opcode).String())
		}
		got[resp.Stream] = true
	}
	for _, s := range streams {
		if !got[s] {
			t.Fatalf("missing response for in-flight stream %d", s)
		}
	}

	if _, err := readFrameWithTimeout(t, clientConn); err == nil {
		t.Fatal("connection should close after all in-flight responses delivered")
	}
	select {
	case err := <-errCh:
		if err != nil {
			t.Fatalf("serve returned error during graceful drain: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("serve did not exit after drain")
	}
}

func TestServeClosesIdleConnectionCleanlyOnDrain(t *testing.T) {
	h := &blockingHandler{inFlight: make(chan struct{}, 1), release: make(chan struct{})}
	server := NewServer("127.0.0.1:0", h, zap.NewNop(), &ServerOptions{Authenticator: &AllowAllAuthenticator{}})
	clientConn, serverConn := net.Pipe()
	defer clientConn.Close()
	defer serverConn.Close()
	conn := NewConn(serverConn, server)
	errCh := make(chan error, 1)
	go func() { errCh <- conn.Serve() }()

	handshake(t, clientConn)
	// No in-flight request: an idle pooled connection.
	conn.beginDrain()

	// The idle connection closes promptly (no waiting for force-close timeout).
	if _, err := readFrameWithTimeout(t, clientConn); err == nil {
		t.Fatal("idle connection should close promptly on drain")
	}
	// And Serve exits cleanly (nil), not with a torn-read error.
	select {
	case err := <-errCh:
		if err != nil {
			t.Fatalf("serve returned error draining idle connection: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("serve did not exit after draining idle connection")
	}
}

// A connection saturated at its in-flight cap with a wedged handler must not
// hang the serve goroutine when the connection is force-closed (e.g. via
// server.Stop()/ForceCloseConnections, which does not pre-mark draining).
func TestServeUnblocksPermitWaiterOnClose(t *testing.T) {
	h := &blockingHandler{inFlight: make(chan struct{}, 4), release: make(chan struct{})}
	defer close(h.release) // unwind the wedged handler on any exit
	server := NewServer("127.0.0.1:0", h, zap.NewNop(), &ServerOptions{
		Authenticator:     &AllowAllAuthenticator{},
		MaxInFlightFrames: 1,
	})
	clientConn, serverConn := net.Pipe()
	defer clientConn.Close()
	defer serverConn.Close()
	conn := NewConn(serverConn, server)
	errCh := make(chan error, 1)
	go func() { errCh <- conn.Serve() }()

	handshake(t, clientConn)
	writeQuery(t, clientConn, 5)
	<-h.inFlight // handler holds the only permit and wedges

	writeQuery(t, clientConn, 6) // serve reads it, then blocks acquiring a permit

	// Wait until the serve loop is actually parked waiting for a permit.
	deadline := time.Now().Add(2 * time.Second)
	for {
		conn.framePermitMu.Lock()
		waiting := conn.waitingNonBatch
		conn.framePermitMu.Unlock()
		if waiting > 0 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("serve loop never parked on permit acquisition")
		}
		time.Sleep(5 * time.Millisecond)
	}

	// Force-close. The serve goroutine must return rather than hang forever on
	// the permit condition.
	_ = conn.Close()

	select {
	case <-errCh:
	case <-time.After(2 * time.Second):
		t.Fatal("serve hung on permit acquisition after force close")
	}
}