import asyncio
import httpx
import time
import io
from PIL import Image

# 配置后端 API 地址
URL = "http://localhost:13002/process-image"

async def send_request(client, task_id):
    # Create a unique dummy image for each request
    file_io = io.BytesIO()
    image = Image.new('RGBA', size=(800, 800), color=(task_id * 50, 0, 0, 255))
    image.save(file_io, 'png')
    file_io.seek(0)
    
    print(f"[Task {task_id}] Sending request...")
    start_time = time.time()
    
    try:
        response = await client.post(
            URL,
            files={"file": ("test.png", file_io, "image/png")},
            data={"width": "600", "height": "600"},
            timeout=60.0
        )
        end_time = time.time()
        print(f"[Task {task_id}] Completed in {end_time - start_time:.2f}s with status {response.status_code}")
        return response.status_code
    except Exception as e:
        print(f"[Task {task_id}] Failed: {str(e)}")
        return None

async def run_stress_test():
    async with httpx.AsyncClient() as client:
        # Fire 3 requests simultaneously
        tasks = [send_request(client, i) for i in range(1, 4)]
        print(f"--- Starting Stress Test: 3 Concurrent Requests ---")
        start_total = time.time()
        results = await asyncio.gather(*tasks)
        end_total = time.time()
        print(f"--- Stress Test Finished ---")
        print(f"Total time for all tasks: {end_total - start_total:.2f}s")
        print(f"Success count: {results.count(200)}/3")

if __name__ == "__main__":
    asyncio.run(run_stress_test())