// Package loadgen runs bounded, reproducible HTTP request workloads. It is a
// capacity-evidence tool and is deliberately separate from service processes.
package loadgen

import (
	"bufio"
	"bytes"
	"context"
	"encoding/base64"
	"errors"
	"io"
	"math/bits"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

var ErrInvalidOptions = errors.New("invalid load generator options")

const maximumResponseDrainBytes = 64 << 10

type Scenario string

const (
	ScenarioHTTP    Scenario = "http"
	ScenarioConnect Scenario = "connect"
)

// Options bounds one HTTP workload. Requests selects a fixed-size run;
// otherwise Duration selects a time-boxed run and Rate caps its start rate.
type Options struct {
	Scenario       Scenario
	TargetURL      string
	ProxyURL       string
	Method         string
	Headers        http.Header
	RequestBody    []byte
	Requests       int
	Duration       time.Duration
	Rate           int
	Concurrency    int
	RequestTimeout time.Duration
	TunnelHold     time.Duration
}

type Report struct {
	Scenario           Scenario
	StartedAt          time.Time
	Duration           time.Duration
	Requests           uint64
	Completed          uint64
	Succeeded          uint64
	Failed             uint64
	Status4xx          uint64
	Status5xx          uint64
	TimeoutErrors      uint64
	RequestErrors      uint64
	TunnelsEstablished uint64
	Throughput         float64
	Latency            LatencyReport
	Runtime            RuntimeReport
}

// LatencyReport uses logarithmic microsecond buckets. Percentiles are upper
// bounds so the collector remains fixed-size even under high QPS.
type LatencyReport struct {
	Samples           uint64
	P50UpperBound     time.Duration
	P95UpperBound     time.Duration
	P99UpperBound     time.Duration
	MaximumUpperBound time.Duration
}

type RuntimeReport struct {
	NumCPU     int
	Goroutines int
	AllocBytes uint64
	SysBytes   uint64
	GCCount    uint32
}

type counters struct {
	requests           atomic.Uint64
	completed          atomic.Uint64
	succeeded          atomic.Uint64
	failed             atomic.Uint64
	status4xx          atomic.Uint64
	status5xx          atomic.Uint64
	timeoutErrors      atomic.Uint64
	requestErrors      atomic.Uint64
	tunnelsEstablished atomic.Uint64
	latency            latencyHistogram
}

type latencyHistogram struct {
	buckets [64]atomic.Uint64
}

func Run(ctx context.Context, options Options) (Report, error) {
	if ctx == nil {
		return Report{}, ErrInvalidOptions
	}
	normalized, err := normalizeOptions(options)
	if err != nil {
		return Report{}, err
	}
	stats := &counters{}
	var execute func(context.Context)
	if normalized.Scenario == ScenarioHTTP {
		client, clientErr := newHTTPClient(normalized)
		if clientErr != nil {
			return Report{}, clientErr
		}
		defer client.CloseIdleConnections()
		execute = func(requestCtx context.Context) { executeHTTP(requestCtx, client, normalized, stats) }
	} else {
		execute = func(requestCtx context.Context) { executeCONNECT(requestCtx, normalized, stats) }
	}

	started := time.Now()
	if normalized.Requests > 0 {
		runFixed(ctx, normalized, execute)
		if err := ctx.Err(); err != nil {
			return report(started, normalized.Scenario, stats), err
		}
	} else {
		runForDuration(ctx, normalized, execute)
	}
	return report(started, normalized.Scenario, stats), nil
}

func normalizeOptions(options Options) (Options, error) {
	normalized := options
	normalized.Scenario = Scenario(strings.ToLower(strings.TrimSpace(string(normalized.Scenario))))
	normalized.Method = strings.ToUpper(strings.TrimSpace(normalized.Method))
	normalized.Headers = normalized.Headers.Clone()
	normalized.RequestBody = bytes.Clone(normalized.RequestBody)
	if normalized.Scenario == "" {
		normalized.Scenario = ScenarioHTTP
	}
	if normalized.Method == "" {
		normalized.Method = http.MethodGet
	}
	if strings.TrimSpace(normalized.TargetURL) != normalized.TargetURL || normalized.TargetURL == "" ||
		normalized.Concurrency <= 0 || normalized.RequestTimeout <= 0 || normalized.Requests < 0 || normalized.Duration < 0 ||
		normalized.Rate < 0 || (normalized.Requests == 0 && normalized.Duration <= 0) ||
		(normalized.Requests > 0 && (normalized.Duration != 0 || normalized.Rate != 0)) {
		return Options{}, ErrInvalidOptions
	}
	target, err := url.Parse(normalized.TargetURL)
	if err != nil || target.Scheme == "" || target.Host == "" || (target.Scheme != "http" && target.Scheme != "https") {
		return Options{}, ErrInvalidOptions
	}
	if normalized.Scenario != ScenarioHTTP && normalized.Scenario != ScenarioConnect {
		return Options{}, ErrInvalidOptions
	}
	if normalized.Scenario == ScenarioHTTP {
		if normalized.TunnelHold != 0 {
			return Options{}, ErrInvalidOptions
		}
		if _, err := parseProxyURL(normalized.ProxyURL, false); err != nil {
			return Options{}, err
		}
		return normalized, nil
	}
	if normalized.ProxyURL == "" || normalized.TunnelHold <= 0 || len(normalized.RequestBody) != 0 || normalized.Method != http.MethodGet {
		return Options{}, ErrInvalidOptions
	}
	if _, err := parseProxyURL(normalized.ProxyURL, true); err != nil {
		return Options{}, err
	}
	return normalized, nil
}

func newHTTPClient(options Options) (*http.Client, error) {
	parsed, err := parseProxyURL(options.ProxyURL, false)
	if err != nil {
		return nil, err
	}
	var proxy func(*http.Request) (*url.URL, error)
	if parsed != nil {
		proxy = http.ProxyURL(parsed)
	}
	transport := &http.Transport{
		Proxy:                 proxy,
		ForceAttemptHTTP2:     false,
		MaxConnsPerHost:       options.Concurrency,
		MaxIdleConns:          options.Concurrency,
		MaxIdleConnsPerHost:   options.Concurrency,
		IdleConnTimeout:       30 * time.Second,
		TLSHandshakeTimeout:   options.RequestTimeout,
		ResponseHeaderTimeout: options.RequestTimeout,
	}
	return &http.Client{Transport: transport}, nil
}

func parseProxyURL(value string, plainHTTPOnly bool) (*url.URL, error) {
	if value == "" {
		return nil, nil
	}
	if strings.TrimSpace(value) != value {
		return nil, ErrInvalidOptions
	}
	parsed, err := url.Parse(value)
	if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") ||
		(plainHTTPOnly && parsed.Scheme != "http") {
		return nil, ErrInvalidOptions
	}
	return parsed, nil
}

func runFixed(ctx context.Context, options Options, execute func(context.Context)) {
	jobs := make(chan struct{}, min(options.Concurrency, options.Requests))
	var workers sync.WaitGroup
	for range options.Concurrency {
		workers.Add(1)
		go func() {
			defer workers.Done()
			for range jobs {
				if ctx.Err() != nil {
					return
				}
				execute(ctx)
			}
		}()
	}
	for range options.Requests {
		select {
		case jobs <- struct{}{}:
		case <-ctx.Done():
			close(jobs)
			workers.Wait()
			return
		}
	}
	close(jobs)
	workers.Wait()
}

func runForDuration(ctx context.Context, options Options, execute func(context.Context)) {
	workloadContext, cancel := context.WithTimeout(ctx, options.Duration)
	defer cancel()
	if options.Rate == 0 {
		var workers sync.WaitGroup
		for range options.Concurrency {
			workers.Add(1)
			go func() {
				defer workers.Done()
				for workloadContext.Err() == nil {
					execute(workloadContext)
				}
			}()
		}
		workers.Wait()
		return
	}

	jobs := make(chan struct{}, options.Concurrency)
	var workers sync.WaitGroup
	for range options.Concurrency {
		workers.Add(1)
		go func() {
			defer workers.Done()
			for {
				select {
				case <-workloadContext.Done():
					return
				case <-jobs:
					execute(workloadContext)
				}
			}
		}()
	}

	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()
	last := time.Now()
	credit := 0.0
	for {
		select {
		case <-workloadContext.Done():
			workers.Wait()
			return
		case now := <-ticker.C:
			credit += float64(options.Rate) * now.Sub(last).Seconds()
			last = now
			tokens := int(credit)
			credit -= float64(tokens)
			if tokens > options.Concurrency {
				tokens = options.Concurrency
				credit = 0
			}
			for range tokens {
				select {
				case jobs <- struct{}{}:
				case <-workloadContext.Done():
					workers.Wait()
					return
				}
			}
		}
	}
}

func executeHTTP(ctx context.Context, client *http.Client, options Options, stats *counters) {
	stats.requests.Add(1)
	started := time.Now()
	requestContext, cancel := context.WithTimeout(ctx, options.RequestTimeout)
	defer cancel()
	request, err := http.NewRequestWithContext(
		requestContext,
		options.Method,
		options.TargetURL,
		bytes.NewReader(options.RequestBody),
	)
	if err != nil {
		stats.failed.Add(1)
		stats.requestErrors.Add(1)
		stats.completed.Add(1)
		stats.latency.Record(time.Since(started))
		return
	}
	request.Header = options.Headers.Clone()
	response, err := client.Do(request)
	latency := time.Since(started)
	stats.latency.Record(latency)
	stats.completed.Add(1)
	if err != nil {
		stats.failed.Add(1)
		if errors.Is(err, context.DeadlineExceeded) || errors.Is(requestContext.Err(), context.DeadlineExceeded) {
			stats.timeoutErrors.Add(1)
		} else {
			stats.requestErrors.Add(1)
		}
		return
	}
	_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, maximumResponseDrainBytes))
	_ = response.Body.Close()
	switch {
	case response.StatusCode >= http.StatusOK && response.StatusCode < http.StatusMultipleChoices:
		stats.succeeded.Add(1)
	case response.StatusCode >= http.StatusInternalServerError:
		stats.failed.Add(1)
		stats.status5xx.Add(1)
	case response.StatusCode >= http.StatusBadRequest:
		stats.failed.Add(1)
		stats.status4xx.Add(1)
	default:
		stats.failed.Add(1)
		stats.requestErrors.Add(1)
	}
}

func executeCONNECT(ctx context.Context, options Options, stats *counters) {
	stats.requests.Add(1)
	started := time.Now()
	proxy, err := parseProxyURL(options.ProxyURL, true)
	if err != nil {
		recordRequestError(started, stats, false)
		return
	}
	target, err := url.Parse(options.TargetURL)
	if err != nil {
		recordRequestError(started, stats, false)
		return
	}
	authority := target.Host
	if target.Port() == "" {
		if target.Scheme == "https" {
			authority = net.JoinHostPort(target.Hostname(), "443")
		} else {
			authority = net.JoinHostPort(target.Hostname(), "80")
		}
	}

	requestContext, cancel := context.WithTimeout(ctx, options.RequestTimeout)
	defer cancel()
	connection, err := (&net.Dialer{}).DialContext(requestContext, "tcp", proxy.Host)
	if err != nil {
		recordRequestError(started, stats, errors.Is(requestContext.Err(), context.DeadlineExceeded))
		return
	}
	defer connection.Close()
	if err := connection.SetDeadline(time.Now().Add(options.RequestTimeout)); err != nil {
		recordRequestError(started, stats, false)
		return
	}
	if err := writeCONNECTRequest(connection, authority, options.Headers, proxy.User); err != nil {
		recordRequestError(started, stats, errors.Is(requestContext.Err(), context.DeadlineExceeded))
		return
	}
	response, err := http.ReadResponse(bufio.NewReader(connection), &http.Request{Method: http.MethodConnect})
	latency := time.Since(started)
	stats.latency.Record(latency)
	if err != nil {
		stats.failed.Add(1)
		stats.completed.Add(1)
		if errors.Is(requestContext.Err(), context.DeadlineExceeded) {
			stats.timeoutErrors.Add(1)
		} else {
			stats.requestErrors.Add(1)
		}
		return
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		stats.failed.Add(1)
		stats.completed.Add(1)
		switch {
		case response.StatusCode >= http.StatusInternalServerError:
			stats.status5xx.Add(1)
		case response.StatusCode >= http.StatusBadRequest:
			stats.status4xx.Add(1)
		default:
			stats.requestErrors.Add(1)
		}
		return
	}
	if err := connection.SetDeadline(time.Time{}); err != nil {
		stats.failed.Add(1)
		stats.requestErrors.Add(1)
		stats.completed.Add(1)
		return
	}
	stats.succeeded.Add(1)
	stats.tunnelsEstablished.Add(1)
	waitForTunnelHold(ctx, options.TunnelHold)
	stats.completed.Add(1)
}

func writeCONNECTRequest(connection net.Conn, authority string, headers http.Header, user *url.Userinfo) error {
	if connection == nil || authority == "" {
		return ErrInvalidOptions
	}
	requestHeaders := headers.Clone()
	requestHeaders.Del("Host")
	if user != nil && requestHeaders.Get("Proxy-Authorization") == "" {
		password, _ := user.Password()
		encoded := base64.StdEncoding.EncodeToString([]byte(user.Username() + ":" + password))
		requestHeaders.Set("Proxy-Authorization", "Basic "+encoded)
	}
	var payload bytes.Buffer
	_, _ = payload.WriteString("CONNECT " + authority + " HTTP/1.1\r\nHost: " + authority + "\r\n")
	if err := requestHeaders.Write(&payload); err != nil {
		return err
	}
	_, _ = payload.WriteString("\r\n")
	_, err := io.Copy(connection, &payload)
	return err
}

func waitForTunnelHold(ctx context.Context, hold time.Duration) {
	timer := time.NewTimer(hold)
	defer timer.Stop()
	select {
	case <-ctx.Done():
	case <-timer.C:
	}
}

func recordRequestError(started time.Time, stats *counters, timeout bool) {
	stats.latency.Record(time.Since(started))
	stats.failed.Add(1)
	stats.completed.Add(1)
	if timeout {
		stats.timeoutErrors.Add(1)
		return
	}
	stats.requestErrors.Add(1)
}

func report(started time.Time, scenario Scenario, stats *counters) Report {
	duration := time.Since(started)
	completed := stats.completed.Load()
	var memory runtime.MemStats
	runtime.ReadMemStats(&memory)
	result := Report{
		Scenario: scenario, StartedAt: started.UTC(), Duration: duration, Requests: stats.requests.Load(), Completed: completed,
		Succeeded: stats.succeeded.Load(), Failed: stats.failed.Load(), Status4xx: stats.status4xx.Load(),
		Status5xx: stats.status5xx.Load(), TimeoutErrors: stats.timeoutErrors.Load(), RequestErrors: stats.requestErrors.Load(),
		TunnelsEstablished: stats.tunnelsEstablished.Load(),
		Latency:            stats.latency.Report(completed),
		Runtime:            RuntimeReport{NumCPU: runtime.NumCPU(), Goroutines: runtime.NumGoroutine(), AllocBytes: memory.Alloc, SysBytes: memory.Sys, GCCount: memory.NumGC},
	}
	if duration > 0 {
		result.Throughput = float64(completed) / duration.Seconds()
	}
	return result
}

func (histogram *latencyHistogram) Record(latency time.Duration) {
	microseconds := latency.Microseconds()
	if microseconds < 1 {
		microseconds = 1
	}
	index := bits.Len64(uint64(microseconds - 1))
	if index >= len(histogram.buckets) {
		index = len(histogram.buckets) - 1
	}
	histogram.buckets[index].Add(1)
}

func (histogram *latencyHistogram) Report(samples uint64) LatencyReport {
	result := LatencyReport{Samples: samples}
	if samples == 0 {
		return result
	}
	result.P50UpperBound = histogram.percentile(samples, 50)
	result.P95UpperBound = histogram.percentile(samples, 95)
	result.P99UpperBound = histogram.percentile(samples, 99)
	result.MaximumUpperBound = histogram.percentile(samples, 100)
	return result
}

func (histogram *latencyHistogram) percentile(samples uint64, percentile uint64) time.Duration {
	target := (samples*percentile + 99) / 100
	seen := uint64(0)
	for index := range histogram.buckets {
		seen += histogram.buckets[index].Load()
		if seen >= target {
			return time.Duration(uint64(1)<<index) * time.Microsecond
		}
	}
	return time.Duration(1<<63 - 1)
}