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

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

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

const maximumResponseDrainBytes = 64 << 10

// 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 {
	TargetURL      string
	ProxyURL       string
	Method         string
	Headers        http.Header
	Requests       int
	Duration       time.Duration
	Rate           int
	Concurrency    int
	RequestTimeout time.Duration
}

type Report struct {
	StartedAt     time.Time
	Duration      time.Duration
	Requests      uint64
	Completed     uint64
	Succeeded     uint64
	Failed        uint64
	Status4xx     uint64
	Status5xx     uint64
	TimeoutErrors uint64
	RequestErrors 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
	latency       latencyHistogram
}

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

func Run(ctx context.Context, options Options) (Report, error) {
	if ctx == nil {
		return Report{}, ErrInvalidOptions
	}
	client, normalized, err := newClient(options)
	if err != nil {
		return Report{}, err
	}
	defer client.CloseIdleConnections()

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

func newClient(options Options) (*http.Client, Options, error) {
	normalized := options
	normalized.Method = strings.ToUpper(strings.TrimSpace(normalized.Method))
	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 nil, Options{}, ErrInvalidOptions
	}
	target, err := url.Parse(normalized.TargetURL)
	if err != nil || target.Scheme == "" || target.Host == "" || (target.Scheme != "http" && target.Scheme != "https") {
		return nil, Options{}, ErrInvalidOptions
	}
	var proxy func(*http.Request) (*url.URL, error)
	if normalized.ProxyURL != "" {
		if strings.TrimSpace(normalized.ProxyURL) != normalized.ProxyURL {
			return nil, Options{}, ErrInvalidOptions
		}
		parsed, err := url.Parse(normalized.ProxyURL)
		if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") {
			return nil, Options{}, ErrInvalidOptions
		}
		proxy = http.ProxyURL(parsed)
	}
	transport := &http.Transport{
		Proxy:                 proxy,
		ForceAttemptHTTP2:     false,
		MaxConnsPerHost:       normalized.Concurrency,
		MaxIdleConns:          normalized.Concurrency,
		MaxIdleConnsPerHost:   normalized.Concurrency,
		IdleConnTimeout:       30 * time.Second,
		TLSHandshakeTimeout:   normalized.RequestTimeout,
		ResponseHeaderTimeout: normalized.RequestTimeout,
	}
	return &http.Client{Transport: transport}, normalized, nil
}

func runFixed(ctx context.Context, client *http.Client, options Options, stats *counters) {
	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
				}
				executeOne(ctx, client, options, stats)
			}
		}()
	}
	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, client *http.Client, options Options, stats *counters) {
	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 {
					executeOne(workloadContext, client, options, stats)
				}
			}()
		}
		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:
					executeOne(workloadContext, client, options, stats)
				}
			}
		}()
	}

	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 executeOne(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, nil)
	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 report(started time.Time, stats *counters) Report {
	duration := time.Since(started)
	completed := stats.completed.Load()
	var memory runtime.MemStats
	runtime.ReadMemStats(&memory)
	result := Report{
		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(),
		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)
}