"""Locust 性能压测脚本 — 天工智能体平台
启动: locust -f deploy/locustfile.py --host=http://localhost:8037
"""
from locust import HttpUser, task, between
import json


class AIPlatformUser(HttpUser):
    wait_time = between(1, 3)

    def on_start(self):
        """登录获取 token。"""
        resp = self.client.post(
            "/api/v1/auth/login",
            headers={"Content-Type": "application/x-www-form-urlencoded"},
            data={"username": "admin", "password": "123456"},
        )
        if resp.status_code == 200:
            self.token = resp.json().get("access_token", "")
        else:
            self.token = ""

    def _headers(self):
        return {"Authorization": f"Bearer {self.token}"} if self.token else {}

    @task(5)
    def get_agents(self):
        self.client.get("/api/v1/agents", headers=self._headers())

    @task(3)
    def agent_chat(self):
        self.client.post(
            "/api/v1/agent-chat/bare",
            json={"message": "你好，请简单介绍一下自己"},
            headers=self._headers(),
        )

    @task(2)
    def get_executions(self):
        self.client.get("/api/v1/executions?limit=10", headers=self._headers())

    @task(2)
    def health_check(self):
        self.client.get("/health")

    @task(1)
    def get_workflows(self):
        self.client.get("/api/v1/workflows?limit=20", headers=self._headers())