from PIL import Image
from io import BytesIO


def tile_image_row(image):

    width, height = image.size
    new_width = width * 2
    new_height = height

    result = Image.new(image.mode, (new_width, new_height))

    result.paste(image, (0, 0))
    result.paste(image, (width, 0))

    return result

def tile_image_column(image):

    width, height = image.size
    new_width = width
    new_height = height * 2
    
    result = Image.new(image.mode, (new_width, new_height))

    result.paste(image, (0, 0))
    result.paste(image, (0, height))

    return result

def tile_image_grid(image):

    temp_image = tile_image_column(image)
    result = tile_image_row(temp_image)

    return result

def tile_image_vert_brick(image):

    width, height = image.size
    
    new_width = width * 2
    new_height = height *2

    result = Image.new(image.mode, (new_width, new_height))

    # First column: Doubled image
    result.paste(image, (0, 0))
    result.paste(image, (0, height))

    # Second column: Shifted doubled image
    top_half = image.crop((0, 0, width, height // 2))
    bottom_half = image.crop((0, height // 2, width, height))
    shifted_image = Image.new(image.mode, (width, height*2))
    shifted_image.paste(bottom_half, (0, 0))
    shifted_image.paste(top_half, (0, height+height // 2))
    shifted_image.paste(image, (0, height // 2))
    result.paste(shifted_image, (width, 0)) 

    return result

def tile_image_horiz_brick(image):

    width, height = image.size
    
    new_width = width * 2
    new_height = height *2

    result = Image.new(image.mode, (new_width, new_height))

    # First column: Doubled image
    result.paste(image, (0, 0))
    result.paste(image, (width, 0))

    # Second column: Shifted doubled image
    left_half = image.crop((0, 0, width//2, height))
    right_half = image.crop((width//2, 0, width, height))
    shifted_image = Image.new(image.mode, (width*2, height))
    shifted_image.paste(right_half, (0, 0))
    shifted_image.paste(left_half, (width+width // 2, 0))
    shifted_image.paste(image, (width//2, 0))
    result.paste(shifted_image, (0, height)) 

    return result


def crop_along_line_to_gif(img, width, height, start_x, start_y, end_x, end_y, num_crops=200, duration=80, loop=0, jpeg_quality=85, scale_factor=0.5):
    # Convert to JPEG in-memory if needed
    if jpeg_quality is not None:
        img = img.convert("RGB")  # Ensure RGB mode
        with BytesIO() as buffer:
            img.save(buffer, format="JPEG", quality=jpeg_quality)
            buffer.seek(0)
            img = Image.open(buffer)
            img.load()

    # Apply scaling factor to image dimensions and coordinates
    if scale_factor != 1.0:
        img = img.resize((int(img.width * scale_factor), int(img.height * scale_factor)), Image.Resampling.LANCZOS)
        width = int(width * scale_factor)
        height = int(height * scale_factor)
        start_x = int(start_x * scale_factor)
        start_y = int(start_y * scale_factor)
        end_x = int(end_x * scale_factor)
        end_y = int(end_y * scale_factor)

    img_width, img_height = img.size
    x_step = (end_x - start_x) / (num_crops - 1)
    y_step = (end_y - start_y) / (num_crops - 1)

    cropped_images = []
    for i in range(num_crops):
        x = start_x + i * x_step
        y = start_y + i * y_step
        
        if 0 <= x < img_width and 0 <= y < img_height:
            # Calculate crop area coordinates
            left = x
            top = y
            right = x + width
            bottom = y + height

            cropped_img = img.crop((left, top, right, bottom))  

            # Apply scaling if scale_factor is not 1.0


            cropped_images.append(cropped_img)

    # Create a BytesIO buffer to store the GIF
    gif_buffer = BytesIO()

    # Save the GIF to the buffer
    first_image = cropped_images[0]
    first_image.info["duration"] = duration
    first_image.save(
    gif_buffer,
    save_all=True,
    append_images=cropped_images[1:],
    loop=loop,
    format="GIF"  # Explicitly specify GIF format
    )

    # Rewind the buffer to the beginning
    gif_buffer.seek(0)

    # Return the buffer containing the GIF
    return gif_buffer

def tile_image(image, tile_type="grid"):
    """Tiles an image based on the specified type using Pillow.

    Args:
        image: The image to tile (PIL Image object).
        tile_type: The type of tiling ("row", "column", "grid", "vertical_brick").

    Returns:
        The tiled image (PIL Image object).
    """
    if tile_type == "row":
        return tile_image_row(image)
    elif tile_type == "column":
        return tile_image_column(image)
    elif tile_type == 'grid':  # Grid
        return tile_image_grid(image)
    elif tile_type == "vertical_brick":
        return tile_image_vert_brick(image)
    elif tile_type == "horizontal_brick":
        return tile_image_horiz_brick(image)
    
def tile_and_convert(image, tile_type='grid'):
    img_width, img_height = image.size

    assert img_width < 2600, "Please only upload 1x1, 2x1, 1x2 or 2x2 tiles downloaded from ideogram.ai"
    if img_width > 1280:
        scale_factor = 0.23
    else:
        scale_factor = 0.46

    tiled_image = tile_image(image, tile_type)

    start_x, start_y, end_x, end_y = 0,0,0,0 # default values
    if tile_type == 'row':
        end_x = img_width
    elif tile_type == 'column':
        end_y = img_height
    else:
        end_x = img_width
        end_y = img_height

    return crop_along_line_to_gif(tiled_image, img_width, img_height, start_x, start_y, end_x, end_y, jpeg_quality=80, scale_factor=scale_factor)