#!/usr/bin/env python3
import asyncio
import aiohttp
import base64
import time
import json
from typing import List

API_URL = "http://127.0.0.1:8008/api/v1/inpaint"


def img_to_b64(path: str) -> str:
    with open(path, "rb") as f:
        return "data:image/png;base64," + base64.b64encode(f.read()).decode("utf-8")


async def post_inpaint(session: aiohttp.ClientSession, image_b64: str, mask_b64: str):
    payload = {
        "image": image_b64,
        "mask": mask_b64,
        "model_name": "simple-lama",
    }
    t0 = time.time()
    async with session.post(API_URL, json=payload) as resp:
        _ = await resp.read()
        dt = time.time() - t0
        return resp.status, dt


async def run_scenario(concurrency: int, total: int, img_path: str, mask_path: str):
    image_b64 = img_to_b64(img_path)
    mask_b64 = img_to_b64(mask_path)
    async with aiohttp.ClientSession() as session:
        tasks = []
        # 워밍업: 배치 2를 유도하기 위해 빠르게 2개
        for _ in range(2):
            tasks.append(asyncio.create_task(post_inpaint(session, image_b64, mask_b64)))
        results = await asyncio.gather(*tasks)
        print("Warmup:", results)

        # 본 테스트: 동시성(concurrency)로 total개 요청 보내기
        pending = 0
        idx = 0
        inflight: List[asyncio.Task] = []
        latencies = []
        while idx < total or inflight:
            while idx < total and len(inflight) < concurrency:
                inflight.append(asyncio.create_task(post_inpaint(session, image_b64, mask_b64)))
                idx += 1
            done, inflight = await asyncio.wait(inflight, return_when=asyncio.FIRST_COMPLETED)
            for d in done:
                code, dt = d.result()
                latencies.append(dt)
        print("count=", len(latencies), "avg=", sum(latencies)/len(latencies), "max=", max(latencies))


if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--img", required=True)
    p.add_argument("--mask", required=True)
    p.add_argument("--concurrency", type=int, default=2)
    p.add_argument("--total", type=int, default=10)
    args = p.parse_args()
    asyncio.run(run_scenario(args.concurrency, args.total, args.img, args.mask))