publisher/tasks/render.py

import os
import uuid
import asyncio
import subprocess
from datetime import datetime


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

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


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

def normalize_context(context):

    if isinstance(context, dict):
        return {
            "video_path": context.get("video_path"),
            "srt": context.get("srt"),
            "subtitles": context.get("subtitles")
        }

    return {
        "video_path": getattr(context, "video_path", None),
        "srt": getattr(context, "srt", None),
        "subtitles": getattr(context, "subtitles", None)
    }


# -------------------------------------------------
# SRT RESOLVER
# -------------------------------------------------

def resolve_srt(ctx):
    """
    Accepts:
    - raw SRT string
    - file path
    - None
    """

    srt = ctx.get("srt")

    if not srt:
        return None

    if isinstance(srt, str) and os.path.exists(srt):
        with open(srt, "r", encoding="utf-8") as f:
            return f.read()

    return srt if isinstance(srt, str) else None


# -------------------------------------------------
# SAFE FFMPEG RENDER ENGINE
# -------------------------------------------------

def run_ffmpeg(video_path, srt_path, output_path):

    cmd = [
        "ffmpeg",
        "-y",
        "-i", video_path,
    ]

    # Subtitle overlay (only if available)
    if srt_path and os.path.exists(srt_path):
        cmd += [
            "-vf",
            f"subtitles={srt_path}"
        ]

    cmd += [
        "-c:v", "libx264",
        "-preset", "veryfast",
        "-crf", "23",
        "-c:a", "aac",
        "-b:a", "128k",
        output_path
    ]

    process = subprocess.run(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True
    )

    if process.returncode != 0:
        raise RuntimeError(process.stderr)

    return output_path


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

async def run(context):

    ctx = normalize_context(context)

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

    try:

        video_path = ctx.get("video_path")
        srt_data = resolve_srt(ctx)

        if not video_path or not os.path.exists(video_path):
            return {
                "status": "error",
                "task": "render",
                "message": "Missing or invalid video_path",
                "stage": "validation"
            }

        # -------------------------------------------------
        # TEMP SRT FILE HANDLING
        # -------------------------------------------------

        srt_path = None

        if srt_data:
            srt_path = os.path.join(OUTPUT_DIR, f"{batch_id}.srt")

            with open(srt_path, "w", encoding="utf-8") as f:
                f.write(srt_data)

        output_path = os.path.join(
            OUTPUT_DIR,
            f"{batch_id}_render.mp4"
        )

        # -------------------------------------------------
        # FFMPEG EXECUTION (THREAD SAFE)
        # -------------------------------------------------

        await asyncio.to_thread(
            run_ffmpeg,
            video_path,
            srt_path,
            output_path
        )

        # -------------------------------------------------
        # CLEANUP OPTIONAL
        # -------------------------------------------------

        if srt_path and os.path.exists(srt_path):
            os.remove(srt_path)

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

        return {
            "status": "success",
            "task": "render",
            "batch_id": batch_id,
            "started_at": started_at,
            "completed_at": datetime.utcnow().isoformat(),
            "output_path": output_path
        }

    except Exception as e:

        return {
            "status": "error",
            "task": "render",
            "batch_id": batch_id,
            "message": str(e),
            "stage": "render_failed"
        }