# -*- coding: utf-8 -*-
"""
Locust 压力测试脚本 - VF React 网站

使用方式:
    # 安装依赖
    pip install locust

    # 启动 Web UI 模式（推荐）
    locust -f locustfile.py --host=https://valuefrontier.cn

    # 命令行模式（无 UI）
    locust -f locustfile.py --host=https://valuefrontier.cn \
           --users 1000 --spawn-rate 50 --run-time 5m --headless

    # 分布式模式（多机压测）
    # 主节点
    locust -f locustfile.py --master --host=https://valuefrontier.cn
    # 从节点（在其他机器运行）
    locust -f locustfile.py --worker --master-host=<master-ip>

Web UI 访问: http://localhost:8089
"""

from locust import HttpUser, task, between, events
import random
import time
import json


class WebsiteUser(HttpUser):
    """模拟普通网站用户行为"""

    # 用户请求间隔（1-3秒，模拟真实用户）
    wait_time = between(1, 3)

    def on_start(self):
        """用户启动时执行（可用于登录）"""
        self.stock_codes = [
            "600000", "600036", "601318", "000001", "000002",
            "300750", "002594", "601888", "600519", "000858"
        ]

    @task(10)
    def view_homepage(self):
        """访问首页（权重 10）"""
        self.client.get("/", name="首页")

    @task(8)
    def get_stock_list(self):
        """获取股票列表（权重 8）"""
        self.client.get("/api/stocks", name="股票列表 API")

    @task(5)
    def get_stock_detail(self):
        """获取个股详情（权重 5）"""
        code = random.choice(self.stock_codes)
        self.client.get(f"/api/stocks/{code}", name="个股详情 API")

    @task(3)
    def get_community_events(self):
        """获取社区事件（权重 3）"""
        self.client.get("/api/community/events", name="社区事件 API")

    @task(2)
    def get_market_overview(self):
        """获取市场概览（权重 2）"""
        self.client.get("/api/market/overview", name="市场概览 API")


class APIUser(HttpUser):
    """模拟 API 密集型用户（重度用户）"""

    wait_time = between(0.5, 1.5)

    def on_start(self):
        self.stock_codes = [
            "600000", "600036", "601318", "000001", "000002",
            "300750", "002594", "601888", "600519", "000858"
        ]

    @task(5)
    def get_realtime_quote(self):
        """获取实时行情"""
        code = random.choice(self.stock_codes)
        self.client.get(f"/api/quote/{code}", name="实时行情 API")

    @task(3)
    def get_kline_data(self):
        """获取 K 线数据"""
        code = random.choice(self.stock_codes)
        self.client.get(f"/api/kline/{code}?period=day", name="K线数据 API")

    @task(2)
    def get_concept_stocks(self):
        """获取概念板块股票"""
        self.client.get("/api/concepts", name="概念板块 API")


class AuthenticatedUser(HttpUser):
    """模拟登录用户行为"""

    wait_time = between(2, 5)

    def on_start(self):
        """登录获取 session"""
        # 如果有测试账号，可以在这里登录
        # response = self.client.post("/api/auth/login", json={
        #     "username": "test_user",
        #     "password": "test_pass"
        # })
        pass

    @task(3)
    def view_portfolio(self):
        """查看投资组合"""
        self.client.get("/api/portfolio", name="投资组合 API")

    @task(2)
    def view_watchlist(self):
        """查看自选股"""
        self.client.get("/api/watchlist", name="自选股 API")


# ==================== 自定义统计 ====================

@events.request.add_listener
def on_request(request_type, name, response_time, response_length, exception, **kwargs):
    """请求完成后的回调（可用于自定义监控）"""
    if exception:
        print(f"❌ 请求失败: {name} - {exception}")
    elif response_time > 1000:  # 超过 1 秒的慢请求
        print(f"⚠️ 慢请求: {name} - {response_time:.0f}ms")


@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    """测试开始时执行"""
    print("=" * 60)
    print("🚀 压力测试开始!")
    print(f"   目标: {environment.host}")
    print("=" * 60)


@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    """测试结束时执行"""
    print("=" * 60)
    print("🛑 压力测试结束!")
    print("=" * 60)