publisher/tasks/generate_thumbnail.py

import os
import uuid
import asyncio
from datetime import datetime

from PIL import Image, ImageDraw, ImageFont


# -------------------------------------------------
# SAFE OUTPUT DIRECTORY
# -------------------------------------------------

OUTPUT_DIR = "jobs/thumbnails"
os.makedirs(OUTPUT_DIR, exist_ok=True)


# -------------------------------------------------
# CONTEXT NORMALIZER
# -------------------------------------------------

def normalize_context(context):

    if isinstance(context, dict):
        return {
            "title": context.get("title"),
            "hook": context.get("hook"),
            "strategy": context.get("strategy", {}),
            "text": context.get("text", ""),
            "video_path": context.get("video_path")
        }

    return {
        "title": getattr(context, "title", None),
        "hook": getattr(context, "hook", None),
        "strategy": getattr(context, "strategy", {}),
        "text": getattr(context, "text", ""),
        "video_path": getattr(context, "video_path", None)
    }


# -------------------------------------------------
# TEXT EXTRACTOR
# -------------------------------------------------

def extract_text(ctx):
    if ctx.get("hook"):
        return ctx["hook"]

    if isinstance(ctx.get("strategy"), dict):
        hook = ctx["strategy"].get("hook")
        if hook:
            return hook

    return ctx.get("text") or "Create engaging content that stands out."


# -------------------------------------------------
# SAFE FONT LOADER
# -------------------------------------------------

def load_font(size):
    """
    Tries system fonts safely.
    Falls back to default PIL font if unavailable.
    """

    try:
        return ImageFont.truetype("arial.ttf", size)
    except Exception:
        return ImageFont.load_default()


# -------------------------------------------------
# THUMBNAIL GENERATOR CORE
# -------------------------------------------------

def build_thumbnail(text, width=1280, height=720):

    img = Image.new("RGB", (width, height), color=(10, 10, 10))
    draw = ImageDraw.Draw(img)

    # Accent style bar
    draw.rectangle([0, 0, 20, height], fill=(232, 255, 71))

    # Title text
    font_large = load_font(64)
    font_small = load_font(36)

    wrapped_text = text[:120]

    draw.text(
        (60, 200),
        wrapped_text,
        font=font_large,
        fill=(232, 232, 232)
    )

    draw.text(
        (60, 320),
        "AI-Generated Content",
        font=font_small,
        fill=(136, 136, 136)
    )

    return img


# -------------------------------------------------
# MAIN ENTRYPOINT
# -------------------------------------------------

async def run(context):

    ctx = normalize_context(context)

    batch_id = str(uuid.uuid4())
    started_at = datetime.utcnow().isoformat()

    try:

        text = extract_text(ctx)

        # -------------------------------------------------
        # BUILD IMAGE (CPU SAFE)
        # -------------------------------------------------

        img = await asyncio.to_thread(build_thumbnail, text)

        file_name = f"{batch_id}_thumbnail.png"
        output_path = os.path.join(OUTPUT_DIR, file_name)

        img.save(output_path)

        # -------------------------------------------------
        # RESPONSE
        # -------------------------------------------------

        return {
            "status": "success",
            "task": "generate-thumbnail",
            "batch_id": batch_id,
            "started_at": started_at,
            "completed_at": datetime.utcnow().isoformat(),
            "thumbnail_path": output_path
        }

    except Exception as e:

        return {
            "status": "error",
            "task": "generate-thumbnail",
            "batch_id": batch_id,
            "message": str(e),
            "stage": "thumbnail_generation_failed"
        }