package native

import (
	"bufio"
	"context"
	"io"
	"net"
	"net/http"
	"sync"
	"testing"
	"time"
)

type pausedWriteListener struct {
	net.Listener
	writing chan struct{}
	release chan struct{}
}

func (l *pausedWriteListener) Accept() (net.Conn, error) {
	conn, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}
	return &pausedWriteConn{Conn: conn, writing: l.writing, release: l.release}, nil
}

type pausedWriteConn struct {
	net.Conn
	writing chan struct{}
	release chan struct{}
	once    sync.Once
}

func (c *pausedWriteConn) Write(data []byte) (int, error) {
	c.once.Do(func() { close(c.writing); <-c.release })
	return c.Conn.Write(data)
}

func TestShutdownDuringResponseWriteClosesNewlyIdleConnection(t *testing.T) {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	wrapped := &pausedWriteListener{Listener: listener, writing: make(chan struct{}), release: make(chan struct{})}
	var once sync.Once
	unblock := func() { once.Do(func() { close(wrapped.release) }) }
	defer unblock()
	defer listener.Close()
	server := New(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, "finished") }))
	server.ShutdownTimeout = 200 * time.Millisecond
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- server.Serve(ctx, wrapped) }()
	client, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()
	client.SetDeadline(time.Now().Add(2 * time.Second))
	if _, err := io.WriteString(client, "GET / HTTP/1.1\r\nHost: localhost\r\n\r\n"); err != nil {
		t.Fatal(err)
	}
	select {
	case <-wrapped.writing:
	case <-time.After(time.Second):
		t.Fatal("response did not reach write")
	}
	// The response has already chosen its Keep-Alive header. Shutdown must
	// still close this connection when its write completes and it becomes idle.
	cancel()
	deadline := time.Now().Add(time.Second)
	for !server.isClosing() {
		if time.Now().After(deadline) {
			t.Fatal("shutdown did not start")
		}
		time.Sleep(time.Millisecond)
	}
	unblock()
	response, err := http.ReadResponse(bufio.NewReader(client), &http.Request{Method: "GET"})
	if err != nil {
		t.Fatal(err)
	}
	body, err := io.ReadAll(response.Body)
	response.Body.Close()
	if err != nil || string(body) != "finished" {
		t.Fatalf("response %q: %v", body, err)
	}
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("completed request left a connection stuck in drain: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("shutdown did not complete")
	}
	if stats := server.Stats(); stats.ActiveConnections != 0 || stats.ActiveRequests != 0 {
		t.Fatalf("resources retained: %+v", stats)
	}
}