import asyncio
import httpx
import base64
from pathlib import Path
import time

# 테스트 설정
BASE_URL = "http://localhost:8008/api/v1/inpaint"
# 테스트에 사용할 이미지와 마스크 경로
IMAGE_PATH = Path(__file__).parent / "test_image.png"
MASK_PATH = Path(__file__).parent / "test_mask.png"
CONCURRENT_REQUESTS = 8  # 동시에 보낼 요청 수 (min_sessions < 요청 수 < max_sessions * 2 정도로 설정)

def encode_image(image_path: Path) -> str:
    """이미지 파일을 base64로 인코딩"""
    try:
        with open(image_path, "rb") as f:
            return base64.b64encode(f.read()).decode('utf-8')
    except FileNotFoundError as e:
        print(f"오류: 테스트 이미지를 찾을 수 없습니다 - {e}")
        print("먼저 tests/generate_test_images.py 를 실행하여 테스트 이미지를 생성해주세요.")
        raise

async def send_request(client: httpx.AsyncClient, request_num: int, image_b64: str, mask_b64: str):
    """단일 API 요청을 보냄"""
    payload = {
        "image": image_b64,
        "mask": mask_b64,
        "model_name": "simple-lama",
        "ldm_steps": 20,
        "ldm_sampler": "plms",
        "hd_strategy": "Original",
        "hd_strategy_crop_margin": 128,
        "hd_strategy_crop_trigger_size": 800,
        "hd_strategy_resize_limit": 1280
    }
    
    print(f"[{request_num:02d}] 요청 시작...")
    start_time = time.time()
    try:
        response = await client.post(BASE_URL, json=payload, timeout=120)
        duration = time.time() - start_time
        
        if response.status_code == 200:
            print(f"✅ [{request_num:02d}] 요청 성공 ({duration:.2f}초)")
            # with open(f"result_{request_num}.png", "wb") as f:
            #     f.write(response.content)
        else:
            print(f"❌ [{request_num:02d}] 요청 실패 - 상태 코드: {response.status_code} ({duration:.2f}초)")
            print(f"   응답: {response.text[:200]}")
            
    except httpx.ReadTimeout:
        duration = time.time() - start_time
        print(f"❌ [{request_num:02d}] 요청 시간 초과 ({duration:.2f}초)")
    except httpx.ConnectError as e:
        print(f"❌ [{request_num:02d}] 연결 실패: {e}")

async def main():
    """메인 실행 함수"""
    print("SimpleLama 부하 테스트 시작...")
    print(f"동시 요청 수: {CONCURRENT_REQUESTS}")
    
    # 이미지 파일 인코딩
    image_b64 = encode_image(IMAGE_PATH)
    mask_b64 = encode_image(MASK_PATH)
    
    async with httpx.AsyncClient() as client:
        tasks = [
            send_request(client, i, image_b64, mask_b64)
            for i in range(CONCURRENT_REQUESTS)
        ]
        await asyncio.gather(*tasks)

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