video_assembler.py

"""
TimeLapseForge — Video Assembler Module
Converts frame sequences into complete videos with transitions,
text overlays, and background audio.
"""

import os
import tempfile
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from typing import List, Dict, Optional, Tuple


class VideoAssembler:
    """Assembles frames into a polished timelapse video."""

    def __init__(self):
        self.font_cache = {}

    def _load_font(self, size: int = 32) -> ImageFont.FreeTypeFont:
        """Load a font with caching and fallbacks."""
        if size in self.font_cache:
            return self.font_cache[size]

        font_paths = [
            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
            "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
            "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
            "/usr/share/fonts/TTF/DejaVuSans-Bold.ttf",
        ]
        for path in font_paths:
            if os.path.exists(path):
                try:
                    font = ImageFont.truetype(path, size)
                    self.font_cache[size] = font
                    return font
                except Exception:
                    continue

        font = ImageFont.load_default()
        self.font_cache[size] = font
        return font

    def add_text_overlay(
        self,
        image: Image.Image,
        text: str,
        position: str = "bottom",
        font_size: int = 32,
        text_color: Tuple[int, ...] = (255, 255, 255),
        bg_color: Tuple[int, ...] = (0, 0, 0, 160),
        margin: int = 10,
    ) -> Image.Image:
        """Add text overlay with semi-transparent background."""
        img = image.copy().convert("RGBA")
        overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
        draw = ImageDraw.Draw(overlay)
        font = self._load_font(font_size)

        try:
            bbox = draw.textbbox((0, 0), text, font=font)
            text_w = bbox[2] - bbox[0]
            text_h = bbox[3] - bbox[1]
        except Exception:
            text_w = len(text) * font_size // 2
            text_h = font_size

        img_w, img_h = img.size
        if position == "top":
            x = (img_w - text_w) // 2
            y = margin
        elif position == "bottom":
            x = (img_w - text_w) // 2
            y = img_h - text_h - margin * 3
        elif position == "top-left":
            x = margin
            y = margin
        elif position == "top-right":
            x = img_w - text_w - margin
            y = margin
        else:
            x = (img_w - text_w) // 2
            y = (img_h - text_h) // 2

        padding = 8
        draw.rectangle(
            [x - padding, y - padding, x + text_w + padding, y + text_h + padding],
            fill=bg_color,
        )
        draw.text((x, y), text, fill=text_color, font=font)

        result = Image.alpha_composite(img, overlay)
        return result.convert("RGB")

    def add_progress_bar(
        self,
        image: Image.Image,
        progress: float,
        bar_height: int = 6,
        bar_color: Tuple[int, ...] = (0, 200, 100),
        bg_color: Tuple[int, ...] = (60, 60, 60),
    ) -> Image.Image:
        """Add a progress bar at the bottom of the frame."""
        img = image.copy()
        draw = ImageDraw.Draw(img)
        w, h = img.size

        draw.rectangle([0, h - bar_height, w, h], fill=bg_color)
        bar_width = int(w * min(progress, 1.0))
        if bar_width > 0:
            draw.rectangle([0, h - bar_height, bar_width, h], fill=bar_color)

        return img

    def prepare_frames(
        self,
        images: List[Image.Image],
        fps: int = 24,
        hold_seconds: float = 2.0,
        add_labels: bool = True,
        labels: Optional[List[str]] = None,
        add_progress: bool = True,
        add_bookend_labels: bool = True,
    ) -> List[np.ndarray]:
        """Prepare frames for video export."""
        frames_per_panel = max(1, int(fps * hold_seconds))
        all_frames = []

        for i, img in enumerate(images):
            frame = img.copy()

            if add_labels and labels and i < len(labels) and labels[i]:
                frame = self.add_text_overlay(
                    frame, labels[i], position="top-left", font_size=24
                )

            if add_bookend_labels:
                if i == 0:
                    frame = self.add_text_overlay(
                        frame, "BEFORE", position="bottom", font_size=40
                    )
                elif i == len(images) - 1:
                    frame = self.add_text_overlay(
                        frame, "AFTER", position="bottom", font_size=40
                    )

            if add_progress:
                progress = (i + 1) / len(images)
                frame = self.add_progress_bar(frame, progress)

            frame_array = np.array(frame.convert("RGB"))
            for _ in range(frames_per_panel):
                all_frames.append(frame_array)

        return all_frames

    def create_video(
        self,
        images: List[Image.Image],
        output_path: Optional[str] = None,
        fps: int = 24,
        hold_seconds: float = 2.0,
        add_labels: bool = True,
        labels: Optional[List[str]] = None,
        add_progress: bool = True,
        add_bookend_labels: bool = True,
        fade_in_frames: int = 15,
        fade_out_frames: int = 15,
    ) -> str:
        """Create a complete video from panel images."""
        import imageio

        if output_path is None:
            output_path = os.path.join(
                tempfile.gettempdir(), "timelapse_output.mp4"
            )

        frames = self.prepare_frames(
            images, fps, hold_seconds, add_labels, labels,
            add_progress, add_bookend_labels,
        )

        if not frames:
            raise ValueError("No frames to assemble!")

        # Apply fade in
        for i in range(min(fade_in_frames, len(frames))):
            alpha = i / max(fade_in_frames, 1)
            frames[i] = (frames[i].astype(np.float32) * alpha).astype(np.uint8)

        # Apply fade out
        for i in range(min(fade_out_frames, len(frames))):
            idx = len(frames) - 1 - i
            alpha = i / max(fade_out_frames, 1)
            frames[idx] = (frames[idx].astype(np.float32) * alpha).astype(np.uint8)

        # Write video
        writer = imageio.get_writer(
            output_path,
            fps=fps,
            codec="libx264",
            quality=8,
            pixelformat="yuv420p",
            output_params=["-preset", "medium"],
        )

        for frame in frames:
            writer.append_data(frame)

        writer.close()
        return output_path

    def add_audio_to_video(
        self,
        video_path: str,
        audio_path: str,
        output_path: Optional[str] = None,
    ) -> str:
        """Add background audio to the video using moviepy."""
        if output_path is None:
            output_path = video_path.replace(".mp4", "_audio.mp4")

        try:
            from moviepy.editor import VideoFileClip, AudioFileClip

            video = VideoFileClip(video_path)
            audio = AudioFileClip(audio_path)

            if audio.duration < video.duration:
                from moviepy.editor import concatenate_audioclips
                loops = int(video.duration / audio.duration) + 1
                audio = concatenate_audioclips([audio] * loops)

            audio = audio.subclip(0, video.duration)
            final = video.set_audio(audio)
            final.write_videofile(
                output_path, codec="libx264",
                audio_codec="aac", logger=None,
            )
            video.close()
            return output_path

        except Exception as e:
            print(f"Audio mixing failed: {e}")
            return video_path

    def create_comparison_image(
        self,
        first: Image.Image,
        last: Image.Image,
        gap: int = 4,
    ) -> Image.Image:
        """Create side-by-side comparison image."""
        w, h = first.size
        comp = Image.new("RGB", (w * 2 + gap, h), (30, 30, 30))
        comp.paste(first.resize((w, h), Image.LANCZOS), (0, 0))
        comp.paste(last.resize((w, h), Image.LANCZOS), (w + gap, 0))

        comp = self.add_text_overlay(comp, "BEFORE", position="top-left", font_size=28)
        # Add AFTER on right side
        draw = ImageDraw.Draw(comp)
        font = self._load_font(28)
        try:
            draw.text((w + gap + 10, 10), "AFTER", fill=(255, 255, 255), font=font)
        except Exception:
            pass

        return comp

    def create_gif(
        self,
        images: List[Image.Image],
        output_path: Optional[str] = None,
        duration_per_frame: int = 500,
        loop: int = 0,
    ) -> str:
        """Create an animated GIF."""
        if output_path is None:
            output_path = os.path.join(tempfile.gettempdir(), "timelapse.gif")

        max_size = (480, 480)
        resized = []
        for img in images:
            r = img.copy()
            r.thumbnail(max_size, Image.LANCZOS)
            resized.append(r.convert("RGB"))

        resized[0].save(
            output_path,
            save_all=True,
            append_images=resized[1:],
            duration=duration_per_frame,
            loop=loop,
            optimize=True,
        )
        return output_path