package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"proxy-pool/internal/loadgen"
)

type loadRun func(context.Context, loadgen.Options) (loadgen.Report, error)

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	os.Exit(execute(ctx, os.Args[1:], loadgen.Run, os.Stdout, os.Stderr))
}

func execute(ctx context.Context, args []string, run loadRun, stdout, stderr io.Writer) int {
	flags := flag.NewFlagSet("proxy-loadgen", flag.ContinueOnError)
	flags.SetOutput(stderr)
	targetURL := flags.String("target", "", "HTTP or HTTPS target URL")
	proxyURL := flags.String("proxy", "", "optional HTTP or HTTPS forward proxy URL")
	method := flags.String("method", http.MethodGet, "HTTP method")
	body := flags.String("body", "", "UTF-8 request body")
	requests := flags.Int("requests", 0, "fixed request count; mutually exclusive with -duration")
	duration := flags.Duration("duration", 0, "time-boxed workload duration")
	rate := flags.Int("rate", 0, "maximum request starts per second for -duration; zero saturates workers")
	concurrency := flags.Int("concurrency", 64, "maximum concurrent requests")
	timeout := flags.Duration("timeout", 10*time.Second, "per-request timeout")
	var headers headerValues
	flags.Var(&headers, "header", "repeatable HTTP header in Name: Value form")
	if err := flags.Parse(args); err != nil {
		if errors.Is(err, flag.ErrHelp) {
			return 0
		}
		return 2
	}
	if flags.NArg() != 0 || ctx == nil || run == nil {
		_, _ = fmt.Fprintln(stderr, "proxy-loadgen: target and a bounded workload are required")
		return 2
	}
	parsedHeaders, err := headers.Header()
	if err != nil {
		_, _ = fmt.Fprintf(stderr, "proxy-loadgen: %v\n", err)
		return 2
	}
	report, err := run(ctx, loadgen.Options{
		TargetURL: *targetURL, ProxyURL: *proxyURL, Method: *method, Headers: parsedHeaders, RequestBody: []byte(*body),
		Requests: *requests, Duration: *duration, Rate: *rate, Concurrency: *concurrency, RequestTimeout: *timeout,
	})
	if err != nil {
		if errors.Is(err, loadgen.ErrInvalidOptions) {
			_, _ = fmt.Fprintln(stderr, "proxy-loadgen: -target, positive -concurrency/-timeout, and exactly one bounded workload are required")
			return 2
		}
		_, _ = fmt.Fprintf(stderr, "proxy-loadgen: %v\n", err)
		return 1
	}
	if err := json.NewEncoder(stdout).Encode(report); err != nil {
		_, _ = fmt.Fprintf(stderr, "proxy-loadgen: encode report: %v\n", err)
		return 1
	}
	return 0
}

type headerValues []string

func (values *headerValues) String() string {
	return strings.Join(*values, ",")
}

func (values *headerValues) Set(value string) error {
	if strings.TrimSpace(value) != value || value == "" {
		return errors.New("invalid header")
	}
	*values = append(*values, value)
	return nil
}

func (values headerValues) Header() (http.Header, error) {
	result := make(http.Header, len(values))
	for _, value := range values {
		name, content, found := strings.Cut(value, ":")
		name = strings.TrimSpace(name)
		content = strings.TrimSpace(content)
		if !found || name == "" || strings.ContainsAny(name, "\r\n") || strings.ContainsAny(content, "\r\n") {
			return nil, errors.New("invalid header")
		}
		result.Add(name, content)
	}
	return result, nil
}