package main

import (
	"context"
	"io"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"time"
)

const (
	authAddr = "127.0.0.1:16060"                 // auth upstream listener
	svcAddr  = "127.0.0.1:18080"                 // service upstream listener
	target   = "http://127.0.0.1:18888/svc/ping" // target URL (through Caddy)
	workers  = 300                               // concurrent load workers
	dur      = 30 * time.Second                  // load duration
)

func main() {
	var bugged, authHits, svcHits atomic.Int64

	// mux with expeced verify endpoint
	authMux := http.NewServeMux()
	authMux.HandleFunc("/verify", func(w http.ResponseWriter, r *http.Request) {
		authHits.Add(1)
		w.WriteHeader(200)
	})

	// mux with unexpected verify endpoint (should never be hit)
	svcMux := http.NewServeMux()
	svcMux.HandleFunc("/verify", func(w http.ResponseWriter, r *http.Request) {
		svcHits.Add(1)
		n := bugged.Add(1)
		slog.Warn("BUG TRIGGERED",
			slog.Int64("count", n),
			slog.String("method", r.Method),
			slog.String("url", r.URL.String()),
			slog.String("remote", r.RemoteAddr))
		// this should not exist, so return 404 to make it obvious in the report
		w.WriteHeader(404)
	})
	// and real endpoint we actually want to hit
	svcMux.HandleFunc("/ping", func(w http.ResponseWriter, r *http.Request) {
		svcHits.Add(1)
		w.WriteHeader(200)
	})

	go func() {
		if err := http.ListenAndServe(authAddr, authMux); err != nil {
			slog.Error("auth upstream failed", slog.String("addr", authAddr), slog.Any("err", err))
			os.Exit(1)
		}
	}()
	go func() {
		if err := http.ListenAndServe(svcAddr, svcMux); err != nil {
			slog.Error("svc upstream failed", slog.String("addr", svcAddr), slog.Any("err", err))
			os.Exit(1)
		}
	}()

	// just wait for internal servers to be ready
	time.Sleep(time.Second)
	slog.Info("upstreams up", slog.String("auth", authAddr), slog.String("svc", svcAddr))

	client := &http.Client{Transport: &http.Transport{MaxIdleConnsPerHost: 100}}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()
	ctx, cancel := context.WithTimeout(ctx, dur)
	defer cancel()

	var codes [600]atomic.Int64
	var errs, total atomic.Int64
	var wg sync.WaitGroup
	slog.Info("load starting", slog.Int("workers", workers), slog.Duration("duration", dur), slog.String("target", target))
	start := time.Now()
	for range workers {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				select {
				case <-ctx.Done():
					return
				default:
				}
				total.Add(1)
				resp, err := client.Get(target)
				if err != nil {
					errs.Add(1)
					continue
				}
				io.Copy(io.Discard, resp.Body)
				resp.Body.Close()
				if resp.StatusCode < 600 {
					codes[resp.StatusCode].Add(1)
				}
			}
		}()
	}
	wg.Wait()

	elapsed := time.Since(start)
	if ctx.Err() == context.Canceled {
		slog.Info("interrupted; printing report", slog.Duration("elapsed", elapsed.Round(time.Millisecond)))
	}

	for i := range codes {
		if v := codes[i].Load(); v > 0 {
			slog.Info("status count", slog.Int("status", i), slog.Int64("count", v))
		}
	}

	slog.Info("load complete",
		slog.Int64("errors", errs.Load()),
		slog.Int64("total", total.Load()),
		slog.Float64("req_per_s", float64(total.Load())/elapsed.Seconds()),
		slog.Duration("elapsed", elapsed.Round(time.Millisecond)))
	slog.Info("counters",
		slog.Int64("bugged", bugged.Load()),
		slog.Int64("auth_hits", authHits.Load()),
		slog.Int64("svc_hits", svcHits.Load()))
	if bugged.Load() > 0 {
		slog.Error("RESULT: BUG REPRODUCED", slog.Int64("bugged", bugged.Load()))
	} else {
		slog.Info("RESULT: clean (no bug observed)")
	}
}