stitching.py

from typing import Tuple
from PIL import Image, ImageDraw, ImageFont, ImageEnhance
from io import BytesIO
import boto3
s3_client = boto3.client("s3")


STYLE_CONFIG = {
    "default": {
        "patch": {
            "background_color": (10, 10, 10, 255),
            "text_color": (255, 255, 255, 255),
            "opacity": 1,
            "top_height": 100,
            "bottom_height": 300,
        },
        "fade": {"color": (10, 10, 10, 255), "strength": 1.7},
        "image": {"contrast": 1, "brightness": 1, "saturation": 1},
    }
}


def wrap_text(
    text: str, font: ImageFont.FreeTypeFont, max_width: int, padding: int = 60
) -> list:
    """Wrap text to fit within the specified width."""
    max_width -= 2 * padding
    words = text.split()
    lines, current_line = [], []

    for word in words:
        test_line = " ".join(current_line + [word])
        text_width = ImageDraw.Draw(
            Image.new("RGBA", (max_width + 2 * padding, 1000))
        ).textbbox((0, 0), test_line, font)[2]

        if text_width <= max_width:
            current_line.append(word)
        else:
            lines.append(" ".join(current_line))
            current_line = [word]

    if current_line:
        lines.append(" ".join(current_line))
    return lines


def apply_style_effects(
    image: Image.Image, visual_style: str = "default"
) -> Image.Image:
    """Apply image style effects such as contrast, brightness, and saturation."""
    style_config = STYLE_CONFIG.get(visual_style, STYLE_CONFIG["default"])
    enhancers = [
        (ImageEnhance.Contrast, style_config["image"]["contrast"]),
        (ImageEnhance.Brightness, style_config["image"]["brightness"]),
        (ImageEnhance.Color, style_config["image"]["saturation"]),
    ]

    for enhancer_class, factor in enhancers:
        image = enhancer_class(image).enhance(factor)

    return image.convert("RGBA")


def create_fade_mask(
    size: Tuple[int, int], visual_style: str = "default"
) -> Image.Image:
    """Create a fade mask with a gradient effect."""
    width, height = size
    style_config = STYLE_CONFIG.get(visual_style, STYLE_CONFIG["default"])
    fade_height = int(height * 0.25)

    mask = Image.new("L", size, 255)
    gradient_data = [
        int(255 * (i / fade_height) * style_config["fade"]["strength"])
        for i in range(fade_height)
    ]

    gradient = Image.new("L", (1, fade_height))
    gradient.putdata(gradient_data)
    gradient = gradient.resize((width, fade_height))

    mask.paste(gradient, (0, 0))
    mask.paste(gradient.transpose(Image.FLIP_TOP_BOTTOM), (0, height - fade_height))

    return mask


def create_text_patch(
    text: str, width: int, height: int, font_path: str, visual_style: str
) -> Image.Image:
    """Create a text patch with wrapping and styling."""
    style_config = STYLE_CONFIG.get(visual_style, STYLE_CONFIG["default"])
    patch = Image.new(
        "RGBA", (width, height), style_config["patch"]["background_color"]
    )
    draw = ImageDraw.Draw(patch)

    font_size = 40
    font = None

    while font_size > 10:
        try:
            font = ImageFont.truetype(font_path, font_size)
            lines = wrap_text(text, font, width - 40, 60)
            font_height = font.getbbox("A")[3] - font.getbbox("A")[1]
            total_text_height = len(lines) * (font_height + 10)

            if total_text_height <= height - 40 and len(lines) > 0:
                break
        except Exception:
            pass
        font_size -= 2

    font = font or ImageFont.load_default()
    lines = wrap_text(text, font, width - 40, 60)
    font_height = font.getbbox("A")[3] - font.getbbox("A")[1]
    total_text_height = len(lines) * (font_height + 10)

    y = (height - total_text_height) // 2

    for line in lines:
        bbox = draw.textbbox((0, 0), line, font=font)
        text_width = bbox[2] - bbox[0]
        x = (width - text_width) // 2
        draw.text((x, y), line, font=font, fill=style_config["patch"]["text_color"])
        y += font_height + 10

    return patch


def process_single_image(
    image_data: bytes,
    text: str,
    font_path: str,
    target_width: int = 1024,
    target_height: int = 1024,
    visual_style: str = "default",
) -> BytesIO:
    """Process a single image with text overlay and styling."""
    img = Image.open(BytesIO(image_data))
    style_config = STYLE_CONFIG.get(visual_style, STYLE_CONFIG["default"])
    top_height, bottom_height = (
        style_config["patch"]["top_height"],
        style_config["patch"]["bottom_height"],
    )

    # Process the main image
    processed = apply_style_effects(img, visual_style).resize(
        (target_width, target_height)
    )

    # Apply fade effect
    mask = create_fade_mask(processed.size, visual_style)
    background = Image.new("RGBA", processed.size, style_config["fade"]["color"])
    processed = Image.composite(processed, background, mask)

    # Create text patches
    bottom_patch = create_text_patch(
        text, target_width, bottom_height, font_path, visual_style
    )
    top_patch = Image.new(
        "RGBA", (target_width, top_height), style_config["patch"]["background_color"]
    )

    # Combine all elements
    final_image = Image.new(
        "RGBA", (target_width, target_height + top_height + bottom_height)
    )
    final_image.paste(top_patch, (0, 0))
    final_image.paste(processed, (0, top_height))
    final_image.paste(bottom_patch, (0, top_height + target_height))

    # Save to BytesIO
    output_buffer = BytesIO()
    final_image.save(output_buffer, format="PNG")
    output_buffer.seek(0)

    return output_buffer