package transport_test

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"sync"
	"testing"
	"time"

	omni "git.yuntuo.hk/k2safe/go-omni"
	"git.yuntuo.hk/k2safe/go-omni/transport/native"
)

// Both servers use the same application, client, request validation and fixed
// concurrency. B/op includes client AND server allocations in this process;
// ns/op is aggregate throughput cost, not individual request latency.
func BenchmarkTransportTCPApplicationRequest(b *testing.B) {
	for _, adapter := range []string{"nethttp", "native"} {
		for _, concurrency := range []int{1, 32, 128} {
			b.Run(fmt.Sprintf("%s/concurrency%d", adapter, concurrency), func(b *testing.B) {
				benchmarkTCP(b, adapter, concurrency)
			})
		}
	}
}

func benchmarkTCP(b *testing.B, adapter string, concurrency int) {
	application := omni.New()
	application.GET("/health", func(ctx *omni.Context) error {
		return ctx.OK(map[string]string{"status": "ok"})
	})
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		b.Fatal(err)
	}
	done := make(chan error, 1)
	var shutdown func(context.Context) error
	if adapter == "native" {
		server := native.New(application)
		server.MaxConnections = 1024
		server.MaxConcurrentRequests = 1024
		shutdown = server.Shutdown
		go func() { done <- server.Serve(context.Background(), listener) }()
	} else {
		server := &http.Server{
			Handler: application, ReadHeaderTimeout: 5 * time.Second,
			ReadTimeout: 30 * time.Second, WriteTimeout: 30 * time.Second,
			IdleTimeout: time.Minute, MaxHeaderBytes: 1 << 20,
		}
		shutdown = server.Shutdown
		go func() {
			err := server.Serve(listener)
			if err == http.ErrServerClosed {
				err = nil
			}
			done <- err
		}()
	}
	client := &http.Client{Transport: &http.Transport{
		MaxIdleConns:        1024,
		MaxIdleConnsPerHost: 1024,
		MaxConnsPerHost:     concurrency,
		IdleConnTimeout:     time.Minute,
	}, Timeout: 5 * time.Second}
	b.Cleanup(func() {
		client.CloseIdleConnections()
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := shutdown(ctx); err != nil {
			b.Error(err)
		}
		select {
		case err := <-done:
			if err != nil {
				b.Error(err)
			}
		case <-ctx.Done():
			b.Error("server did not stop")
		}
	})
	endpoint := "http://" + listener.Addr().String() + "/health"
	request := func() error {
		response, err := client.Get(endpoint)
		if err != nil {
			return err
		}
		_, readErr := io.Copy(io.Discard, response.Body)
		closeErr := response.Body.Close()
		if response.StatusCode != http.StatusOK || response.Header.Get("X-Request-ID") == "" {
			return fmt.Errorf("unexpected response: status=%d", response.StatusCode)
		}
		if readErr != nil {
			return readErr
		}
		return closeErr
	}
	if err := request(); err != nil {
		b.Fatal(err)
	}
	errors := make(chan error, concurrency)
	var workers sync.WaitGroup
	b.ReportAllocs()
	b.ResetTimer()
	for worker := 0; worker < concurrency; worker++ {
		workers.Add(1)
		go func(worker int) {
			defer workers.Done()
			for index := worker; index < b.N; index += concurrency {
				if err := request(); err != nil {
					errors <- err
					return
				}
			}
		}(worker)
	}
	workers.Wait()
	b.StopTimer()
	close(errors)
	for err := range errors {
		b.Error(err)
	}
}