package cql

import (
	"net"
	"testing"
	"time"

	"go.uber.org/zap"
)

func TestServerBeginDrainStopsAcceptingConnections(t *testing.T) {
	server := NewServer("127.0.0.1:0", noopHandler{}, zap.NewNop(), &ServerOptions{
		ProtocolVersion: ProtoVersion4,
	})
	if err := server.Start(); err != nil {
		t.Fatalf("server start failed: %v", err)
	}
	defer server.Stop()

	addr := server.Addr()
	if addr == nil {
		t.Fatalf("server addr is nil")
	}

	conn, err := net.Dial("tcp", addr.String())
	if err != nil {
		t.Fatalf("dial failed: %v", err)
	}
	defer conn.Close()

	_ = waitForConnectionInfo(t, server, 2*time.Second)
	active := server.BeginDrain()
	if active < 1 {
		t.Fatalf("expected at least 1 active connection when draining, got %d", active)
	}
	if !server.IsDraining() {
		t.Fatalf("expected draining state")
	}

	if _, err := net.DialTimeout("tcp", addr.String(), 200*time.Millisecond); err == nil {
		t.Fatalf("expected dial to fail after drain")
	}
}

func TestServerBeginDrainMarksConnectionsDraining(t *testing.T) {
	server := NewServer("127.0.0.1:0", nil, zap.NewNop(), nil)
	serverConn, clientConn := net.Pipe()
	defer serverConn.Close()
	defer clientConn.Close()

	conn := NewConn(serverConn, server)

	server.mu.Lock()
	server.connections["pipe"] = conn
	server.mu.Unlock()

	server.BeginDrain()
	if !conn.isDraining() {
		t.Fatalf("expected connection to be marked draining")
	}
}

func TestServerForceCloseConnections(t *testing.T) {
	server := NewServer("127.0.0.1:0", nil, zap.NewNop(), nil)
	serverConn, clientConn := net.Pipe()
	defer clientConn.Close()

	conn := NewConn(serverConn, server)

	server.mu.Lock()
	server.connections["pipe"] = conn
	server.mu.Unlock()

	closed := server.ForceCloseConnections()
	if closed != 1 {
		t.Fatalf("expected 1 closed connection, got %d", closed)
	}

	_ = clientConn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	buf := make([]byte, 1)
	_, err := clientConn.Read(buf)
	if err == nil {
		t.Fatalf("expected client side read to fail after force close")
	}
}

func TestConnDrainStaggerSpread(t *testing.T) {
	if got := connDrainStagger(0); got != 0 {
		t.Fatalf("first connection should drain immediately, got %s", got)
	}
	if got := connDrainStagger(1); got != connDrainStaggerStep {
		t.Fatalf("expected %s for index 1, got %s", connDrainStaggerStep, got)
	}
	// Spread is capped so total drain stays well within wait_timeout.
	max := time.Duration(connDrainStaggerBuckets-1) * connDrainStaggerStep
	if got := connDrainStagger(connDrainStaggerBuckets + 1); got > max {
		t.Fatalf("stagger %s exceeds cap %s", got, max)
	}
}

func TestServerBeginDrainEventuallyMarksAllDraining(t *testing.T) {
	server := NewServer("127.0.0.1:0", nil, zap.NewNop(), nil)
	conns := make([]*Conn, 0, 5)
	for i := 0; i < 5; i++ {
		serverConn, clientConn := net.Pipe()
		t.Cleanup(func() { _ = serverConn.Close(); _ = clientConn.Close() })
		c := NewConn(serverConn, server)
		server.mu.Lock()
		server.connections[clientConn.LocalAddr().String()+"#"+itoa(i)] = c
		server.mu.Unlock()
		conns = append(conns, c)
	}

	server.BeginDrain()

	deadline := time.Now().Add(2 * time.Second)
	for {
		all := true
		for _, c := range conns {
			if !c.isDraining() {
				all = false
				break
			}
		}
		if all {
			return
		}
		if time.Now().After(deadline) {
			t.Fatal("not all connections entered drain within 2s")
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func itoa(i int) string {
	return string(rune('0' + i))
}