publisher/tasks/batch_runner.py

import uuid
import asyncio
from datetime import datetime

from utils.logger import logger
from utils.job_queue import create_job, get_job


# =====================================================
# VALIDATION
# =====================================================

def normalize_payload(payload: dict | None):
    payload = payload or {}

    items = payload.get("items", [])

    if not isinstance(items, list):
        raise ValueError("batch items must be a list")

    normalized = []

    for i, item in enumerate(items):
        if not isinstance(item, dict):
            raise ValueError(f"batch item {i} must be an object")

        normalized.append({
            "id": item.get("id", str(uuid.uuid4())),
            "task": item.get("task"),
            "video_path": item.get("video_path"),
            "source": item.get("source"),
            "payload": item.get("payload", {}),
            "webhook": item.get("webhook"),
        })

    if not normalized:
        raise ValueError("batch cannot be empty")

    return normalized


# =====================================================
# SAFE TASK EXECUTION WRAPPER
# =====================================================

async def execute_single(task_executor, item, index: int):

    try:
        logger.info(f"[BATCH] Executing item {index} → {item['task']}")

        result = await task_executor(
            item["task"],
            {
                **(item.get("payload") or {}),
                "video_path": item.get("video_path"),
                "source": item.get("source"),
            },
            item.get("webhook"),
        )

        return {
            "index": index,
            "id": item["id"],
            "task": item["task"],
            "status": "success",
            "result": result,
        }

    except Exception as e:

        logger.exception(f"[BATCH ERROR] item {index}")

        return {
            "index": index,
            "id": item["id"],
            "task": item.get("task"),
            "status": "failed",
            "error": str(e),
        }


# =====================================================
# CONCURRENCY CONTROLLER
# =====================================================

async def run_concurrent(tasks, executor, max_concurrency: int = 3):

    semaphore = asyncio.Semaphore(max_concurrency)

    async def bound(item, index):
        async with semaphore:
            return await execute_single(executor, item, index)

    return await asyncio.gather(
        *[bound(item, i) for i, item in enumerate(tasks)]
    )


# =====================================================
# MAIN ENTRYPOINT (REGISTRY COMPATIBLE)
# =====================================================

async def run(payload: dict | None = None, context: dict | None = None):

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

    try:

        items = normalize_payload(payload)

        logger.info(f"[BATCH] Starting batch_id={batch_id}, items={len(items)}")

        # -------------------------------------------------
        # EXECUTOR HOOK (inject from registry context)
        # -------------------------------------------------

        def executor(task_name, task_payload, webhook=None):
            """
            This is intentionally abstract so it can plug into:
            - registry executor
            - legacy executor
            - FastAPI layer
            """

            from core.execution.executor import execute_task

            return execute_task(task_name, task_payload, webhook)

        # -------------------------------------------------
        # RUN BATCH
        # -------------------------------------------------

        results = await run_concurrent(items, executor)

        success_count = len([r for r in results if r["status"] == "success"])
        failed_count = len(results) - success_count

        # -------------------------------------------------
        # OUTPUT CONTRACT
        # -------------------------------------------------

        return {
            "status": "completed",
            "task": "batch_runner",

            "batch_id": batch_id,
            "started_at": started_at,
            "completed_at": datetime.utcnow().isoformat(),

            "total": len(items),
            "success": success_count,
            "failed": failed_count,

            "results": results,
        }

    except Exception as e:

        logger.exception("[BATCH FATAL ERROR]")

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