Build: Async Batch Generator

import asyncio
import json
import anthropic
from dotenv import load_dotenv
from logger import get_logger

load_dotenv()
logger = get_logger(__name__)

client = anthropic.AsyncAnthropic()

PLATFORMS = ["Kling", "Runway", "Veo"]


async def generate_one(shot, platform):
    """Generate a single prompt. Returns a dict with the result or error info."""
    try:
        message = await client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=150,
            messages=[{
                "role": "user",
                "content": f"Write a {platform} AI video prompt for: {shot}. "
                           f"Under 75 words. Only the prompt."
            }]
        )
        logger.info(f"Generated: [{platform}] {shot[:40]}...")
        return {
            "platform": platform,
            "shot": shot,
            "prompt": message.content[0].text,
            "status": "success",
        }
    except Exception as e:
        logger.error(f"Failed: [{platform}] {shot[:40]}... — {e}")
        return {
            "platform": platform,
            "shot": shot,
            "prompt": None,
            "status": f"error: {e}",
        }


async def batch_generate(shots, platforms=None):
    """Generate prompts for all shots × all platforms in parallel."""
    platforms = platforms or PLATFORMS

    # Build a flat list of all tasks — every combination of shot × platform.
    tasks = []
    for shot in shots:
        for platform in platforms:
            tasks.append(generate_one(shot, platform))

    print(f"Generating {len(tasks)} prompts ({len(shots)} shots × {len(platforms)} platforms)...")

    # asyncio.gather runs them ALL in parallel and returns results in order.
    results = await asyncio.gather(*tasks)
    return results


async def main():
    """Interactive mode — enter shots, generate in batch."""
    print("=== DVP Batch Prompt Generator ===\n")
    print("Enter shots one per line. Type 'done' when finished.\n")

    shots = []
    while True:
        shot = input(f"  Shot {len(shots) + 1} (or 'done'): ")
        if shot.strip().lower() == "done":
            break
        if shot.strip():
            shots.append(shot.strip())

    if not shots:
        print("No shots entered.")
        return

    # Run the async batch generation
    results = await batch_generate(shots)

    # Display results grouped by shot
    print(f"\n{'=' * 60}")
    current_shot = None
    for r in results:
        if r["shot"] != current_shot:
            current_shot = r["shot"]
            print(f"\n  SHOT: {current_shot}")
            print(f"  {'-' * 50}")

        if r["status"] == "success":
            print(f"\n  [{r['platform']}]")
            print(f"  {r['prompt']}\n")
        else:
            print(f"\n  [{r['platform']}] FAILED: {r['status']}\n")

    # Offer to save results
    successful = [r for r in results if r["status"] == "success"]
    print(f"\n{len(successful)} of {len(results)} prompts generated successfully.")

    if successful and input("Save all to database? (y/n): ").lower() == "y":
        from database import Database
        db = Database("prompt_vault.db")
        for r in successful:
            db.add_prompt(r["platform"], r["shot"], r["prompt"])
        db.close()
        print(f"Saved {len(successful)} prompts to database.")


# In a .py file, run async code with asyncio.run().
# This starts the event loop, runs main(), and cleans up.
if __name__ == "__main__":
    asyncio.run(main())