publisher/tasks/batch.py

import asyncio
import uuid
from datetime import datetime

# Optional queue integration (safe fallback if not present)
try:
    from utils.job_queue import create_job
except Exception:
    create_job = None


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

def normalize_context(context):
    if isinstance(context, dict):
        return {
            "items": context.get("items", []),
            "webhook": context.get("webhook"),
            "mode": context.get("mode", "sequential")
        }
    return {
        "items": getattr(context, "items", []),
        "webhook": getattr(context, "webhook", None),
        "mode": getattr(context, "mode", "sequential")
    }


# -------------------------------------------------
# SINGLE TASK EXECUTOR WRAPPER
# -------------------------------------------------

async def execute_single(task_name, payload):
    """
    Uses registry executor if available, otherwise returns structured fallback.
    """

    try:
        from core.execution.executor import execute_task

        return await execute_task(
            task_name,
            payload
        )

    except Exception as e:
        return {
            "task": task_name,
            "status": "failed",
            "error": str(e)
        }


# -------------------------------------------------
# MAIN BATCH RUNNER
# -------------------------------------------------

async def run(context):

    ctx = normalize_context(context)
    items = ctx["items"]

    if not items:
        return {
            "status": "error",
            "message": "Batch requires 'items' list"
        }

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

    results = []
    failed = 0

    # -------------------------------------------------
    # MODE: SEQUENTIAL EXECUTION
    # -------------------------------------------------

    if ctx["mode"] == "sequential":

        for i, item in enumerate(items):

            task_name = item.get("task")
            payload = item.get("payload", {})

            if not task_name:
                results.append({
                    "index": i,
                    "status": "skipped",
                    "reason": "missing task"
                })
                continue

            result = await execute_single(task_name, payload)

            if isinstance(result, dict) and result.get("status") == "failed":
                failed += 1

            results.append({
                "index": i,
                "task": task_name,
                "result": result
            })

    # -------------------------------------------------
    # MODE: PARALLEL EXECUTION
    # -------------------------------------------------

    elif ctx["mode"] == "parallel":

        async def run_item(i, item):
            task_name = item.get("task")
            payload = item.get("payload", {})

            if not task_name:
                return {
                    "index": i,
                    "status": "skipped"
                }

            result = await execute_single(task_name, payload)

            return {
                "index": i,
                "task": task_name,
                "result": result
            }

        results = await asyncio.gather(
            *[run_item(i, item) for i, item in enumerate(items)]
        )

    else:
        return {
            "status": "error",
            "message": f"Unsupported mode: {ctx['mode']}"
        }

    # -------------------------------------------------
    # JOB QUEUE INTEGRATION (OPTIONAL)
    # -------------------------------------------------

    job_id = None
    if create_job:
        try:
            job_id = create_job(
                video_path=None,
                webhook=ctx["webhook"],
                metadata={
                    "batch_id": batch_id,
                    "total": len(items),
                    "failed": failed
                }
            )
        except Exception:
            job_id = None

    # -------------------------------------------------
    # FINAL RESPONSE
    # -------------------------------------------------

    return {
        "status": "success",
        "batch_id": batch_id,
        "job_id": job_id,
        "started_at": started_at,
        "completed_at": datetime.utcnow().isoformat(),
        "mode": ctx["mode"],
        "total": len(items),
        "failed": failed,
        "results": results
    }