package http2

import (
	"bytes"
	"context"
	"crypto/sha256"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/http"
	"runtime"
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

type mixedConnectionCounts struct {
	health        atomic.Uint64
	slowUploads   atomic.Uint64
	slowDownloads atomic.Uint64
}

// TestTLSMixedSoak shares each of eight connections between fourteen fast
// workers, one paced uploader and one paced downloader. The 8 MiB download is
// larger than x/net v0.25.0's 4 MiB initial stream window, requiring window
// updates to finish. This does not exhaust the connection-level receive window.
func TestTLSMixedSoak(t *testing.T) {
	const connectionCount, workersPerConnection = 8, 16
	const uploadBytes, downloadBytes = 256 << 10, 8 << 20
	started := time.Now()
	duration := soakDuration(t)
	logSoakResources(t, started, "before", 0, 0, Stats{})
	t.Logf("scenario=mixed connections=%d workersPerConnection=%d uploadBytes=%d uploadChunkBytes=4096 uploadDelay=2ms downloadBytes=%d downloadChunkBytes=65536 downloadDelay=20ms",
		connectionCount, workersPerConnection, uploadBytes, downloadBytes)
	downloadPayload := bytes.Repeat([]byte("x"), downloadBytes)
	expectedDownload := sha256.Sum256(downloadPayload)
	expectedUpload := sha256.Sum256(downloadPayload[:uploadBytes])
	certificate, roots := testCertificate(t)
	var serverFailures atomic.Uint64
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		writer.Header().Set("X-Soak-Request-Id", request.Header.Get("X-Soak-Request-Id"))
		switch request.URL.Path {
		case "/slowUpload":
			digest := sha256.New()
			count, err := io.Copy(digest, request.Body)
			if err != nil || count != uploadBytes || !bytes.Equal(digest.Sum(nil), expectedUpload[:]) {
				serverFailures.Add(1)
				http.Error(writer, "incomplete upload", http.StatusBadRequest)
				return
			}
			_, _ = io.WriteString(writer, "uploaded")
		case "/slowDownload":
			writer.Header().Set("Content-Length", strconv.Itoa(downloadBytes))
			if n, err := writer.Write(downloadPayload); err != nil || n != downloadBytes {
				serverFailures.Add(1)
			}
		default:
			_, _ = io.WriteString(writer, "healthy")
		}
	}))
	server.MaxConnections = connectionCount
	server.MaxConcurrentRequests = connectionCount * workersPerConnection
	server.MaxConcurrentStreams = workersPerConnection
	server.MaxBodyBytes = 1 << 20
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() {
		done <- server.ServeTLS(context.Background(), listener, &tls.Config{
			Certificates: []tls.Certificate{certificate}, MinVersion: tls.VersionTLS12,
		})
	}()
	var transports []*xhttp2.Transport
	var clients []*http.Client
	for index := 0; index < connectionCount; index++ {
		transport := &xhttp2.Transport{
			TLSClientConfig:            &tls.Config{RootCAs: roots, ServerName: "localhost", MinVersion: tls.VersionTLS12, NextProtos: []string{"h2"}},
			StrictMaxConcurrentStreams: true,
		}
		transports = append(transports, transport)
		clients = append(clients, &http.Client{Transport: transport, Timeout: 20 * time.Second})
	}
	var stopOnce sync.Once
	stop := func() {
		stopOnce.Do(func() {
			for _, transport := range transports {
				transport.CloseIdleConnections()
			}
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			if err := server.Shutdown(ctx); err != nil {
				t.Errorf("mixed soak shutdown: %v", err)
			}
			_ = listener.Close()
			if err := serveResult(t, done); err != nil {
				t.Errorf("mixed soak serve: %v", err)
			}
		})
	}
	defer stop()

	var counts [connectionCount]mixedConnectionCounts
	var requests, failures atomic.Uint64
	var failureOnce sync.Once
	var firstFailure string
	recordFailure := func(err error) {
		failures.Add(1)
		failureOnce.Do(func() { firstFailure = err.Error() })
	}
	logCounts := func(phase string) {
		var health, uploads, downloads uint64
		for i := range counts {
			health += counts[i].health.Load()
			uploads += counts[i].slowUploads.Load()
			downloads += counts[i].slowDownloads.Load()
		}
		t.Logf("mixedTotals phase=%s health=%d slowUploads=%d slowDownloads=%d uploadedBytes=%d downloadedBytes=%d serverFailures=%d",
			phase, health, uploads, downloads, uploads*uploadBytes, downloads*downloadBytes, serverFailures.Load())
	}
	monitorStop, monitorDone := make(chan struct{}), make(chan struct{})
	go func() {
		defer close(monitorDone)
		ticker := time.NewTicker(15 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				logSoakResources(t, started, "running", requests.Load(), failures.Load(), server.Stats())
				logCounts("running")
			case <-monitorStop:
				return
			}
		}
	}()
	var monitorOnce sync.Once
	stopMonitor := func() { monitorOnce.Do(func() { close(monitorStop); <-monitorDone }) }
	defer stopMonitor()

	deadline := time.Now().Add(duration)
	var workers sync.WaitGroup
	for worker := 0; worker < connectionCount*workersPerConnection; worker++ {
		workers.Add(1)
		go func(worker int) {
			defer workers.Done()
			connection := worker / workersPerConnection
			kind := worker % workersPerConnection
			for sequence := uint64(0); time.Now().Before(deadline) && failures.Load() == 0; sequence++ {
				ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
				path, method := "/health", http.MethodGet
				var body io.Reader
				if kind == 0 {
					path, method = "/slowUpload", http.MethodPost
					body = &pacedUpload{ctx: ctx, remaining: uploadBytes}
				} else if kind == 1 {
					path = "/slowDownload"
				}
				requestID := strconv.Itoa(worker) + ":" + strconv.FormatUint(sequence, 10)
				request, err := http.NewRequestWithContext(ctx, method, "https://"+listener.Addr().String()+path, body)
				if err == nil {
					request.Header.Set("X-Soak-Request-Id", requestID)
					err = checkMixedResponse(ctx, clients[connection], request, kind, downloadBytes, expectedDownload)
				}
				cancel()
				if err != nil {
					recordFailure(fmt.Errorf("connection=%d path=%s: %w", connection, path, err))
					continue
				}
				requests.Add(1)
				switch kind {
				case 0:
					counts[connection].slowUploads.Add(1)
				case 1:
					counts[connection].slowDownloads.Add(1)
				default:
					counts[connection].health.Add(1)
				}
			}
		}(worker)
	}
	workers.Wait()
	stopMonitor()
	stop()
	logSoakResources(t, started, "after", requests.Load(), failures.Load(), server.Stats())
	logCounts("after")
	var completed uint64
	for index := range counts {
		c := &counts[index]
		completed += c.health.Load() + c.slowUploads.Load() + c.slowDownloads.Load()
		t.Logf("connection=%d health=%d slowUploads=%d slowDownloads=%d", index, c.health.Load(), c.slowUploads.Load(), c.slowDownloads.Load())
		if c.health.Load() == 0 || c.slowUploads.Load() == 0 || c.slowDownloads.Load() == 0 {
			t.Errorf("connection %d did not complete all workload types", index)
		}
	}
	if completed != requests.Load() {
		t.Errorf("mixed workload accounting: completed=%d requests=%d", completed, requests.Load())
	}
	// Release workload fixtures and idle client pools before measuring retained
	// objects. Two GCs also allow sync.Pool's previous generation to be collected.
	downloadPayload, clients, transports = nil, nil, nil
	runtime.GC()
	runtime.GC()
	logSoakResources(t, started, "afterReclaim", requests.Load(), failures.Load(), server.Stats())
	if failures.Load() != 0 || serverFailures.Load() != 0 {
		t.Fatalf("mixed soak failures=%d serverFailures=%d firstFailure=%q", failures.Load(), serverFailures.Load(), firstFailure)
	}
	if stats := server.Stats(); stats.AcceptedConnections != connectionCount || stats.ActiveConnections != 0 || stats.ActiveRequests != 0 || stats.ProtocolErrors != 0 || stats.RejectedRequests != 0 || stats.RejectedConnections != 0 {
		t.Fatalf("mixed soak connections, errors or recovery: %+v", stats)
	}
}

func checkMixedResponse(ctx context.Context, client *http.Client, request *http.Request, kind, downloadBytes int, expectedHash [sha256.Size]byte) error {
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK || response.ProtoMajor != 2 || response.TLS == nil || response.TLS.NegotiatedProtocol != "h2" || response.Header.Get("X-Soak-Request-Id") != request.Header.Get("X-Soak-Request-Id") {
		return fmt.Errorf("unexpected status, protocol or request ID: %s", response.Status)
	}
	if kind != 1 {
		body, err := io.ReadAll(io.LimitReader(response.Body, 1024))
		expected := "healthy"
		if kind == 0 {
			expected = "uploaded"
		}
		if err != nil || string(body) != expected {
			return fmt.Errorf("invalid response %q: %v", body, err)
		}
		return nil
	}
	digest := sha256.New()
	buffer := make([]byte, 64<<10)
	total := 0
	for {
		n, err := response.Body.Read(buffer)
		total += n
		if total > downloadBytes {
			return fmt.Errorf("download exceeds %d bytes", downloadBytes)
		}
		if n > 0 {
			_, _ = digest.Write(buffer[:n])
			if pauseErr := pauseMixedClient(ctx, 20*time.Millisecond); pauseErr != nil {
				return pauseErr
			}
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
	}
	if total != downloadBytes || !bytes.Equal(digest.Sum(nil), expectedHash[:]) {
		return fmt.Errorf("download length or digest mismatch: %d bytes", total)
	}
	return nil
}

type pacedUpload struct {
	ctx       context.Context
	remaining int
}

func (r *pacedUpload) Read(buffer []byte) (int, error) {
	if r.remaining == 0 {
		return 0, io.EOF
	}
	if len(buffer) == 0 {
		return 0, nil
	}
	if err := pauseMixedClient(r.ctx, 2*time.Millisecond); err != nil {
		return 0, err
	}
	n := len(buffer)
	if n > 4096 {
		n = 4096
	}
	if n > r.remaining {
		n = r.remaining
	}
	for index := 0; index < n; index++ {
		buffer[index] = 'x'
	}
	r.remaining -= n
	return n, nil
}

func pauseMixedClient(ctx context.Context, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-timer.C:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}