// perfcheck compares real TCP servers in separate processes. It is a local
// measurement tool, not part of the framework's public CLI or server API.
package main

import (
	"bufio"
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"math/big"
	"math/bits"
	"net"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

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

const expectedBody = `{"code":"OK","message":"success","data":{"status":"ok"},"requestId":"perfCheck"}`

type snapshot struct {
	Address         string `json:"address,omitempty"`
	HeapAlloc       uint64 `json:"heapAlloc"`
	HeapInuse       uint64 `json:"heapInuse"`
	StackInuse      uint64 `json:"stackInuse"`
	TotalAlloc      uint64 `json:"totalAlloc"`
	Mallocs         uint64 `json:"mallocs"`
	NumGC           uint32 `json:"numGc"`
	Goroutines      int    `json:"goroutines"`
	Connections     int64  `json:"connections"`
	Accepted        int64  `json:"accepted"`
	FileDescriptors int    `json:"fileDescriptors,omitempty"`
}

type result struct {
	Adapter                    string   `json:"adapter"`
	Scenario                   string   `json:"scenario"`
	Concurrency                int      `json:"concurrency"`
	Run                        int      `json:"run"`
	Requests                   uint64   `json:"requests,omitempty"`
	Errors                     uint64   `json:"errors"`
	FirstError                 string   `json:"firstError,omitempty"`
	ElapsedSeconds             float64  `json:"elapsedSeconds,omitempty"`
	RequestsPerSecond          float64  `json:"requestsPerSecond,omitempty"`
	P50Millis                  float64  `json:"p50Millis,omitempty"`
	P95Millis                  float64  `json:"p95Millis,omitempty"`
	P99Millis                  float64  `json:"p99Millis,omitempty"`
	ServerBytesPerRequest      float64  `json:"serverBytesPerRequest,omitempty"`
	ServerAllocsPerRequest     float64  `json:"serverAllocsPerRequest,omitempty"`
	SampledPeakHeapAlloc       uint64   `json:"sampledPeakHeapAlloc,omitempty"`
	SampledPeakFileDescriptors int      `json:"sampledPeakFileDescriptors,omitempty"`
	IdleHeapDelta              int64    `json:"idleHeapDelta,omitempty"`
	IdleHeapPerConnection      float64  `json:"idleHeapPerConnection,omitempty"`
	Before                     snapshot `json:"before"`
	During                     snapshot `json:"during"`
	After                      snapshot `json:"after"`
	AfterReclaim               snapshot `json:"afterReclaim"`
}

type report struct {
	CreatedAt  string   `json:"createdAt"`
	GoVersion  string   `json:"goVersion"`
	Platform   string   `json:"platform"`
	GOMAXPROCS int      `json:"gomaxprocs"`
	Duration   string   `json:"duration"`
	TLS        bool     `json:"tls"`
	Notes      []string `json:"notes"`
	Results    []result `json:"results"`
}

type trackedListener struct {
	net.Listener
	active, accepted atomic.Int64
}

type trackedConn struct {
	net.Conn
	once   sync.Once
	active *atomic.Int64
}

func (l *trackedListener) Accept() (net.Conn, error) {
	c, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}
	l.active.Add(1)
	l.accepted.Add(1)
	return &trackedConn{Conn: c, active: &l.active}, nil
}

func (c *trackedConn) Close() error {
	err := c.Conn.Close()
	c.once.Do(func() { c.active.Add(-1) })
	return err
}

func testTLSConfig() (*tls.Config, error) {
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return nil, err
	}
	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 120))
	if err != nil {
		return nil, err
	}
	now := time.Now()
	template := &x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: "localhost"},
		NotBefore:    now.Add(-time.Minute),
		NotAfter:     now.Add(time.Hour),
		KeyUsage:     x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		DNSNames:     []string{"localhost"},
	}
	der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
	if err != nil {
		return nil, err
	}
	return &tls.Config{
		Certificates: []tls.Certificate{{Certificate: [][]byte{der}, PrivateKey: key}},
		MinVersion:   tls.VersionTLS12,
	}, nil
}

func main() {
	serverMode := flag.String("server", "", "internal server subprocess: nethttp or native")
	tlsEnabled := flag.Bool("tls", false, "run the comparison over TLS")
	duration := flag.Duration("duration", 2*time.Second, "measured duration per traffic case")
	concurrencyText := flag.String("concurrency", "1,32,128", "traffic concurrency levels")
	idleText := flag.String("idle", "128,1024", "established idle Keep-Alive connection levels")
	runs := flag.Int("count", 3, "repetitions; adapter order alternates each repetition")
	output := flag.String("output", "", "optional JSON report path")
	flag.Parse()
	var err error
	if *serverMode != "" {
		err = serve(*serverMode, *tlsEnabled)
	} else {
		err = compare(*duration, *concurrencyText, *idleText, *runs, *output, *tlsEnabled)
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func serve(adapter string, tlsEnabled bool) error {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return err
	}
	tracked := &trackedListener{Listener: listener}
	var serveListener net.Listener = tracked
	if tlsEnabled {
		config, configErr := testTLSConfig()
		if configErr != nil {
			listener.Close()
			return configErr
		}
		serveListener = tls.NewListener(tracked, config)
	}
	application := omni.New()
	application.GET("/health", func(c *omni.Context) error {
		return c.OK(map[string]string{"status": "ok"})
	})
	done := make(chan error, 1)
	var shutdown func(context.Context) error
	switch adapter {
	case "native":
		server := native.New(application)
		server.MaxConnections = 8192
		server.MaxConcurrentRequests = 4096
		shutdown = server.Shutdown
		go func() { done <- server.Serve(context.Background(), serveListener) }()
	case "nethttp":
		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(serveListener)
			if errors.Is(err, http.ErrServerClosed) {
				err = nil
			}
			done <- err
		}()
	default:
		listener.Close()
		return fmt.Errorf("unknown server %q", adapter)
	}
	stop := func() error {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := shutdown(ctx); err != nil {
			return err
		}
		select {
		case err := <-done:
			return err
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	encoder := json.NewEncoder(os.Stdout)
	if err := encoder.Encode(snapshot{Address: listener.Addr().String()}); err != nil {
		return err
	}
	decoder := json.NewDecoder(os.Stdin)
	for {
		var command string
		if err := decoder.Decode(&command); err != nil || command == "stop" {
			return stop()
		}
		if command == "gc" {
			runtime.GC()
		}
		var memory runtime.MemStats
		runtime.ReadMemStats(&memory)
		if err := encoder.Encode(snapshot{
			HeapAlloc: memory.HeapAlloc, HeapInuse: memory.HeapInuse,
			StackInuse: memory.StackInuse, TotalAlloc: memory.TotalAlloc,
			Mallocs: memory.Mallocs, NumGC: memory.NumGC,
			Goroutines: runtime.NumGoroutine(), Connections: tracked.active.Load(), Accepted: tracked.accepted.Load(),
			FileDescriptors: openFileDescriptorCount(),
		}); err != nil {
			return err
		}
	}
}

func openFileDescriptorCount() int {
	entries, err := os.ReadDir("/proc/self/fd")
	if err != nil {
		return 0
	}
	return len(entries)
}

type child struct {
	cmd     *exec.Cmd
	input   io.WriteCloser
	encoder *json.Encoder
	decoder *json.Decoder
	address string
	tls     bool
	cancel  context.CancelFunc
	closed  bool
}

func startChild(adapter string) (*child, error) {
	return startChildMode(adapter, false)
}

func startChildMode(adapter string, tlsEnabled bool) (*child, error) {
	executable, err := os.Executable()
	if err != nil {
		return nil, err
	}
	// This bounds even a broken server or protocol read; it is not a load
	// duration. Keep it above the supported 30-minute soak plus drain time.
	ctx, cancel := context.WithTimeout(context.Background(), 35*time.Minute)
	args := []string{"--server", adapter}
	if tlsEnabled {
		args = append(args, "--tls")
	}
	c := &child{cmd: exec.CommandContext(ctx, executable, args...), cancel: cancel, tls: tlsEnabled}
	c.cmd.Stderr = os.Stderr
	if c.input, err = c.cmd.StdinPipe(); err != nil {
		cancel()
		return nil, err
	}
	output, err := c.cmd.StdoutPipe()
	if err != nil {
		cancel()
		c.input.Close()
		return nil, err
	}
	if err := c.cmd.Start(); err != nil {
		cancel()
		c.input.Close()
		output.Close()
		return nil, err
	}
	c.encoder, c.decoder = json.NewEncoder(c.input), json.NewDecoder(output)
	var ready snapshot
	if err := c.decoder.Decode(&ready); err != nil {
		cancel()
		c.cmd.Wait()
		return nil, err
	}
	c.address = ready.Address
	return c, nil
}

func (c *child) snapshot(gc bool) (snapshot, error) {
	command := "snapshot"
	if gc {
		command = "gc"
	}
	var value snapshot
	if err := c.encoder.Encode(command); err != nil {
		return value, err
	}
	err := c.decoder.Decode(&value)
	return value, err
}

func (c *child) close() error {
	if c.closed {
		return nil
	}
	c.closed = true
	defer c.cancel()
	err := c.encoder.Encode("stop")
	c.input.Close()
	return errors.Join(err, c.cmd.Wait())
}

// Histograms use 16 subdivisions per power of two. Percentiles are bucket
// upper bounds (at most about 6.25% rounding), with bounded client memory.
type histogram [1024]uint64

func (h *histogram) add(d time.Duration) {
	n := uint64(max(d.Nanoseconds(), 1))
	exponent := bits.Len64(n) - 1
	base := uint64(1) << exponent
	sub := (n - base) * 16 / base
	h[exponent*16+int(sub)]++
}

func (h *histogram) percentile(count uint64, percent uint64) float64 {
	target := (count*percent + 99) / 100
	if target == 0 {
		return 0
	}
	var seen uint64
	for bucket, n := range h {
		seen += n
		if seen >= target {
			base := uint64(1) << (bucket / 16)
			return float64(base+uint64(bucket%16+1)*base/16) / float64(time.Millisecond)
		}
	}
	return 0
}

func exchange(client *http.Client, address string) error {
	return exchangeMode(client, address, false)
}

func exchangeMode(client *http.Client, address string, tlsEnabled bool) error {
	scheme := "http"
	if tlsEnabled {
		scheme = "https"
	}
	req, err := http.NewRequest(http.MethodGet, scheme+"://"+address+"/health", nil)
	if err != nil {
		return err
	}
	req.Header.Set("X-Request-ID", "perfCheck")
	response, err := client.Do(req)
	if err != nil {
		return err
	}
	body, readErr := io.ReadAll(io.LimitReader(response.Body, 4096))
	closeErr := response.Body.Close()
	if response.StatusCode != 200 || response.Header.Get("X-Request-ID") != "perfCheck" || strings.TrimSpace(string(body)) != expectedBody {
		return fmt.Errorf("invalid response status=%d body=%q", response.StatusCode, body)
	}
	return errors.Join(readErr, closeErr)
}

func traffic(c *child, concurrency int, duration time.Duration) (result, error) {
	r := result{Scenario: "traffic", Concurrency: concurrency}
	transport := &http.Transport{
		MaxIdleConns: concurrency, MaxIdleConnsPerHost: concurrency,
		MaxConnsPerHost: concurrency, IdleConnTimeout: time.Minute,
	}
	if c.tls {
		transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // test certificate is self-signed
	}
	client := &http.Client{Transport: transport, Timeout: 5 * time.Second}
	defer client.CloseIdleConnections()
	// Warm every worker before forcing a baseline GC in the server.
	warmErrors := make(chan error, concurrency)
	for i := 0; i < concurrency; i++ {
		go func() { warmErrors <- exchangeMode(client, c.address, c.tls) }()
	}
	for i := 0; i < concurrency; i++ {
		if err := <-warmErrors; err != nil {
			return r, err
		}
	}
	var err error
	if r.Before, err = c.snapshot(true); err != nil {
		return r, err
	}
	type workerResult struct {
		latencies     histogram
		count, errors uint64
		firstError    string
	}
	workers := make([]workerResult, concurrency)
	start := make(chan struct{})
	var group sync.WaitGroup
	var deadline time.Time
	for i := range workers {
		group.Add(1)
		go func(w *workerResult) {
			defer group.Done()
			<-start
			for time.Now().Before(deadline) {
				began := time.Now()
				err := exchangeMode(client, c.address, c.tls)
				w.latencies.add(time.Since(began))
				w.count++
				if err != nil {
					w.errors++
					if w.firstError == "" {
						w.firstError = err.Error()
					}
				}
			}
		}(&workers[i])
	}
	done := make(chan struct{})
	go func() { group.Wait(); close(done) }()
	begin := time.Now()
	deadline = begin.Add(duration)
	close(start)
	ticker := time.NewTicker(200 * time.Millisecond)
	defer ticker.Stop()
	var sampleErr error
	loading := true
	for loading {
		select {
		case <-done:
			loading = false
		case <-ticker.C:
			s, err := c.snapshot(false)
			if err != nil {
				sampleErr = err
			} else if s.HeapAlloc > r.SampledPeakHeapAlloc {
				r.SampledPeakHeapAlloc = s.HeapAlloc
			}
			if s.FileDescriptors > r.SampledPeakFileDescriptors {
				r.SampledPeakFileDescriptors = s.FileDescriptors
			}
		}
	}
	r.ElapsedSeconds = time.Since(begin).Seconds()
	if sampleErr != nil {
		return r, sampleErr
	}
	var latencies histogram
	for _, w := range workers {
		r.Requests += w.count
		r.Errors += w.errors
		if r.FirstError == "" {
			r.FirstError = w.firstError
		}
		for i, n := range w.latencies {
			latencies[i] += n
		}
	}
	if r.During, err = c.snapshot(false); err != nil {
		return r, err
	}
	r.SampledPeakHeapAlloc = max(r.SampledPeakHeapAlloc, r.During.HeapAlloc)
	r.SampledPeakFileDescriptors = max(r.SampledPeakFileDescriptors, r.During.FileDescriptors)
	if r.Requests == 0 {
		return r, errors.New("no requests completed")
	}
	r.RequestsPerSecond = float64(r.Requests-r.Errors) / r.ElapsedSeconds
	r.P50Millis = latencies.percentile(r.Requests, 50)
	r.P95Millis = latencies.percentile(r.Requests, 95)
	r.P99Millis = latencies.percentile(r.Requests, 99)
	r.ServerBytesPerRequest = float64(r.During.TotalAlloc-r.Before.TotalAlloc) / float64(r.Requests)
	r.ServerAllocsPerRequest = float64(r.During.Mallocs-r.Before.Mallocs) / float64(r.Requests)
	client.CloseIdleConnections()
	r.After, err = drained(c)
	if err == nil {
		r.AfterReclaim, err = c.snapshot(true)
	}
	return r, err
}

func idle(c *child, count int) (result, error) {
	r := result{Scenario: "idle", Concurrency: count}
	var err error
	if r.Before, err = c.snapshot(true); err != nil {
		return r, err
	}
	connections := make([]net.Conn, 0, count)
	defer func() {
		for _, conn := range connections {
			conn.Close()
		}
	}()
	for i := 0; i < count; i++ {
		conn, err := net.DialTimeout("tcp", c.address, 5*time.Second)
		if err != nil {
			return r, err
		}
		if c.tls {
			secure := tls.Client(conn, &tls.Config{InsecureSkipVerify: true, ServerName: "localhost"}) // test certificate is self-signed
			conn = secure
		}
		connections = append(connections, conn)
		conn.SetDeadline(time.Now().Add(5 * time.Second))
		if _, err := io.WriteString(conn, "GET /health HTTP/1.1\r\nHost: localhost\r\nX-Request-ID: perfCheck\r\n\r\n"); err != nil {
			return r, err
		}
		response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodGet})
		if err != nil {
			return r, err
		}
		body, readErr := io.ReadAll(response.Body)
		response.Body.Close()
		if readErr != nil || response.StatusCode != 200 || strings.TrimSpace(string(body)) != expectedBody || response.Close {
			return r, fmt.Errorf("invalid Keep-Alive response: status=%d error=%v", response.StatusCode, readErr)
		}
		conn.SetDeadline(time.Time{})
	}
	if r.During, err = c.snapshot(true); err != nil {
		return r, err
	}
	if r.During.Connections != int64(count) {
		return r, fmt.Errorf("idle connections: want %d, have %d", count, r.During.Connections)
	}
	if r.During.FileDescriptors > 0 && r.Before.FileDescriptors > 0 && r.During.FileDescriptors < r.Before.FileDescriptors+count {
		return r, fmt.Errorf("idle file descriptors: want at least %d, have %d", r.Before.FileDescriptors+count, r.During.FileDescriptors)
	}
	r.IdleHeapDelta = int64(r.During.HeapAlloc) - int64(r.Before.HeapAlloc)
	r.IdleHeapPerConnection = float64(r.IdleHeapDelta) / float64(count)
	for _, conn := range connections {
		conn.Close()
	}
	connections = nil
	r.After, err = drained(c)
	if err == nil {
		r.AfterReclaim, err = c.snapshot(true)
	}
	return r, err
}

func drained(c *child) (snapshot, error) {
	deadline := time.Now().Add(5 * time.Second)
	for {
		s, err := c.snapshot(false)
		if err != nil {
			return s, err
		}
		if s.Connections == 0 {
			return c.snapshot(true)
		}
		if time.Now().After(deadline) {
			return s, fmt.Errorf("%d connections remained after clients closed", s.Connections)
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func levels(text string) ([]int, error) {
	var values []int
	if text == "" {
		return values, nil
	}
	for _, item := range strings.Split(text, ",") {
		n, err := strconv.Atoi(strings.TrimSpace(item))
		if err != nil || n < 1 || n > 4096 {
			return nil, fmt.Errorf("invalid level %q (expected 1..4096)", item)
		}
		values = append(values, n)
	}
	return values, nil
}

func compare(duration time.Duration, concurrencyText, idleText string, count int, output string, tlsEnabled bool) error {
	concurrency, err := levels(concurrencyText)
	if err != nil {
		return err
	}
	idleLevels, err := levels(idleText)
	if err != nil {
		return err
	}
	if count < 1 || duration < 200*time.Millisecond || duration > 30*time.Minute || len(concurrency)+len(idleLevels) == 0 {
		return errors.New("count must be positive, duration 200ms..30m, and at least one scenario selected")
	}
	report := report{
		CreatedAt: time.Now().UTC().Format(time.RFC3339), GoVersion: runtime.Version(),
		Platform: runtime.GOOS + "/" + runtime.GOARCH, GOMAXPROCS: runtime.GOMAXPROCS(0), Duration: duration.String(), TLS: tlsEnabled,
		Notes: []string{
			"Loopback HTTP/1.1 health response, fixed requestId, optional TLS, no access logs or database.",
			"Each case starts a fresh server process; load client runs in the parent process on the same machine.",
			"Closed-loop concurrency; throughput is successful responses / actual elapsed time, including drain.",
			"Latency percentiles include failures and are histogram upper bounds with <=6.25% rounding.",
			"Memory is server-only runtime.MemStats (not RSS or kernel socket memory); sampling interval 200ms.",
			"FileDescriptors is read from /proc/self/fd on Linux and omitted on other platforms.",
			"Idle memory measured after a forced GC; delta includes heap growth amortized across connections, excluding stacks.",
			"After is after closing clients and one forced GC; afterReclaim is a second GC to observe pooled-cache release.",
			"Native limits: 8192 connections / 4096 active requests. These cases do not benchmark overload rejection.",
		},
	}
	var caseErrors []error
	for run := 1; run <= count; run++ {
		adapters := []string{"nethttp", "native"}
		if run%2 == 0 {
			adapters[0], adapters[1] = adapters[1], adapters[0]
		}
		for _, scenario := range []struct {
			name   string
			levels []int
		}{{"traffic", concurrency}, {"idle", idleLevels}} {
			for _, level := range scenario.levels {
				for _, adapter := range adapters {
					c, err := startChildMode(adapter, tlsEnabled)
					if err != nil {
						return err
					}
					var value result
					if scenario.name == "traffic" {
						value, err = traffic(c, level, duration)
					} else {
						value, err = idle(c, level)
					}
					err = errors.Join(err, c.close())
					value.Adapter, value.Run = adapter, run
					if err != nil {
						value.Errors++
						value.FirstError = err.Error()
					}
					report.Results = append(report.Results, value)
					if value.Errors > 0 {
						caseErrors = append(caseErrors, fmt.Errorf("%s %s/%d run %d: %s", adapter, scenario.name, level, run, value.FirstError))
					}
					if scenario.name == "traffic" {
						fmt.Printf("%s concurrency=%d run=%d: %.0f req/s p95=%.3fms p99=%.3fms server=%.0f B/req errors=%d\n", adapter, level, run, value.RequestsPerSecond, value.P95Millis, value.P99Millis, value.ServerBytesPerRequest, value.Errors)
					} else {
						fmt.Printf("%s idle=%d run=%d: server heap delta=%.2f MiB (%.0f B/conn), goroutines=%d->%d->%d, fds=%d->%d->%d errors=%d\n", adapter, level, run, float64(value.IdleHeapDelta)/(1<<20), value.IdleHeapPerConnection, value.Before.Goroutines, value.During.Goroutines, value.After.Goroutines, value.Before.FileDescriptors, value.During.FileDescriptors, value.After.FileDescriptors, value.Errors)
					}
				}
			}
		}
	}
	if output != "" {
		data, err := json.MarshalIndent(report, "", "  ")
		if err != nil {
			return err
		}
		if err := os.WriteFile(output, append(data, '\n'), 0644); err != nil {
			return err
		}
		fmt.Printf("Report: %s\n", output)
	}
	return errors.Join(caseErrors...)
}