package http2

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"testing"
	"time"

	xhttp2 "golang.org/x/net/http2"
	"golang.org/x/net/http2/hpack"
)

// A zero receive window simulates a client that has stopped consuming DATA.
// Headers need no flow-control credit, so a 204 on another stream should still
// complete while the large response waits for a WINDOW_UPDATE or cancellation.
func TestSlowResponseStreamIsolationAndRecovery(t *testing.T) {
	payload := strings.Repeat("x", 32<<10)
	started := make(chan struct{}, 2)
	written := make(chan error, 2)
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.URL.Path == "/health" {
			writer.WriteHeader(http.StatusNoContent)
			return
		}
		if request.URL.Path == "/ready" {
			_, _ = io.WriteString(writer, "ready")
			return
		}
		writer.Header().Set("Content-Length", fmt.Sprint(len(payload)))
		writer.(http.Flusher).Flush()
		started <- struct{}{}
		_, err := io.WriteString(writer, payload)
		written <- err
	}))
	server.MaxConnections = 2
	server.MaxConcurrentRequests = 2
	address, stop := startServer(t, server)
	defer stop()

	slow := newFramePeer(t, address, xhttp2.Setting{ID: xhttp2.SettingInitialWindowSize, Val: 0})
	defer slow.Close()
	slow.framer.ReadMetaHeaders = hpack.NewDecoder(4096, nil)
	slow.get(t, 1, "/large")
	waitSignal(t, started)

	// The same TCP connection can still carry response headers for other streams.
	slow.get(t, 3, "/health")
	readNoContent(t, slow, 3)
	select {
	case err := <-written:
		t.Fatalf("large response escaped the zero stream window: %v", err)
	default:
	}

	// A second connection also remains usable while the first handler is blocked.
	fast := newFramePeer(t, address)
	defer fast.Close()
	fast.framer.ReadMetaHeaders = hpack.NewDecoder(4096, nil)
	fast.get(t, 1, "/ready")
	if body := fast.body(t, 1); body != "ready" {
		t.Fatalf("other connection response: %q", body)
	}
	if stats := server.Stats(); stats.ActiveConnections != 2 || stats.RejectedRequests != 0 {
		t.Fatalf("slow response affected other traffic: %+v", stats)
	}

	// Restoring stream credit must resume the existing response without reconnecting.
	if err := slow.framer.WriteWindowUpdate(1, uint32(len(payload))); err != nil {
		t.Fatal(err)
	}
	if body := slow.body(t, 1); body != payload {
		t.Fatalf("resumed response mismatch: got %d bytes, want %d", len(body), len(payload))
	}
	select {
	case err := <-written:
		if err != nil {
			t.Fatalf("resumed write: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("resumed write did not return")
	}

	// A new zero-window stream must release its handler when the client cancels.
	slow.get(t, 5, "/large")
	waitSignal(t, started)
	if err := slow.framer.WriteRSTStream(5, xhttp2.ErrCodeCancel); err != nil {
		t.Fatal(err)
	}
	select {
	case err := <-written:
		if err == nil {
			t.Fatal("canceled zero-window write unexpectedly succeeded")
		}
	case <-time.After(time.Second):
		t.Fatal("RST_STREAM did not unblock the writer")
	}

	slow.get(t, 7, "/health")
	readNoContent(t, slow, 7)
	deadline := time.Now().Add(time.Second)
	for server.Stats().ActiveRequests != 0 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if stats := server.Stats(); stats.ActiveRequests != 0 || stats.RSTStreams != 1 || stats.ProtocolErrors != 0 || stats.RejectedRequests != 0 {
		t.Fatalf("flow control did not recover: %+v", stats)
	}
}

// Leave stream credit available while exhausting the shared connection window.
// DATA must pause, but headers and traffic on a different TCP connection must
// remain usable until connection-level credit is restored.
func TestConnectionWindowIsolationAndRecovery(t *testing.T) {
	const initialConnectionWindow = 65535
	payload := strings.Repeat("x", 96<<10)
	written := make(chan error, 1)
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		switch request.URL.Path {
		case "/health":
			writer.WriteHeader(http.StatusNoContent)
		case "/ready":
			_, _ = io.WriteString(writer, "ready")
		default:
			writer.Header().Set("Content-Length", fmt.Sprint(len(payload)))
			writer.(http.Flusher).Flush()
			_, err := io.WriteString(writer, payload)
			written <- err
		}
	}))
	server.MaxConnections = 2
	server.MaxConcurrentRequests = 2
	address, stop := startServer(t, server)
	defer stop()

	slow := newFramePeer(t, address, xhttp2.Setting{ID: xhttp2.SettingInitialWindowSize, Val: 128 << 10})
	defer slow.Close()
	slow.framer.ReadMetaHeaders = hpack.NewDecoder(4096, nil)
	slow.get(t, 1, "/large")
	var received strings.Builder
	for received.Len() < initialConnectionWindow {
		switch frame := slow.read(t).(type) {
		case *xhttp2.DataFrame:
			if frame.StreamID != 1 || frame.StreamEnded() {
				t.Fatalf("unexpected DATA before connection credit was restored: %+v", frame)
			}
			received.Write(frame.Data())
		case *xhttp2.RSTStreamFrame:
			t.Fatalf("unexpected server reset: %+v", frame)
		}
	}
	if received.Len() != initialConnectionWindow {
		t.Fatalf("server exceeded connection window: got %d bytes", received.Len())
	}

	slow.get(t, 3, "/health")
	readNoContent(t, slow, 3)
	select {
	case err := <-written:
		t.Fatalf("large response escaped exhausted connection window: %v", err)
	default:
	}

	fast := newFramePeer(t, address)
	defer fast.Close()
	fast.get(t, 1, "/ready")
	if body := fast.body(t, 1); body != "ready" {
		t.Fatalf("other connection response: %q", body)
	}

	remaining := len(payload) - received.Len()
	if err := slow.framer.WriteWindowUpdate(0, uint32(remaining)); err != nil {
		t.Fatal(err)
	}
	received.WriteString(slow.body(t, 1))
	if received.String() != payload {
		t.Fatalf("resumed connection response mismatch: got %d bytes, want %d", received.Len(), len(payload))
	}
	select {
	case err := <-written:
		if err != nil {
			t.Fatalf("resumed connection write: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("connection WINDOW_UPDATE did not unblock writer")
	}

	if err := slow.framer.WriteWindowUpdate(0, uint32(len("ready"))); err != nil {
		t.Fatal(err)
	}
	slow.get(t, 5, "/ready")
	if body := slow.body(t, 5); body != "ready" {
		t.Fatalf("reused connection response: %q", body)
	}
	waitHTTP2Stats(t, server, func(stats Stats) bool { return stats.ActiveRequests == 0 })
	stats := server.Stats()
	if stats.AcceptedConnections != 2 || stats.ActiveConnections != 2 || stats.RejectedConnections != 0 || stats.RejectedRequests != 0 || stats.ProtocolErrors != 0 || stats.RSTStreams != 0 {
		t.Fatalf("connection flow control did not recover: %+v", stats)
	}
}

func TestShutdownReleasesFlowControlledWriter(t *testing.T) {
	started := make(chan struct{})
	written := make(chan error, 1)
	requestContext := make(chan context.Context, 1)
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		requestContext <- request.Context()
		writer.(http.Flusher).Flush()
		close(started)
		_, err := io.WriteString(writer, strings.Repeat("x", 32<<10))
		written <- err
	}))
	address, done := launchServer(t, server, context.Background())
	peer := newFramePeer(t, address, xhttp2.Setting{ID: xhttp2.SettingInitialWindowSize, Val: 0})
	defer peer.Close()
	peer.get(t, 1, "/blocked")
	waitSignal(t, started)
	requestCtx := <-requestContext
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	stopped := make(chan error, 1)
	go func() { stopped <- server.Shutdown(ctx) }()
	for {
		if frame, ok := peer.read(t).(*xhttp2.GoAwayFrame); ok {
			if frame.ErrCode != xhttp2.ErrCodeNo || frame.LastStreamID != 1 {
				t.Fatalf("unexpected GOAWAY for blocked writer: %+v", frame)
			}
			break
		}
	}
	if err := serveResult(t, stopped); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("shutdown with exhausted receive window: %v", err)
	}
	select {
	case err := <-written:
		if err == nil {
			t.Fatal("blocked write unexpectedly succeeded without window credit")
		}
	case <-time.After(time.Second):
		t.Fatal("shutdown deadline did not release flow-controlled writer")
	}
	if requestCtx.Err() == nil {
		t.Fatal("forced shutdown did not cancel request context")
	}
	if err := serveResult(t, done); err != nil {
		t.Fatal(err)
	}
	waitHTTP2Stats(t, server, func(stats Stats) bool { return stats.ActiveRequests == 0 && stats.ActiveConnections == 0 })
	if stats := server.Stats(); stats.GoAwayFrames != 1 || stats.ProtocolErrors != 0 {
		t.Fatalf("blocked connection did not close cleanly: %+v", stats)
	}
}

func readNoContent(t *testing.T, peer *framePeer, streamID uint32) {
	t.Helper()
	for {
		switch frame := peer.read(t).(type) {
		case *xhttp2.MetaHeadersFrame:
			if frame.StreamID != streamID {
				continue
			}
			if frame.PseudoValue("status") != "204" || !frame.StreamEnded() {
				t.Fatalf("stream %d: expected complete 204, got %+v", streamID, frame)
			}
			return
		case *xhttp2.DataFrame:
			t.Fatalf("unexpected DATA while receive window is exhausted on stream %d", frame.StreamID)
		case *xhttp2.RSTStreamFrame:
			t.Fatalf("unexpected server reset: %+v", frame)
		}
	}
}