// Command loadtest 是平台容量压测器：闭环加压、阶梯并发，输出吞吐/延迟分位曲线。
//
// 它只通过网关 HTTP API 打真实流量（登录 → POST /tasks → SSE 流等收尾），
// 故压的是「网关→NATS→调度→工具RTT→状态回流」全链路。配合 dispatcher 的 LLM_FORCE_STUB=1
// + LLM_STUB_*_MS=0 可绕开真实 LLM 推理，量出平台自身吞吐天花板（而非被模型节奏掩盖）。
//
// 用法：
//
//	go run ./cmd/loadtest -url http://localhost:8080 -email keytest@local.dev -pass keytest123456 \
//	  -levels 1,2,4,8,16,32,64 -dur 12s
package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

func main() {
	var (
		base     = flag.String("url", "http://localhost:8080", "网关地址")
		email    = flag.String("email", "keytest@local.dev", "登录邮箱")
		pass     = flag.String("pass", "keytest123456", "登录密码")
		levelStr = flag.String("levels", "1,2,4,8,16,32,64", "并发阶梯（逗号分隔）")
		dur      = flag.Duration("dur", 12*time.Second, "每个阶梯加压时长")
		prompt   = flag.String("prompt", "用一句话介绍杭州", "任务输入")
	)
	flag.Parse()

	token, err := login(*base, *email, *pass)
	if err != nil {
		fmt.Printf("登录失败: %v\n", err)
		os.Exit(1)
	}
	body := taskBody(*prompt)

	fmt.Printf("== 平台容量压测 ==  目标 %s  每阶梯 %s  (FORCE_STUB 下为平台天花板)\n", *base, *dur)
	fmt.Printf("%-6s %-10s %-10s %-9s %-9s %-9s %-7s\n", "并发", "完成", "吞吐/s", "p50ms", "p95ms", "max ms", "错误")
	for _, ls := range strings.Split(*levelStr, ",") {
		c, err := strconv.Atoi(strings.TrimSpace(ls))
		if err != nil || c <= 0 {
			continue
		}
		r := runLevel(*base, token, body, c, *dur)
		fmt.Printf("%-6d %-10d %-10.1f %-9d %-9d %-9d %-7d\n",
			c, r.done, float64(r.done)/dur.Seconds(), r.p50, r.p95, r.max, r.errs)
	}
}

type result struct {
	done, errs    int64
	p50, p95, max int64
}

// runLevel 在并发 c 下闭环加压 dur：c 个 worker 不停 submit→等终态→再来，统计完成数与 e2e 延迟。
func runLevel(base, token string, body []byte, c int, dur time.Duration) result {
	ctx, cancel := context.WithTimeout(context.Background(), dur)
	defer cancel()
	var done, errs int64
	var mu sync.Mutex
	var lats []int64
	var wg sync.WaitGroup
	for i := 0; i < c; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			cl := &http.Client{Timeout: 60 * time.Second}
			for ctx.Err() == nil {
				t0 := time.Now()
				if runOne(ctx, cl, base, token, body) {
					ms := time.Since(t0).Milliseconds()
					mu.Lock()
					lats = append(lats, ms)
					mu.Unlock()
					atomic.AddInt64(&done, 1)
				} else {
					atomic.AddInt64(&errs, 1)
				}
			}
		}()
	}
	wg.Wait()
	return result{done: done, errs: errs, p50: pct(lats, 50), p95: pct(lats, 95), max: pct(lats, 100)}
}

// runOne 提交一个任务并经 SSE token 流等到收尾（流结束即任务完成）；返回是否成功。
// 用流而非轮询：每任务 1 条推送连接，避免轮询放大把网关自身压成瓶颈，量出真实管线吞吐。
func runOne(ctx context.Context, cl *http.Client, base, token string, body []byte) bool {
	id, err := submit(ctx, cl, base, token, body)
	if err != nil || id == "" {
		return false
	}
	req, _ := http.NewRequestWithContext(ctx, "GET", base+"/api/v1/tasks/"+id+"/stream", nil)
	req.Header.Set("Authorization", "Bearer "+token)
	resp, err := cl.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	// 读到 EOF：网关在 token 流结束（任务收尾）时关闭 SSE。
	_, err = io.Copy(io.Discard, resp.Body)
	return err == nil && ctx.Err() == nil
}

func submit(ctx context.Context, cl *http.Client, base, token string, body []byte) (string, error) {
	req, _ := http.NewRequestWithContext(ctx, "POST", base+"/api/v1/tasks", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+token)
	resp, err := cl.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusAccepted {
		io.Copy(io.Discard, resp.Body)
		return "", fmt.Errorf("submit %d", resp.StatusCode)
	}
	var r struct {
		TaskID string `json:"task_id"`
	}
	json.NewDecoder(resp.Body).Decode(&r)
	return r.TaskID, nil
}

func login(base, email, pass string) (string, error) {
	b, _ := json.Marshal(map[string]string{"email": email, "password": pass})
	resp, err := http.Post(base+"/api/v1/auth/login", "application/json", bytes.NewReader(b))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	var r struct {
		Token string `json:"token"`
	}
	json.NewDecoder(resp.Body).Decode(&r)
	if r.Token == "" {
		return "", fmt.Errorf("no token (status %d)", resp.StatusCode)
	}
	return r.Token, nil
}

func taskBody(prompt string) []byte {
	graph := map[string]any{
		"nodes": []any{
			map[string]any{"id": "i", "kind": "input", "config": map[string]any{"text": prompt}},
			map[string]any{"id": "a", "kind": "agent", "config": map[string]any{"system": "你是简洁的助手"}},
		},
		"edges": []any{map[string]any{"source": "i", "target": "a"}},
	}
	b, _ := json.Marshal(map[string]any{"graph": graph, "meta": map[string]any{}})
	return b
}

// pct 返回延迟切片的第 p 百分位（p=100 即 max）。
func pct(v []int64, p int) int64 {
	if len(v) == 0 {
		return 0
	}
	sort.Slice(v, func(i, j int) bool { return v[i] < v[j] })
	idx := (p * (len(v) - 1)) / 100
	return v[idx]
}