#!/usr/bin/env python3
"""
테스트용 샘플 이미지 생성 스크립트
간단한 테스트 이미지와 마스크를 생성합니다.
"""
import os
import sys
import numpy as np
from PIL import Image, ImageDraw
import cv2

# 프로젝트 루트를 Python 경로에 추가
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

def create_test_image(width=512, height=512):
    """테스트용 샘플 이미지 생성"""
    # 그라데이션 배경
    image = np.zeros((height, width, 3), dtype=np.uint8)
    
    # 수평 그라데이션 (파란색에서 빨간색)
    for x in range(width):
        ratio = x / width
        blue = int(255 * (1 - ratio))
        red = int(255 * ratio)
        image[:, x] = [blue, 0, red]
    
    # 중앙에 원 그리기
    center_x, center_y = width // 2, height // 2
    radius = min(width, height) // 4
    
    # 원 그리기
    cv2.circle(image, (center_x, center_y), radius, (0, 255, 0), -1)
    
    # 텍스트 추가
    cv2.putText(image, "TEST", (center_x - 50, center_y + 10), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
    
    return image

def create_test_mask(width=512, height=512):
    """테스트용 마스크 생성 (중앙 원형 영역)"""
    mask = np.zeros((height, width), dtype=np.uint8)
    
    center_x, center_y = width // 2, height // 2
    radius = min(width, height) // 4
    
    # 중앙 원형 영역을 마스크로 설정 (흰색)
    cv2.circle(mask, (center_x, center_y), radius, 255, -1)
    
    return mask

def create_complex_test_image(width=512, height=512):
    """복잡한 테스트 이미지 생성"""
    image = np.zeros((height, width, 3), dtype=np.uint8)
    
    # 체크무늬 패턴
    square_size = 32
    for y in range(0, height, square_size):
        for x in range(0, width, square_size):
            if (x // square_size + y // square_size) % 2 == 0:
                image[y:y+square_size, x:x+square_size] = [100, 100, 100]
            else:
                image[y:y+square_size, x:x+square_size] = [200, 200, 200]
    
    # 여러 도형 추가
    # 사각형
    cv2.rectangle(image, (100, 100), (200, 200), (0, 255, 0), -1)
    
    # 삼각형
    pts = np.array([[300, 100], [350, 200], [250, 200]], np.int32)
    cv2.fillPoly(image, [pts], (255, 0, 0))
    
    # 타원
    cv2.ellipse(image, (400, 150), (50, 30), 45, 0, 360, (0, 255, 255), -1)
    
    return image

def create_complex_mask(width=512, height=512):
    """복잡한 마스크 생성"""
    mask = np.zeros((height, width), dtype=np.uint8)
    
    # 사각형 마스크
    cv2.rectangle(mask, (100, 100), (200, 200), 255, -1)
    
    # 삼각형 마스크
    pts = np.array([[300, 100], [350, 200], [250, 200]], np.int32)
    cv2.fillPoly(mask, [pts], 255)
    
    # 타원 마스크
    cv2.ellipse(mask, (400, 150), (50, 30), 45, 0, 360, 255, -1)
    
    return mask

def save_test_images():
    """테스트 이미지들을 저장합니다."""
    # 현재 스크립트 위치에서 상대 경로로 테스트 디렉토리 찾기
    script_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.dirname(script_dir)
    images_dir = os.path.join(project_root, "images")
    
    # 테스트 디렉토리 생성
    os.makedirs(images_dir, exist_ok=True)
    
    print(f"🖼️ 테스트 이미지 생성 중...")
    print(f"   저장 위치: {images_dir}")
    
    # 1. 기본 테스트 이미지 (512x512)
    print("  - 기본 테스트 이미지 생성...")
    test_image = create_test_image(512, 512)
    test_mask = create_test_mask(512, 512)
    
    Image.fromarray(cv2.cvtColor(test_image, cv2.COLOR_BGR2RGB)).save(
        os.path.join(images_dir, "test_image_512.png")
    )
    Image.fromarray(test_mask).save(
        os.path.join(images_dir, "test_mask_512.png")
    )
    
    # 2. 복잡한 테스트 이미지 (512x512)
    print("  - 복잡한 테스트 이미지 생성...")
    complex_image = create_complex_test_image(512, 512)
    complex_mask = create_complex_mask(512, 512)
    
    Image.fromarray(cv2.cvtColor(complex_image, cv2.COLOR_BGR2RGB)).save(
        os.path.join(images_dir, "complex_image_512.png")
    )
    Image.fromarray(complex_mask).save(
        os.path.join(images_dir, "complex_mask_512.png")
    )
    
    # 3. 작은 테스트 이미지 (256x256)
    print("  - 작은 테스트 이미지 생성...")
    small_image = create_test_image(256, 256)
    small_mask = create_test_mask(256, 256)
    
    Image.fromarray(cv2.cvtColor(small_image, cv2.COLOR_BGR2RGB)).save(
        os.path.join(images_dir, "test_image_256.png")
    )
    Image.fromarray(small_mask).save(
        os.path.join(images_dir, "test_mask_256.png")
    )
    
    # 4. 큰 테스트 이미지 (1024x1024)
    print("  - 큰 테스트 이미지 생성...")
    large_image = create_test_image(1024, 1024)
    large_mask = create_test_mask(1024, 1024)
    
    Image.fromarray(cv2.cvtColor(large_image, cv2.COLOR_BGR2RGB)).save(
        os.path.join(images_dir, "test_image_1024.png")
    )
    Image.fromarray(large_mask).save(
        os.path.join(images_dir, "large_mask_1024.png")
    )
    
    print(f"✅ 테스트 이미지 생성 완료!")
    print(f"   - 기본 이미지: test_image_512.png, test_mask_512.png")
    print(f"   - 복잡한 이미지: complex_image_512.png, complex_mask_512.png")
    print(f"   - 작은 이미지: test_image_256.png, test_mask_256.png")
    print(f"   - 큰 이미지: test_image_1024.png, large_mask_1024.png")

if __name__ == "__main__":
    save_test_images()