// Command orbit-load runs a bounded, non-retrying task admission sample.
// It is a measurement tool for an isolated environment, not a production
// scheduler and not a capacity promise.
package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"math"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	orbit "git.yuntuo.hk/k2safe/orbit/sdk/go"
)

type credentials struct {
	Token   string `json:"token"`
	QueueID string `json:"queue_id"`
}

type sampleConfig struct {
	URL            string
	Token          string
	QueueID        string
	Requests       int
	Concurrency    int
	Timeout        time.Duration
	Handler        string
	HandlerVersion string
	PayloadVersion string
}

type sampleResult struct {
	duration time.Duration
	code     string
}

type loadReport struct {
	StartedAt   time.Time      `json:"started_at"`
	FinishedAt  time.Time      `json:"finished_at"`
	Requests    int            `json:"requests"`
	Concurrency int            `json:"concurrency"`
	Success     int            `json:"success"`
	Errors      int            `json:"errors"`
	P50MS       float64        `json:"p50_ms"`
	P95MS       float64        `json:"p95_ms"`
	P99MS       float64        `json:"p99_ms"`
	MaxMS       float64        `json:"max_ms"`
	ErrorCodes  map[string]int `json:"error_codes,omitempty"`
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func run() error {
	cfg, err := parseConfig(os.Args[1:])
	if err != nil {
		return err
	}
	client, err := orbit.NewClient(cfg.URL, orbit.Options{
		HTTPClient: &http.Client{Timeout: cfg.Timeout},
		Token:      cfg.Token,
	})
	if err != nil {
		return err
	}
	runID, err := newRunID()
	if err != nil {
		return err
	}
	started := time.Now().UTC()
	results := runSamples(client, cfg, runID)
	finished := time.Now().UTC()
	report := makeReport(cfg, started, finished, results)
	if err := json.NewEncoder(os.Stdout).Encode(report); err != nil {
		return err
	}
	if report.Errors != 0 {
		return errors.New("load sample had failed requests; inspect error_codes and rerun in an isolated environment")
	}
	return nil
}

func parseConfig(args []string) (sampleConfig, error) {
	flags := flag.NewFlagSet("orbit-load", flag.ContinueOnError)
	flags.SetOutput(os.Stderr)
	cfg := sampleConfig{
		URL:            os.Getenv("ORBIT_URL"),
		Token:          os.Getenv("ORBIT_TOKEN"),
		QueueID:        os.Getenv("ORBIT_QUEUE_ID"),
		Requests:       100,
		Concurrency:    8,
		Timeout:        10 * time.Second,
		Handler:        "load.noop",
		HandlerVersion: "v1",
		PayloadVersion: "v1",
	}
	credentialsPath := os.Getenv("ORBIT_CREDENTIALS_FILE")
	flags.StringVar(&cfg.URL, "url", cfg.URL, "Orbit HTTP URL")
	flags.StringVar(&cfg.Token, "token", cfg.Token, "Bearer token; prefer ORBIT_CREDENTIALS_FILE")
	flags.StringVar(&cfg.QueueID, "queue", cfg.QueueID, "target queue ID")
	flags.StringVar(&credentialsPath, "credentials", credentialsPath, "private bootstrap credentials JSON file")
	flags.IntVar(&cfg.Requests, "requests", cfg.Requests, "number of independent admissions, 1..100000")
	flags.IntVar(&cfg.Concurrency, "concurrency", cfg.Concurrency, "maximum concurrent HTTP requests, 1..1024")
	flags.DurationVar(&cfg.Timeout, "timeout", cfg.Timeout, "per-request HTTP timeout")
	flags.StringVar(&cfg.Handler, "handler", cfg.Handler, "handler name frozen in each task")
	flags.StringVar(&cfg.HandlerVersion, "handler-version", cfg.HandlerVersion, "handler version frozen in each task")
	flags.StringVar(&cfg.PayloadVersion, "payload-version", cfg.PayloadVersion, "payload contract version frozen in each task")
	if err := flags.Parse(args); err != nil {
		return sampleConfig{}, err
	}
	if flags.NArg() != 0 {
		return sampleConfig{}, errors.New("orbit-load does not accept positional arguments")
	}
	if credentialsPath != "" {
		file, err := os.ReadFile(credentialsPath)
		if err != nil {
			return sampleConfig{}, fmt.Errorf("read credentials: %w", err)
		}
		if len(file) > 16<<10 {
			return sampleConfig{}, errors.New("credentials file exceeds 16 KiB")
		}
		var saved credentials
		if err := json.Unmarshal(file, &saved); err != nil {
			return sampleConfig{}, errors.New("invalid credentials file")
		}
		if cfg.Token == "" {
			cfg.Token = saved.Token
		}
		if cfg.QueueID == "" {
			cfg.QueueID = saved.QueueID
		}
	}
	if cfg.URL == "" {
		return sampleConfig{}, errors.New("-url or ORBIT_URL is required")
	}
	if cfg.Token == "" {
		return sampleConfig{}, errors.New("-token, ORBIT_TOKEN or ORBIT_CREDENTIALS_FILE is required")
	}
	if !validID(cfg.QueueID) {
		return sampleConfig{}, errors.New("-queue or ORBIT_QUEUE_ID must be a 32-character lowercase hexadecimal ID")
	}
	if cfg.Requests < 1 || cfg.Requests > 100000 {
		return sampleConfig{}, errors.New("-requests must be between 1 and 100000")
	}
	if cfg.Concurrency < 1 || cfg.Concurrency > 1024 {
		return sampleConfig{}, errors.New("-concurrency must be between 1 and 1024")
	}
	if cfg.Timeout <= 0 || cfg.Timeout > time.Hour {
		return sampleConfig{}, errors.New("-timeout must be greater than zero and no more than one hour")
	}
	for name, value := range map[string]string{
		"handler":         cfg.Handler,
		"handler-version": cfg.HandlerVersion,
		"payload-version": cfg.PayloadVersion,
	} {
		if strings.TrimSpace(value) == "" || strings.ContainsAny(value, "\r\n") {
			return sampleConfig{}, fmt.Errorf("-%s must be a non-empty single line", name)
		}
	}
	if cfg.Concurrency > cfg.Requests {
		cfg.Concurrency = cfg.Requests
	}
	return cfg, nil
}

func runSamples(client *orbit.Client, cfg sampleConfig, runID string) []sampleResult {
	jobs := make(chan int)
	results := make(chan sampleResult, cfg.Requests)
	var wait sync.WaitGroup
	for worker := 0; worker < cfg.Concurrency; worker++ {
		wait.Add(1)
		go func() {
			defer wait.Done()
			for sequence := range jobs {
				start := time.Now()
				requestCtx, cancel := context.WithTimeout(context.Background(), cfg.Timeout)
				payload, _ := json.Marshal(struct {
					RunID    string `json:"run_id"`
					Sequence int    `json:"sequence"`
				}{RunID: runID, Sequence: sequence})
				_, err := client.SubmitTask(requestCtx, fmt.Sprintf("orbit-load-%s-%06d", runID, sequence), orbit.SubmitTaskRequest{
					QueueID:        cfg.QueueID,
					Handler:        cfg.Handler,
					HandlerVersion: cfg.HandlerVersion,
					PayloadVersion: cfg.PayloadVersion,
					Payload:        payload,
					TaskTimeoutMS:  60 * 60 * 1000,
					RecoverySafety: orbit.RecoveryReconcileOnly,
					Retry:          orbit.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: 10 * 1000},
				})
				cancel()
				code := ""
				if err != nil {
					code = errorCode(err)
				}
				results <- sampleResult{duration: time.Since(start), code: code}
			}
		}()
	}
	for sequence := 0; sequence < cfg.Requests; sequence++ {
		jobs <- sequence
	}
	close(jobs)
	wait.Wait()
	close(results)
	out := make([]sampleResult, 0, cfg.Requests)
	for result := range results {
		out = append(out, result)
	}
	return out
}

func makeReport(cfg sampleConfig, started, finished time.Time, results []sampleResult) loadReport {
	durations := make([]time.Duration, 0, len(results))
	errorCodes := make(map[string]int)
	success := 0
	for _, result := range results {
		durations = append(durations, result.duration)
		if result.code == "" {
			success++
		} else {
			errorCodes[result.code]++
		}
	}
	sort.Slice(durations, func(i, j int) bool { return durations[i] < durations[j] })
	return loadReport{
		StartedAt: started, FinishedAt: finished,
		Requests: cfg.Requests, Concurrency: cfg.Concurrency,
		Success: success, Errors: len(results) - success,
		P50MS: percentileMS(durations, 0.50), P95MS: percentileMS(durations, 0.95),
		P99MS: percentileMS(durations, 0.99), MaxMS: percentileMS(durations, 1),
		ErrorCodes: errorCodes,
	}
}

func percentileMS(values []time.Duration, fraction float64) float64 {
	if len(values) == 0 {
		return 0
	}
	index := int(math.Ceil(float64(len(values))*fraction)) - 1
	if index < 0 {
		index = 0
	}
	if index >= len(values) {
		index = len(values) - 1
	}
	return float64(values[index]) / float64(time.Millisecond)
}

func errorCode(err error) string {
	var apiErr *orbit.APIError
	if errors.As(err, &apiErr) {
		if apiErr.Code != "" {
			return apiErr.Code
		}
		return fmt.Sprintf("HTTP_%d", apiErr.StatusCode)
	}
	var httpErr *orbit.HTTPError
	if errors.As(err, &httpErr) {
		return fmt.Sprintf("HTTP_%d", httpErr.StatusCode)
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return "TIMEOUT"
	}
	return "TRANSPORT"
}

func newRunID() (string, error) {
	var bytes [16]byte
	if _, err := rand.Read(bytes[:]); err != nil {
		return "", fmt.Errorf("generate load run id: %w", err)
	}
	return hex.EncodeToString(bytes[:]), nil
}

func validID(value string) bool {
	if len(value) != 32 {
		return false
	}
	for _, ch := range value {
		if !(ch >= '0' && ch <= '9' || ch >= 'a' && ch <= 'f') {
			return false
		}
	}
	return true
}