publisher/tasks/generate_metadata.py

import asyncio
from datetime import datetime
import uuid


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

def normalize_context(context):
    """
    Accepts registry dict or legacy object input.
    """

    if isinstance(context, dict):
        return {
            "strategy": context.get("strategy", {}),
            "transcript": context.get("transcript", {}),
            "platform": context.get("platform", "tiktok"),
            "video_path": context.get("video_path")
        }

    return {
        "strategy": getattr(context, "strategy", {}),
        "transcript": getattr(context, "transcript", {}),
        "platform": getattr(context, "platform", "tiktok"),
        "video_path": getattr(context, "video_path", None)
    }


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

def extract_text(transcript):
    """
    Extracts usable text from transcript structure safely.
    """

    if isinstance(transcript, dict):
        segments = transcript.get("segments", [])
        if segments:
            return " ".join([s.get("text", "") for s in segments])

    if isinstance(transcript, str):
        return transcript

    return ""


# -------------------------------------------------
# METADATA GENERATOR CORE
# -------------------------------------------------

def build_metadata(text, strategy, platform):
    """
    Deterministic metadata generator (no external API required).
    """

    hook = ""
    if isinstance(strategy, dict):
        hook = strategy.get("hook", "")

    if not hook:
        hook = text[:120] + "..." if text else "Discover powerful insights in this video."

    title = hook[:70].strip()

    description = (
        f"{hook}\n\n"
        f"Watch till the end for key insights.\n"
        f"Optimized for {platform}."
    )

    tags = [
        "content",
        "viral",
        "shorts",
        platform,
        "ai generated",
        "social media"
    ]

    hashtags = [
        "#ViralContent",
        "#ContentCreator",
        "#Shorts",
        f"#{platform.capitalize()}",
        "#AIContent"
    ]

    return {
        "title": title,
        "description": description,
        "tags": tags,
        "hashtags": hashtags
    }


# -------------------------------------------------
# OPTIONAL LLM ENHANCEMENT (SAFE WRAPPER)
# -------------------------------------------------

async def enhance_with_llm(base_metadata, context):
    """
    Optional enhancement layer.
    Never breaks pipeline if API missing.
    """

    try:
        import os

        if not os.getenv("GEMINI_API_KEY"):
            return base_metadata

        # Lazy import to avoid startup crashes
        import google.generativeai as genai

        genai.configure(api_key=os.environ["GEMINI_API_KEY"])

        model = genai.GenerativeModel("gemini-1.5-flash")

        prompt = f"""
Improve this social media metadata for virality:

TITLE: {base_metadata['title']}
DESCRIPTION: {base_metadata['description']}
TAGS: {base_metadata['tags']}
HASHTAGS: {base_metadata['hashtags']}

Return STRICT JSON with:
title, description, tags, hashtags
"""

        response = await model.generate_content_async(prompt)

        import json
        cleaned = response.text.strip().replace("```json", "").replace("```", "")
        data = json.loads(cleaned)

        return data

    except Exception:
        return base_metadata


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

async def run(context):

    ctx = normalize_context(context)

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

    try:

        strategy = ctx["strategy"]
        transcript = ctx["transcript"]
        platform = ctx["platform"]

        text = extract_text(transcript)

        # -------------------------------------------------
        # BASE METADATA
        # -------------------------------------------------

        base_metadata = build_metadata(text, strategy, platform)

        # -------------------------------------------------
        # OPTIONAL ENHANCEMENT
        # -------------------------------------------------

        final_metadata = await enhance_with_llm(base_metadata, ctx)

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

        return {
            "status": "success",
            "task": "generate-metadata",
            "batch_id": batch_id,
            "started_at": started_at,
            "completed_at": datetime.utcnow().isoformat(),
            "platform": platform,
            "metadata": final_metadata
        }

    except Exception as e:

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