from locust import HttpUser, task, between, tag
import random
import string
import json
import io
from PIL import Image  # 需要安装Pillow库: pip install pillow


class GradioImageUser(HttpUser):
    wait_time = between(1, 3)  # 用户操作间隔时间

    def on_start(self):
        """用户启动时的初始化操作"""
        # 生成随机字符串作为上传ID和会话标识
        self.upload_id = ''.join(random.choices(string.ascii_lowercase + string.digits, k=10))

        # 生成随机测试图片(200x200像素)作为上传内容
        self.test_image = self.generate_test_image()

    def generate_test_image(self):
        """生成随机测试图片并返回字节流"""
        # 创建200x200的随机彩色图片
        img = Image.new('RGB', (200, 200), color=(
            random.randint(0, 255),
            random.randint(0, 255),
            random.randint(0, 255)
        ))

        # 将图片保存到字节流
        img_byte_arr = io.BytesIO()
        img.save(img_byte_arr, format='JPEG')  # 保存为JPEG格式
        img_byte_arr.seek(0)  # 移动到字节流开头

        return img_byte_arr

    @task(2)  # 权重为2，执行频率更高
    @tag("upload")
    def upload_image(self):
        """模拟图片上传操作"""
        # 构建multipart/form-data格式的图片上传请求
        # 使用随机文件名，模拟不同用户上传不同图片
        random_filename = f"test_image_{random.randint(1000, 9999)}.jpg"

        files = {
            'file': (
                random_filename,
                self.test_image,
                'image/jpeg'  # 图片MIME类型
            )
        }

        # 发送图片上传请求
        with self.client.post(
                f"/gradio_api/upload?upload_id={self.upload_id}",
                files=files,
                catch_response=True
        ) as response:
            if response.status_code != 200:
                response.failure(f"上传失败，状态码: {response.status_code}")
            else:
                try:
                    # 尝试解析JSON响应
                    json_response = response.json()
                    # 可以根据实际API响应添加更多验证
                    response.success()
                except json.JSONDecodeError:
                    response.failure("响应不是有效的JSON格式")

    @task(1)  # 权重为1
    @tag("queue")
    def join_queue(self):
        """模拟加入队列操作"""
        # 准备队列请求的数据
        payload = {
            "data": [
                f"uploaded_file:{self.upload_id}",  # 关联之前上传的图片
                True  # 保持与实际请求一致的参数
            ],
            "event_data": None,
            "fn_index": 0,  # 根据实际功能索引调整
            "session_hash": self.upload_id  # 使用相同的upload_id作为会话标识
        }

        # 发送加入队列请求
        with self.client.post(
                "/gradio_api/queue/join?",
                json=payload,
                headers={"Content-Type": "application/json"},
                catch_response=True
        ) as response:
            if response.status_code != 200:
                response.failure(f"加入队列失败，状态码: {response.status_code}")
            else:
                try:
                    json_response = response.json()
                    response.success()
                except json.JSONDecodeError:
                    response.failure("响应不是有效的JSON格式")