src/generation/plan_generator.py

# src/generation/plan_generator.py
import json
from src.core.ai_helpers import generate_with_ai
from src.utils.smart_logger import logger

def plan_video(project_config: dict, gemini_client):
    """Workflow Step 1: Create a video plan."""
    logger.start_operation("plan_video", "Generating video plan")
    
    user_prompt = project_config.get("prompt")
    if not user_prompt:
        logger.end_operation("plan_video", "No prompt found")
        return

    logger.quick_log("info", f"Prompt: {user_prompt[:50]}...")
    
    if not gemini_client:
        logger.end_operation("plan_video", "Gemini client not initialized")
        return
    
    # Test AI connection first
    test_response = generate_with_ai(gemini_client, "Hello", is_json=False)
    if not test_response:
        logger.end_operation("plan_video", "AI connection test failed")
        return
    
    logger.quick_log("success", "AI connection verified")
    
    director_prompt = f"""
    You are a Director AI for creating educational Manim videos. Based on the user's prompt, create a comprehensive video plan.

    User Prompt: "{user_prompt}"

    IMPORTANT: You must return a JSON OBJECT (starting with {{), NOT a JSON ARRAY (starting with [).

    Generate a JSON object with EXACTLY this structure:
    {{
        "video_title": "Video Title",
        "video_description": "Brief description.",
        "scenes": [
            {{
                "scene_name": "IntroScene",
                "description": "Detailed description of visuals and animations.",
                "narration": "The complete narration text for this scene."
            }}
        ]
    }}

    CRITICAL REQUIREMENTS:
    1. Start your response with {{ (curly brace), NOT [ (square bracket)
    2. Create 2-4 concise scenes.
    3. Scene names must be valid Python class names (no spaces, special chars).
    4. Use only basic geometric shapes: circles, rectangles, squares, triangles, lines, arrows
    5. Colors only: RED, BLUE, GREEN, YELLOW, WHITE, BLACK, GRAY, ORANGE, PURPLE, PINK
    6. Always white background with dark text
    7. Return ONLY the JSON object, no other text.
    """
    
    try:
        response_text = generate_with_ai(gemini_client, director_prompt, is_json=True)
        if not response_text:
            logger.end_operation("plan_video", "No AI response received")
            return

        # Parse and validate the JSON response
        video_plan = json.loads(response_text)
        
        # Handle case where AI returns a list instead of dict
        if isinstance(video_plan, list):
            logger.quick_log("warn", "AI returned list, converting to proper format")
            # If it's a list of scenes, wrap it in the proper structure
            if len(video_plan) > 0 and isinstance(video_plan[0], dict):
                video_plan = {
                    "video_title": "Generated Video",
                    "video_description": "Educational video created from user prompt",
                    "scenes": video_plan
                }
            else:
                logger.error_with_context("Invalid list format", "List doesn't contain valid scenes")
                logger.end_operation("plan_video", "Invalid list structure")
                return
        
        # Validate the structure
        if not isinstance(video_plan, dict):
            logger.error_with_context("Invalid plan format", "Expected dict, got " + type(video_plan).__name__)
            logger.end_operation("plan_video", "Invalid plan structure")
            return
        
        scenes = video_plan.get('scenes', [])
        if not isinstance(scenes, list):
            logger.error_with_context("Invalid scenes format", "Expected list, got " + type(scenes).__name__)
            logger.end_operation("plan_video", "Invalid scenes structure")
            return
        
        # Validate each scene has required fields
        for i, scene in enumerate(scenes):
            if not isinstance(scene, dict):
                logger.error_with_context(f"Scene {i} invalid", "Expected dict, got " + type(scene).__name__)
                logger.end_operation("plan_video", "Invalid scene structure")
                return
            
            # Ensure required fields exist
            if not scene.get('scene_name'):
                scene['scene_name'] = f"Scene{i+1}"
            if not scene.get('description'):
                scene['description'] = "Scene description"
            if not scene.get('narration'):
                scene['narration'] = ""
        
        scenes_count = len(scenes)
        project_config["video_plan"] = video_plan
        logger.end_operation("plan_video", f"Plan created with {scenes_count} scenes")
        
    except json.JSONDecodeError as e:
        logger.error_with_context("Failed to parse AI response as JSON", str(e)[:30])
        logger.end_operation("plan_video", "JSON parsing failed")
    except Exception as e:
        logger.error_with_context("Plan generation failed", str(e)[:50])
        logger.end_operation("plan_video", "Plan generation failed")