from concurrent.futures import ProcessPoolExecutor, as_completed
import os
from typing import Optional
from tqdm import tqdm
from src.logger import get_logger
from PIL import Image, ImageFile

logger = get_logger(__name__)

ImageFile.LOAD_TRUNCATED_IMAGES = False

def resize_image(
    input_path: str,
    output_path: str,
    max_width: int = 3840,
    max_height: int = 2160,
    max_pixels: Optional[int] = None,
):
    with Image.open(input_path) as img:
        img.verify()

    with Image.open(input_path) as img:
        width, height = img.size
        if max_pixels is not None and (width * height) > max_pixels:
            scaling_factor = (max_pixels / (width * height)) ** 0.5
            new_width = int(width * scaling_factor)
            new_height = int(height * scaling_factor)
        else:
            scaling_factor = min(max_width / width, max_height / height, 1)
            new_width = int(width * scaling_factor)
            new_height = int(height * scaling_factor)

        if (new_width, new_height) != (width, height):
            img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
            img.save(output_path)
        else:
            img.save(output_path)

def process_image(args_tuple):
    image_path, output_path, max_width, max_height, max_pixels = args_tuple
    try:
        resize_image(
            image_path,
            output_path,
            max_width=max_width,
            max_height=max_height,
            max_pixels=max_pixels,
        )
        return None
    except Exception as e:
        return (image_path, str(e))


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_images_dir", type=str, required=True)
    parser.add_argument("--output_images_dir", type=str, required=True)
    parser.add_argument("--max_width_or_height", type=int, default=2560)
    parser.add_argument("--max_pixels", type=int, default=None)
    parser.add_argument("--workers", type=int, default=16)
    args = parser.parse_args()

    logger.info(f"Reading images from {args.input_images_dir}")

    tasks = []
    for root, _, files in os.walk(args.input_images_dir):
        for filename in files:
            if not filename.lower().endswith((".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tiff")):
                continue
            
            image_path = os.path.join(root, filename)
            relative_path = os.path.relpath(image_path, args.input_images_dir)
            output_path = os.path.join(args.output_images_dir, relative_path)
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            tasks.append((image_path, output_path, args.max_width_or_height, args.max_width_or_height, args.max_pixels))

    with ProcessPoolExecutor(max_workers=args.workers) as executor:
        futures = [executor.submit(process_image, task) for task in tqdm(tasks, desc="Scheduling image processing")]
        for f in tqdm(as_completed(futures), total=len(futures), desc="Resizing images"):
            result = f.result()
            if result:
                image_path, error = result
                logger.error(f"Error processing {image_path}: {error}")

    
