package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

// 压力测试工具 - 模拟实际业务场景

type TestConfig struct {
	BaseURL     string
	Concurrency int
	Duration    time.Duration
	Scenario    string
}

type TestResult struct {
	TotalRequests   int64
	SuccessRequests int64
	FailedRequests  int64
	TotalLatency    int64 // 毫秒
	MinLatency      int64
	MaxLatency      int64
	Errors          map[string]int64
}

var (
	baseURL     = flag.String("url", "http://localhost:8080", "API 基础地址")
	concurrency = flag.Int("c", 10, "并发数")
	duration    = flag.Int("d", 60, "测试时长（秒）")
	scenario    = flag.String("s", "mixed", "测试场景: list_listings, create_order, chat, wallet, mixed")
)

func main() {
	flag.Parse()

	config := TestConfig{
		BaseURL:     *baseURL,
		Concurrency: *concurrency,
		Duration:    time.Duration(*duration) * time.Second,
		Scenario:    *scenario,
	}

	fmt.Printf("=== 压力测试配置 ===\n")
	fmt.Printf("目标地址: %s\n", config.BaseURL)
	fmt.Printf("并发数: %d\n", config.Concurrency)
	fmt.Printf("测试时长: %d 秒\n", *duration)
	fmt.Printf("测试场景: %s\n", config.Scenario)
	fmt.Printf("==================\n\n")

	result := runTest(config)
	printResult(result)
}

func runTest(config TestConfig) *TestResult {
	result := &TestResult{
		Errors:     make(map[string]int64),
		MinLatency: int64(^uint64(0) >> 1), // Max int64
	}

	var wg sync.WaitGroup
	stopChan := make(chan struct{})

	// 启动并发workers
	for i := 0; i < config.Concurrency; i++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()
			worker(workerID, config, result, stopChan)
		}(i)
	}

	// 等待测试时长
	time.Sleep(config.Duration)
	close(stopChan)

	wg.Wait()

	return result
}

func worker(id int, config TestConfig, result *TestResult, stopChan chan struct{}) {
	client := &http.Client{
		Timeout: 10 * time.Second,
	}

	for {
		select {
		case <-stopChan:
			return
		default:
			executeScenario(client, config, result)
		}
	}
}

func executeScenario(client *http.Client, config TestConfig, result *TestResult) {
	switch config.Scenario {
	case "list_listings":
		testListListings(client, config.BaseURL, result)
	case "create_order":
		testCreateOrder(client, config.BaseURL, result)
	case "chat":
		testChatMessages(client, config.BaseURL, result)
	case "wallet":
		testWalletLedger(client, config.BaseURL, result)
	case "mixed":
		// 混合场景：按实际业务比例分配
		r := rand.Intn(100)
		switch {
		case r < 40: // 40% 查询商品列表
			testListListings(client, config.BaseURL, result)
		case r < 60: // 20% 查询订单
			testListOrders(client, config.BaseURL, result)
		case r < 75: // 15% 查询钱包流水
			testWalletLedger(client, config.BaseURL, result)
		case r < 85: // 10% 聊天消息
			testChatMessages(client, config.BaseURL, result)
		case r < 95: // 10% 创建订单
			testCreateOrder(client, config.BaseURL, result)
		default: // 5% 支付
			testPayOrder(client, config.BaseURL, result)
		}
	default:
		testHealthCheck(client, config.BaseURL, result)
	}
}

// ============================================
// 测试场景实现
// ============================================

func testHealthCheck(client *http.Client, baseURL string, result *TestResult) {
	start := time.Now()
	resp, err := client.Get(baseURL + "/health")
	latency := time.Since(start).Milliseconds()

	atomic.AddInt64(&result.TotalRequests, 1)
	updateLatency(result, latency)

	if err != nil {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, "health_check_error: "+err.Error())
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode == 200 {
		atomic.AddInt64(&result.SuccessRequests, 1)
	} else {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, fmt.Sprintf("health_check_status_%d", resp.StatusCode))
	}
}

func testListListings(client *http.Client, baseURL string, result *TestResult) {
	// 模拟不同的查询条件
	page := rand.Intn(10) + 1
	pageSize := []int{10, 20, 50}[rand.Intn(3)]
	url := fmt.Sprintf("%s/api/listings?page=%d&page_size=%d", baseURL, page, pageSize)

	start := time.Now()
	resp, err := client.Get(url)
	latency := time.Since(start).Milliseconds()

	atomic.AddInt64(&result.TotalRequests, 1)
	updateLatency(result, latency)

	if err != nil {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, "list_listings_error: "+err.Error())
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode == 200 {
		atomic.AddInt64(&result.SuccessRequests, 1)
	} else {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, fmt.Sprintf("list_listings_status_%d", resp.StatusCode))
	}
}

func testListOrders(client *http.Client, baseURL string, result *TestResult) {
	// 需要登录token，这里模拟匿名访问（会返回401）
	page := rand.Intn(5) + 1
	url := fmt.Sprintf("%s/api/orders?page=%d&page_size=20", baseURL, page)

	start := time.Now()
	resp, err := client.Get(url)
	latency := time.Since(start).Milliseconds()

	atomic.AddInt64(&result.TotalRequests, 1)
	updateLatency(result, latency)

	if err != nil {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, "list_orders_error: "+err.Error())
		return
	}
	defer resp.Body.Close()

	// 401是预期的（未登录）
	if resp.StatusCode == 200 || resp.StatusCode == 401 {
		atomic.AddInt64(&result.SuccessRequests, 1)
	} else {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, fmt.Sprintf("list_orders_status_%d", resp.StatusCode))
	}
}

func testCreateOrder(client *http.Client, baseURL string, result *TestResult) {
	// 模拟创建订单（需要登录，会返回401）
	listingID := rand.Intn(1000) + 1
	payload := map[string]interface{}{
		"listing_id":               listingID,
		"estimated_duration_hours": 24,
	}

	body, _ := json.Marshal(payload)
	start := time.Now()
	resp, err := client.Post(
		baseURL+"/api/orders",
		"application/json",
		bytes.NewBuffer(body),
	)
	latency := time.Since(start).Milliseconds()

	atomic.AddInt64(&result.TotalRequests, 1)
	updateLatency(result, latency)

	if err != nil {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, "create_order_error: "+err.Error())
		return
	}
	defer resp.Body.Close()

	// 401是预期的（未登录）
	if resp.StatusCode == 200 || resp.StatusCode == 401 {
		atomic.AddInt64(&result.SuccessRequests, 1)
	} else {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, fmt.Sprintf("create_order_status_%d", resp.StatusCode))
	}
}

func testPayOrder(client *http.Client, baseURL string, result *TestResult) {
	orderID := rand.Intn(1000) + 1
	url := fmt.Sprintf("%s/api/orders/%d/pay", baseURL, orderID)

	payload := map[string]interface{}{
		"provider": "mock",
	}

	body, _ := json.Marshal(payload)
	start := time.Now()
	resp, err := client.Post(url, "application/json", bytes.NewBuffer(body))
	latency := time.Since(start).Milliseconds()

	atomic.AddInt64(&result.TotalRequests, 1)
	updateLatency(result, latency)

	if err != nil {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, "pay_order_error: "+err.Error())
		return
	}
	defer resp.Body.Close()

	// 401是预期的（未登录）
	if resp.StatusCode == 200 || resp.StatusCode == 401 {
		atomic.AddInt64(&result.SuccessRequests, 1)
	} else {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, fmt.Sprintf("pay_order_status_%d", resp.StatusCode))
	}
}

func testWalletLedger(client *http.Client, baseURL string, result *TestResult) {
	page := rand.Intn(10) + 1
	url := fmt.Sprintf("%s/api/wallet/ledger?page=%d&page_size=20", baseURL, page)

	start := time.Now()
	resp, err := client.Get(url)
	latency := time.Since(start).Milliseconds()

	atomic.AddInt64(&result.TotalRequests, 1)
	updateLatency(result, latency)

	if err != nil {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, "wallet_ledger_error: "+err.Error())
		return
	}
	defer resp.Body.Close()

	// 401是预期的（未登录）
	if resp.StatusCode == 200 || resp.StatusCode == 401 {
		atomic.AddInt64(&result.SuccessRequests, 1)
	} else {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, fmt.Sprintf("wallet_ledger_status_%d", resp.StatusCode))
	}
}

func testChatMessages(client *http.Client, baseURL string, result *TestResult) {
	conversationID := rand.Intn(100) + 1
	url := fmt.Sprintf("%s/api/chats/%d/messages?page=1&page_size=50", baseURL, conversationID)

	start := time.Now()
	resp, err := client.Get(url)
	latency := time.Since(start).Milliseconds()

	atomic.AddInt64(&result.TotalRequests, 1)
	updateLatency(result, latency)

	if err != nil {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, "chat_messages_error: "+err.Error())
		return
	}
	defer resp.Body.Close()

	// 401是预期的（未登录）
	if resp.StatusCode == 200 || resp.StatusCode == 401 {
		atomic.AddInt64(&result.SuccessRequests, 1)
	} else {
		atomic.AddInt64(&result.FailedRequests, 1)
		recordError(result, fmt.Sprintf("chat_messages_status_%d", resp.StatusCode))
	}
}

// ============================================
// 辅助函数
// ============================================

func updateLatency(result *TestResult, latency int64) {
	atomic.AddInt64(&result.TotalLatency, latency)

	// 更新最小延迟
	for {
		current := atomic.LoadInt64(&result.MinLatency)
		if latency >= current {
			break
		}
		if atomic.CompareAndSwapInt64(&result.MinLatency, current, latency) {
			break
		}
	}

	// 更新最大延迟
	for {
		current := atomic.LoadInt64(&result.MaxLatency)
		if latency <= current {
			break
		}
		if atomic.CompareAndSwapInt64(&result.MaxLatency, current, latency) {
			break
		}
	}
}

var errorMutex sync.Mutex

func recordError(result *TestResult, errMsg string) {
	errorMutex.Lock()
	defer errorMutex.Unlock()
	result.Errors[errMsg]++
}

func printResult(result *TestResult) {
	fmt.Printf("\n=== 压力测试结果 ===\n")
	fmt.Printf("总请求数: %d\n", result.TotalRequests)
	fmt.Printf("成功请求: %d (%.2f%%)\n",
		result.SuccessRequests,
		float64(result.SuccessRequests)/float64(result.TotalRequests)*100)
	fmt.Printf("失败请求: %d (%.2f%%)\n",
		result.FailedRequests,
		float64(result.FailedRequests)/float64(result.TotalRequests)*100)

	if result.TotalRequests > 0 {
		avgLatency := result.TotalLatency / result.TotalRequests
		fmt.Printf("\n延迟统计:\n")
		fmt.Printf("  最小延迟: %d ms\n", result.MinLatency)
		fmt.Printf("  平均延迟: %d ms\n", avgLatency)
		fmt.Printf("  最大延迟: %d ms\n", result.MaxLatency)
	}

	if len(result.Errors) > 0 {
		fmt.Printf("\n错误统计:\n")
		for err, count := range result.Errors {
			fmt.Printf("  %s: %d 次\n", err, count)
		}
	}

	qps := float64(result.TotalRequests) / float64(*duration)
	fmt.Printf("\nQPS: %.2f\n", qps)
	fmt.Printf("==================\n")
}