main.py

import os
import tempfile
from io import BytesIO
from pathlib import Path
from typing import List, Literal, Optional, Tuple, Union
from urllib.parse import urlparse

import cairosvg
import numpy as np
import pygments
import requests
from PIL import Image, ImageDraw, ImageFont
from pygments.formatters import ImageFormatter
from pygments.lexers import PythonLexer
from pygments.styles import get_style_by_name


def create_gradient_background(
    width: int,
    height: int,
    start_color: Tuple[int, int, int],
    end_color: Tuple[int, int, int],
    frame_num: int = 0,
) -> Image.Image:
    """Create animated gradient background with wave effects."""
    # Pre-calculate wave patterns using numpy operations
    x = np.arange(width)
    y = np.arange(height)
    X, Y = np.meshgrid(x, y)

    wave1 = 30 * np.sin(Y * 0.02 + frame_num * 0.1)
    wave2 = 15 * np.sin(Y * 0.03 - frame_num * 0.05)
    wave3 = 10 * np.cos(X * 0.02 + frame_num * 0.08)
    wave = wave1 + wave2 + wave3

    # Vectorized calculation of progress values
    base_progress = Y / height
    wave_offset = wave / height
    progress = np.clip(base_progress + wave_offset, 0, 1)

    # Vectorized color interpolation
    r = np.clip(
        (start_color[0] + (end_color[0] - start_color[0]) * progress), 0, 255
    ).astype(np.uint8)
    g = np.clip(
        (start_color[1] + (end_color[1] - start_color[1]) * progress), 0, 255
    ).astype(np.uint8)
    b = np.clip(
        (start_color[2] + (end_color[2] - start_color[2]) * progress), 0, 255
    ).astype(np.uint8)

    # Create RGB array and convert to image
    rgb_array = np.stack([r, g, b], axis=2)
    return Image.fromarray(rgb_array)


def create_window_background(
    width: int,
    height: int,
    style_name: str,
    filename: Optional[str] = None,
    font_size: int = 24,
) -> Image.Image:
    """Create window background with title bar and control buttons."""
    # Get background color from style
    style_obj = get_style_by_name(style_name)
    bg_color = style_obj.background_color
    if bg_color.startswith("#"):
        bg_color = tuple(int(bg_color[i : i + 2], 16) for i in (1, 3, 5))
    else:
        bg_color = (40, 40, 40)  # Default dark background

    window = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    draw = ImageDraw.Draw(window)

    title_bar_height = 40
    radius = 10
    draw.rounded_rectangle([(0, 0), (width, height)], radius, fill=bg_color)

    # Draw window control buttons
    circle_y = title_bar_height // 2
    draw.ellipse((20, circle_y - 6, 32, circle_y + 6), fill=(255, 95, 87))
    draw.ellipse((40, circle_y - 6, 52, circle_y + 6), fill=(255, 189, 46))
    draw.ellipse((60, circle_y - 6, 72, circle_y + 6), fill=(39, 201, 63))

    if filename:
        try:
            font = ImageFont.truetype("Arial Bold", int(font_size * 0.5))
        except:
            font = ImageFont.load_default()
        text_width = draw.textlength(filename, font=font)
        text_x = (width - text_width) // 2
        draw.text((text_x, circle_y - 6), filename, fill=(200, 200, 200), font=font)

    return window


def load_and_resize_image(
    image_path: str, target_width: int, window_padding: int
) -> Tuple[Optional[Image.Image], int]:
    """Load and resize an image from path or URL."""
    try:
        if urlparse(image_path).scheme in ("http", "https"):
            response = requests.get(image_path)
            img = Image.open(BytesIO(response.content))
        else:
            img = Image.open(image_path)

        # Scale image to match target width
        aspect = img.width / img.height
        new_width = target_width - 2 * window_padding
        new_height = int(new_width / aspect)
        img = img.resize((new_width, new_height))

        # Add padding
        height_with_padding = img.height + 2 * window_padding
        padded_img = Image.new(
            "RGBA",
            (img.width + 2 * window_padding, height_with_padding),
            (0, 0, 0, 0),
        )
        padded_img.paste(img, (window_padding, window_padding))

        return padded_img, height_with_padding
    except Exception as e:
        print(f"Warning: Could not load image: {e}")
        return None, 0


def create_code_gif(
    code: Union[str, Path],
    output_file: str | None = None,
    style: str = "monokai",
    font_size: int = 24,
    start_delay: float = 0.5,
    end_delay: float = 1.0,
    acceleration: float = 0.8,
    line_numbers: bool = True,
    gradient_start: Tuple[int, int, int] = (45, 49, 66),
    gradient_end: Tuple[int, int, int] = (239, 129, 132),
    title: Optional[str] = None,
    filename: Optional[str] = None,
    favicon: Optional[str] = None,
    photo: Optional[str] = None,
    photo_position: Literal["above", "below"] = "above",
    comments: Optional[List[str]] = None,
    comments_position: Literal["above", "below"] = "above",
    aspect_ratio: float = 16 / 9,
) -> None:
    """Creates an animated GIF of code being typed out with increasing speed."""

    # Initialize constants
    min_padding = 20
    window_padding = 20
    title_bar_height = 40
    title_font_size = int(font_size * 2.5)
    comment_font_size = int(font_size * 2)
    title_height = title_font_size * 2 if title else 0
    comment_height = comment_font_size * 2 if comments else 0

    # Load and process code
    code_str = code.read_text() if isinstance(code, Path) else code

    # Set up syntax highlighting
    lexer = PythonLexer()
    formatter = ImageFormatter(
        style=style, line_numbers=line_numbers, font_size=font_size
    )

    # Get code dimensions
    with tempfile.NamedTemporaryFile(suffix=".png") as tmp:
        tmp.write(pygments.highlight(code_str, lexer, formatter))
        tmp.flush()
        with Image.open(tmp.name) as img:
            code_width, code_height = img.size

    total_width = code_width + 2 * window_padding
    final_height = code_height + title_bar_height + 2 * window_padding

    # Load photo if provided
    photo_img, photo_height = (
        load_and_resize_image(photo, total_width, window_padding)
        if photo
        else (None, 0)
    )

    # Calculate background dimensions
    content_height = (
        title_height + photo_height + comment_height + final_height + 4 * min_padding
    )
    background_width = max(
        total_width + 2 * min_padding, int(content_height * aspect_ratio)
    )
    background_height = content_height
    window_x = (background_width - total_width) // 2

    # Load favicon
    logo = None
    if favicon:
        logo_size = int(min(background_width, background_height) * 0.1)
        try:
            if favicon.lower().endswith(".svg"):
                if urlparse(favicon).scheme in ("http", "https"):
                    response = requests.get(favicon)
                    png_data = cairosvg.svg2png(
                        bytestring=response.content,
                        output_width=logo_size,
                        output_height=logo_size,
                    )
                else:
                    png_data = cairosvg.svg2png(
                        url=favicon,
                        output_width=logo_size,
                        output_height=logo_size,
                    )
                logo = Image.open(BytesIO(png_data))
            else:
                logo, _ = load_and_resize_image(favicon, logo_size, 0)
                if logo:
                    logo.thumbnail((logo_size, logo_size))
        except Exception as e:
            print(f"Warning: Could not load favicon: {e}")

    # Pre-create fonts
    try:
        title_font = ImageFont.truetype("Arial Bold", title_font_size)
        comment_font = ImageFont.truetype("Arial", comment_font_size)
    except:
        title_font = comment_font = ImageFont.load_default()

    # Generate frames
    frames = []
    with tempfile.TemporaryDirectory() as tmpdir:
        code_lines = code_str.split("\n")
        num_frames = len(code_lines)
        frames_per_comment = num_frames // len(comments) if comments else 0

        # Pre-create window background
        window = create_window_background(
            total_width, final_height, style, filename, font_size
        )

        for i in range(num_frames):
            current_code = "\n".join(code_lines[: i + 1])
            highlighted = pygments.highlight(current_code, lexer, formatter)

            temp_path = os.path.join(tmpdir, f"frame_{i}.png")
            with open(temp_path, "wb") as f:
                f.write(highlighted)

            code_img = Image.open(temp_path)
            background = create_gradient_background(
                background_width, background_height, gradient_start, gradient_end, i
            )

            current_y = min_padding

            # Add title
            if title:
                draw = ImageDraw.Draw(background)
                text_width = draw.textlength(title, font=title_font)
                text_x = (background_width - text_width) // 2
                draw.text(
                    (text_x, current_y), title, fill=(255, 255, 255), font=title_font
                )
                current_y += title_height

            # Add photo above
            if photo_img and photo_position == "above":
                photo_x = (background_width - photo_img.width) // 2
                background.paste(photo_img, (photo_x, current_y), photo_img)
                current_y += photo_height

            # Add comments above
            if comments and comments_position == "above":
                draw = ImageDraw.Draw(background)
                comment_idx = min(i // frames_per_comment, len(comments) - 1)
                comment = comments[comment_idx]
                text_width = draw.textlength(comment, font=comment_font)
                text_x = (background_width - text_width) // 2
                draw.text(
                    (text_x, current_y),
                    comment,
                    fill=(255, 255, 255),
                    font=comment_font,
                )
                current_y += comment_height

            # Add code window
            window_copy = window.copy()
            window_copy.paste(
                code_img, (window_padding, window_padding + title_bar_height)
            )
            background.paste(window_copy, (window_x, current_y), window_copy)
            current_y += final_height

            # Add photo below
            if photo_img and photo_position == "below":
                photo_x = (background_width - photo_img.width) // 2
                current_y += min_padding
                background.paste(photo_img, (photo_x, current_y), photo_img)
                current_y += photo_height

            # Add comments below
            if comments and comments_position == "below":
                draw = ImageDraw.Draw(background)
                comment_idx = min(i // frames_per_comment, len(comments) - 1)
                comment = comments[comment_idx]
                text_width = draw.textlength(comment, font=comment_font)
                text_x = (background_width - text_width) // 2
                current_y += min_padding
                draw.text(
                    (text_x, current_y),
                    comment,
                    fill=(255, 255, 255),
                    font=comment_font,
                )

            # Add logo
            if logo:
                logo_x = background_width - logo.width - min_padding
                logo_y = background_height - logo.height - min_padding
                background.paste(
                    logo, (logo_x, logo_y), logo if logo.mode == "RGBA" else None
                )

            frames.append(background)

    # Calculate frame delays
    delays = np.array(
        [
            start_delay * (acceleration ** (i / (len(frames) - 1)))
            for i in range(len(frames))
        ]
    )
    delays = np.clip(delays, end_delay, None)

    if output_file is not None:
        # Save the animated GIF
        frames[0].save(
            output_file,
            save_all=True,
            append_images=frames[1:],
            duration=[int(d * 1000) for d in delays],
            loop=0,
        )
    return frames