import pygame
import os
import shutil

def create_spritesheet(image_paths, output_path, rows=1, padding=0):
    """
    Combine multiple images into a single spritesheet.
    
    Args:
        image_paths (list): List of paths to source images
        output_path (str): Path to save the combined spritesheet
        rows (int): Number of rows in the spritesheet
        padding (int): Pixel padding between images
    """
    if not image_paths:
        raise ValueError("No images provided")
    
    # Load all images
    images = []
    max_width = 0
    max_height = 0
    
    for path in image_paths:
        try:
            img = pygame.image.load(path)#.convert_alpha()
            images.append(img)
            max_width = max(max_width, img.get_width())
            max_height = max(max_height, img.get_height())
        except pygame.error as e:
            print(f"Failed to load {path}: {e}")
            continue
    
    if not images:
        raise ValueError("No valid images loaded")
    
    # Calculate spritesheet dimensions
    cols = (len(images) + rows - 1) // rows
    sheet_width = cols * (max_width + padding) - padding
    sheet_height = rows * (max_height + padding) - padding
    
    # Create blank spritesheet surface
    spritesheet = pygame.Surface((sheet_width, sheet_height), pygame.SRCALPHA)
    
    # Blit images onto spritesheet
    for i, img in enumerate(images):
        row = i // cols
        col = i % cols
        x = col * (max_width + padding)
        y = row * (max_height + padding)
        spritesheet.blit(img, (x, y))
    
    # Save the spritesheet
    print(f"dir: {os.getcwd()}")
    pygame.image.save(spritesheet, output_path)
    print(f"Spritesheet saved to {output_path}")

    return max_width, max_height




pygame.init()

tool_root_dir = os.getcwd()

# Because os.rmdir() only works on empty directories
if os.path.isdir(f"{tool_root_dir}/output"): shutil.rmtree(f"{tool_root_dir}/output")
os.mkdir(f"{tool_root_dir}/output")

spritesheet_names = os.listdir(f"{tool_root_dir}/input")

for spritesheet_name in spritesheet_names:

    source_dir=f"{tool_root_dir}/input/{spritesheet_name}"
    target_dir=f"{tool_root_dir}/output/"

    filenames = os.listdir(f"{source_dir}")
    
    # List of image paths to combine
    image_files = filenames

    cols=8
    rows=len(image_files)//cols + 1

    print(f"Dir: {os.getcwd()}")
    print(f"Target: {target_dir}")

    ## Create and save spritesheet
    os.chdir(source_dir)
    (tile_width, tile_height)  = create_spritesheet(image_files, f"{target_dir}/{spritesheet_name}_spritesheet.png", rows, padding=1)
    spritesheet_filename = f"{target_dir}/{spritesheet_name}_{tile_width}x{tile_height}_spritesheet.png"
    os.rename(f"{target_dir}/{spritesheet_name}_spritesheet.png", spritesheet_filename)

    ## Create+Store metadata
    mapping = ["x_col,y_col,tilename"]

    for index, filename in enumerate(image_files):
        #x_col = index // cols
        #y_row = index - x_col * cols
        y_row = index // cols
        x_col = index - y_row * cols

        print(f"[{spritesheet_name}] : {index} - [{x_col},{y_row}] -> {filename}")
        tilename = filename.split(".")[0]
        mapping.append(  f"{x_col},{y_row},{tilename}" )

    print(f"Wrote {spritesheet_filename}")


    mapping_filename = f"{target_dir}/{spritesheet_name}_{tile_width}x{tile_height}_mapping.txt"
    f = open(mapping_filename,"w")
    f.write("\n".join(mapping))
    f.close()
    print(f"Wrote {mapping_filename}")

    os.chdir(f"{tool_root_dir}")