from PIL import Image, ImageDraw
from pathlib import Path

# 설정
WIDTH, HEIGHT = 800, 1000
IMAGE_FILENAME = "test_image.png"
MASK_FILENAME = "test_mask.png"
SAVE_DIR = Path(__file__).parent

def generate_images():
    """테스트용 이미지와 마스크를 생성합니다."""
    
    # 1. 테스트 이미지 생성 (간단한 그라데이션)
    print(f"'{IMAGE_FILENAME}' 생성 중 ({WIDTH}x{HEIGHT})...")
    image = Image.new("RGB", (WIDTH, HEIGHT))
    draw = ImageDraw.Draw(image)
    
    for i in range(HEIGHT):
        r = int(255 * (i / HEIGHT))
        g = int(255 * (1 - i / HEIGHT))
        b = 128
        draw.line([(0, i), (WIDTH, i)], fill=(r, g, b))
        
    image_path = SAVE_DIR / IMAGE_FILENAME
    image.save(image_path)
    print(f"✅ 이미지를 성공적으로 저장했습니다: {image_path}")

    # 2. 마스크 이미지 생성 (중앙에 20% 크기의 사각형)
    print(f"'{MASK_FILENAME}' 생성 중 ({WIDTH}x{HEIGHT})...")
    mask = Image.new("L", (WIDTH, HEIGHT), 0)  # 검은색 배경
    draw = ImageDraw.Draw(mask)
    
    # 중앙에 약 20% 크기의 흰색 사각형 그리기
    mask_width = int(WIDTH * 0.45) # 800 * 0.45 = 360
    mask_height = int(HEIGHT * 0.45) # 1000 * 0.45 = 450
    # 360 * 450 = 162,000. Total 800,000. 162000/800000 = 0.2025 (약 20%)
    
    x1 = (WIDTH - mask_width) // 2
    y1 = (HEIGHT - mask_height) // 2
    x2 = x1 + mask_width
    y2 = y1 + mask_height
    
    draw.rectangle([x1, y1, x2, y2], fill=255) # 흰색
    
    mask_path = SAVE_DIR / MASK_FILENAME
    mask.save(mask_path)
    print(f"✅ 마스크를 성공적으로 저장했습니다: {mask_path}")


if __name__ == "__main__":
    generate_images()