import json
import os
import shutil
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from collections import defaultdict

def copy_single_image(src_dst_tuple):
    """
    단일 이미지 파일을 복사하는 함수.

    Parameters:
        src_dst_tuple (tuple): (소스 경로, 목적지 경로) 튜플.

    Returns:
        str or None: 에러 메시지 또는 None.
    """
    src_path, dst_path = src_dst_tuple
    try:
        if not os.path.exists(src_path):
            return f"Warning: Source image does not exist: {src_path}"
        
        # 생성할 디렉토리가 없으면 생성
        os.makedirs(os.path.dirname(dst_path), exist_ok=True)
        
        # 이미 존재하는 파일은 건너뛰기
        if os.path.exists(dst_path):
            return f"Skipping already exists: {dst_path}"
        
        shutil.copy2(src_path, dst_path)
        return None
    except Exception as e:
        return f"Error copying {src_path} to {dst_path}: {e}"

def copy_images(json_path, images_src_dir, images_dst_dir, max_workers=8):
    """
    JSON 파일에 명시된 이미지들을 소스 디렉토리에서 목적지 디렉토리로 병렬로 복사하는 함수.

    Parameters:
        json_path (str): JSON 파일 경로.
        images_src_dir (str): 원본 이미지 디렉토리 경로.
        images_dst_dir (str): 목적지 이미지 디렉토리 경로.
        max_workers (int): 병렬 처리할 스레드 수.
    """
    with open(json_path, 'r') as f:
        coco = json.load(f)
    
    images = coco['images']
    total = len(images)
    print(f"Total images to copy from {json_path}: {total}")
    
    # 소스와 목적지 경로 목록 생성
    src_dst_list = []
    for img in images:
        file_name = img['file_name']
        src_path = os.path.join(images_src_dir, file_name)
        dst_path = os.path.join(images_dst_dir, file_name)
        src_dst_list.append((src_path, dst_path))
    
    # 병렬 복사 시작
    copied = 0
    skipped = 0
    missing_files = []
    errors = []
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(copy_single_image, paths): paths for paths in src_dst_list}
        for idx, future in enumerate(as_completed(futures), 1):
            result = future.result()
            if result:
                if result.startswith("Warning"):
                    missing_files.append(result)
                elif result.startswith("Skipping"):
                    skipped += 1
                else:
                    errors.append(result)
            else:
                copied += 1
            
            # 진행 상황 출력
            if idx % 100 == 0 or idx == total:
                print(f"Processed {idx}/{total} images")
    
    # 복사 완료 후 요약 정보 출력
    print("\n=== 복사 요약 ===")
    print(f"총 처리된 이미지: {total}")
    print(f"복사된 이미지: {copied}")
    print(f"이미 이미 존재하여 건너뛴 이미지: {skipped}")
    print(f"누락된 이미지: {len(missing_files)}")
    if missing_files:
        print("누락된 이미지 목록:")
        for warning in missing_files:
            print(f"  - {warning}")
    if errors:
        print(f"복사 중 에러 발생한 파일 수: {len(errors)}")
        for error in errors:
            print(f"  - {error}")

def main():
    """
    메인 함수: train.json과 val.json을 기반으로 이미지를 복사합니다.
    """
    # 데이터 루트 경로 설정
    data_root = '/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl_all_combined/'
    
    # 원본 이미지 디렉토리
    images_src_dir = os.path.join(data_root, 'images')
    #images_src_dir = os.path.join(data_root, 'train_original')
    
    # JSON 파일 경로
    train_json = os.path.join(data_root, 'train.json')
    val_json = os.path.join(data_root, 'val.json')
    
    # 목적지 이미지 디렉토리
    train_images_dst_dir = os.path.join(data_root, 'train', 'images')
    val_images_dst_dir = os.path.join(data_root, 'val', 'images')
    
    print("=== Train 이미지 복사 시작 ===")
    copy_images(train_json, images_src_dir, train_images_dst_dir, max_workers=16)
    print("=== Train 이미지 복사 완료 ===\n")
    
    print("=== Val 이미지 복사 시작 ===")
    copy_images(val_json, images_src_dir, val_images_dst_dir, max_workers=16)
    print("=== Val 이미지 복사 완료 ===\n")
    
    print("이미지 분리가 완료되었습니다.")

if __name__ == "__main__":
    main()