import os
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor
import multiprocessing
import subprocess
import cv2
from PIL import Image

# 이미지 파일 경로 설정
input_dir = '/Users/gohyunjun/Downloads/unprocessed_1007'  # 원본 이미지 폴더 경로
output_dir = '/Users/gohyunjun/Downloads/unprocessed_1007_optim'  # 최적화된 이미지 저장 경로

# 디렉토리가 없으면 생성
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 이미지 리사이즈 함수 (비율 유지, 라코스 보간법 사용)
def resize_image_with_aspect_ratio_cv(img, target_width, target_height):
    original_height, original_width = img.shape[:2]
    ratio = min(target_width / original_width, target_height / original_height)
    new_width = int(original_width * ratio)
    new_height = int(original_height * ratio)
    return cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_LANCZOS4)

# pngquant로 압축하는 함수
def compress_with_pngquant(output_path):
    try:
        # pngquant를 호출하여 이미지 압축 ( --force로 기존 파일 덮어쓰기, --ext로 확장자 지정)
        subprocess.run(['pngquant', '--force', '--ext', '.png', '--quality', '65-80', output_path], check=True)
        print(f"Compressed with pngquant: {output_path}")
    except subprocess.CalledProcessError as e:
        print(f"An error occurred with pngquant on {output_path}: {e}")

# 이미지 처리 함수
def compress_and_resize_image(args):
    input_path, output_path = args
    try:
        # OpenCV로 이미지 읽기
        img = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)

        # 이미지 리사이즈 (비율 유지)
        img = resize_image_with_aspect_ratio_cv(img, 800, 450)

        # OpenCV에서 읽은 이미지를 RGB로 변환
        img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

        # PNG로 저장 (원본 파일, 나중에 pngquant로 압축)
        img_pil.save(output_path, format='PNG', optimize=True)

        # pngquant로 이미지 압축
        compress_with_pngquant(output_path)

    except Exception as e:
        print(f"An error occurred with {input_path}: {e}")

if __name__ == "__main__":
    # 입력 디렉토리의 모든 PNG 파일에 대해 실행
    image_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.png')]

    # 입력 및 출력 경로 생성
    tasks = []
    for filename in image_files:
        input_path = os.path.join(input_dir, filename)
        output_path = os.path.join(output_dir, filename)
        tasks.append((input_path, output_path))

    # 워커 수를 코어 수의 2배로 설정
    num_workers = multiprocessing.cpu_count() * 2

    # 프로세스 풀 생성 및 작업 병렬 실행
    with ProcessPoolExecutor(max_workers=num_workers) as executor:
        list(tqdm(executor.map(compress_and_resize_image, tasks), total=len(tasks), desc="Processing images"))

    print('All images have been processed.')