// 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"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"math/bits"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

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

const (
	maximumResponseDrainBytes   = 64 << 10
	maximumExtractResponseBytes = 256 << 10
)

type Scenario string

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

// 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
	ExtractCount       int
	ExtractFulfillment string
}

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
	ExtractResponsesValidated uint64
	ExtractReturned           uint64
	ExtractValidationFailures uint64
	Throughput                float64
	RateStartsGenerated       uint64
	RateStartsDropped         uint64
	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
	extractResponsesValidated atomic.Uint64
	extractReturned           atomic.Uint64
	extractValidationFailures atomic.Uint64
	rateStartsGenerated       atomic.Uint64
	rateStartsDropped         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 || normalized.Scenario == ScenarioExtract {
		client, clientErr := newHTTPClient(normalized)
		if clientErr != nil {
			return Report{}, clientErr
		}
		defer client.CloseIdleConnections()
		if normalized.Scenario == ScenarioHTTP {
			execute = func(requestCtx context.Context) { executeHTTP(requestCtx, client, normalized, stats) }
		} else {
			workloadID, workloadErr := newExtractWorkloadID()
			if workloadErr != nil {
				return Report{}, workloadErr
			}
			workload := &extractWorkload{id: workloadID, options: normalized, stats: stats}
			execute = func(requestCtx context.Context) { workload.execute(requestCtx, client) }
		}
	} 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, stats)
	}
	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 == "" {
		if normalized.Scenario == ScenarioExtract {
			normalized.Method = http.MethodPost
		} else {
			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 && normalized.Scenario != ScenarioExtract {
		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.Scenario == ScenarioExtract {
		if normalized.ProxyURL != "" || normalized.TunnelHold != 0 || len(normalized.RequestBody) != 0 ||
			normalized.Method != http.MethodPost || normalized.ExtractCount <= 0 || normalized.ExtractCount > 64 ||
			(normalized.ExtractFulfillment != "" && strings.TrimSpace(normalized.ExtractFulfillment) != normalized.ExtractFulfillment) {
			return Options{}, ErrInvalidOptions
		}
		if normalized.ExtractFulfillment == "" {
			normalized.ExtractFulfillment = "partial"
		}
		if normalized.ExtractFulfillment != "partial" && normalized.ExtractFulfillment != "allOrNothing" {
			return Options{}, ErrInvalidOptions
		}
		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), 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 {
					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():
			waitForRateWorkers(&workers, jobs, stats)
			return
		case now := <-ticker.C:
			credit += float64(options.Rate) * now.Sub(last).Seconds()
			last = now
			tokens := int(credit)
			credit -= float64(tokens)
			if tokens <= 0 {
				continue
			}
			stats.rateStartsGenerated.Add(uint64(tokens))
			if tokens > options.Concurrency {
				stats.rateStartsDropped.Add(uint64(tokens - options.Concurrency))
				tokens = options.Concurrency
				credit = 0
			}
			for index := 0; index < tokens; index++ {
				select {
				case jobs <- struct{}{}:
				case <-workloadContext.Done():
					stats.rateStartsDropped.Add(uint64(tokens - index))
					waitForRateWorkers(&workers, jobs, stats)
					return
				}
			}
		}
	}
}

func waitForRateWorkers(workers *sync.WaitGroup, jobs <-chan struct{}, stats *counters) {
	workers.Wait()
	stats.rateStartsDropped.Add(uint64(len(jobs)))
}

type requestBuilder func(context.Context) (*http.Request, responseValidator, error)

type responseValidator func(*http.Response) error

func executeHTTP(ctx context.Context, client *http.Client, options Options, stats *counters) {
	executeHTTPRequest(ctx, client, options.RequestTimeout, stats, func(requestContext context.Context) (*http.Request, responseValidator, error) {
		request, err := http.NewRequestWithContext(
			requestContext,
			options.Method,
			options.TargetURL,
			bytes.NewReader(options.RequestBody),
		)
		if err != nil {
			return nil, nil, err
		}
		request.Header = options.Headers.Clone()
		return request, discardResponse, nil
	})
}

func executeHTTPRequest(ctx context.Context, client *http.Client, timeout time.Duration, stats *counters, build requestBuilder) {
	stats.requests.Add(1)
	started := time.Now()
	requestContext, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	request, validate, err := build(requestContext)
	if err != nil {
		recordRequestError(started, stats, false)
		return
	}
	response, err := client.Do(request)
	latency := time.Since(started)
	stats.latency.Record(latency)
	if err != nil {
		stats.failed.Add(1)
		stats.completed.Add(1)
		if errors.Is(err, context.DeadlineExceeded) || errors.Is(requestContext.Err(), context.DeadlineExceeded) {
			stats.timeoutErrors.Add(1)
		} else {
			stats.requestErrors.Add(1)
		}
		return
	}
	defer response.Body.Close()
	switch {
	case response.StatusCode >= http.StatusOK && response.StatusCode < http.StatusMultipleChoices:
		if validate == nil {
			validate = discardResponse
		}
		if err := validate(response); err != nil {
			stats.failed.Add(1)
			stats.requestErrors.Add(1)
			stats.completed.Add(1)
			return
		}
		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)
	}
	stats.completed.Add(1)
}

func discardResponse(response *http.Response) error {
	if response == nil || response.Body == nil {
		return ErrInvalidOptions
	}
	_, err := io.Copy(io.Discard, io.LimitReader(response.Body, maximumResponseDrainBytes))
	return err
}

type extractWorkload struct {
	id      string
	options Options
	stats   *counters
	next    atomic.Uint64
}

func newExtractWorkloadID() (string, error) {
	var value [12]byte
	if _, err := rand.Read(value[:]); err != nil {
		return "", err
	}
	return "loadgen-" + hex.EncodeToString(value[:]), nil
}

func (workload *extractWorkload) execute(ctx context.Context, client *http.Client) {
	sequence := workload.next.Add(1)
	requestID := workload.id + "-" + strconv.FormatUint(sequence, 10)
	executeHTTPRequest(ctx, client, workload.options.RequestTimeout, workload.stats, func(requestContext context.Context) (*http.Request, responseValidator, error) {
		body, err := json.Marshal(struct {
			Count       int    `json:"count"`
			Fulfillment string `json:"fulfillment"`
		}{Count: workload.options.ExtractCount, Fulfillment: workload.options.ExtractFulfillment})
		if err != nil {
			return nil, nil, err
		}
		request, err := http.NewRequestWithContext(requestContext, http.MethodPost, workload.options.TargetURL, bytes.NewReader(body))
		if err != nil {
			return nil, nil, err
		}
		request.Header = workload.options.Headers.Clone()
		if request.Header == nil {
			request.Header = make(http.Header)
		}
		request.Header.Set("Content-Type", "application/json")
		request.Header.Set("X-Request-ID", requestID)
		request.Header.Set("Idempotency-Key", requestID)
		return request, func(response *http.Response) error {
			return validateExtractResponse(response, requestID, workload.options, workload.stats)
		}, nil
	})
}

func validateExtractResponse(response *http.Response, requestID string, options Options, stats *counters) error {
	if response == nil || response.Body == nil {
		stats.extractValidationFailures.Add(1)
		return ErrInvalidOptions
	}
	body, err := io.ReadAll(io.LimitReader(response.Body, maximumExtractResponseBytes+1))
	if err != nil || len(body) > maximumExtractResponseBytes {
		stats.extractValidationFailures.Add(1)
		if err != nil {
			return err
		}
		return ErrInvalidOptions
	}
	var payload struct {
		RequestID string          `json:"requestId"`
		Requested int             `json:"requested"`
		Returned  int             `json:"returned"`
		Proxies   json.RawMessage `json:"proxies"`
	}
	if err := json.Unmarshal(body, &payload); err != nil || payload.RequestID != requestID || payload.Requested != options.ExtractCount ||
		len(payload.Proxies) == 0 || bytes.Equal(bytes.TrimSpace(payload.Proxies), []byte("null")) ||
		(options.ExtractFulfillment == "allOrNothing" && payload.Returned != options.ExtractCount) {
		stats.extractValidationFailures.Add(1)
		if err != nil {
			return err
		}
		return ErrInvalidOptions
	}
	var proxies []struct {
		ID string `json:"id"`
	}
	if err := json.Unmarshal(payload.Proxies, &proxies); err != nil || payload.Returned < 0 || payload.Returned > options.ExtractCount ||
		payload.Returned != len(proxies) {
		stats.extractValidationFailures.Add(1)
		if err != nil {
			return err
		}
		return ErrInvalidOptions
	}
	ids := make(map[string]struct{}, len(proxies))
	for _, proxy := range proxies {
		if proxy.ID == "" {
			stats.extractValidationFailures.Add(1)
			return ErrInvalidOptions
		}
		if _, duplicate := ids[proxy.ID]; duplicate {
			stats.extractValidationFailures.Add(1)
			return ErrInvalidOptions
		}
		ids[proxy.ID] = struct{}{}
	}
	stats.extractResponsesValidated.Add(1)
	stats.extractReturned.Add(uint64(payload.Returned))
	return nil
}

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(),
		ExtractResponsesValidated: stats.extractResponsesValidated.Load(), ExtractReturned: stats.extractReturned.Load(),
		ExtractValidationFailures: stats.extractValidationFailures.Load(),
		RateStartsGenerated:       stats.rateStartsGenerated.Load(), RateStartsDropped: stats.rateStartsDropped.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)
}