src/generation/code_generator.py

# src/generation/code_generator.py
import json
import tempfile
from src.core.ai_helpers import generate_with_ai
from src.utils.helpers import fix_invalid_colors
from src.utils.smart_logger import logger

def generate_scene_code(project_config: dict, gemini_client, temp_dir_name: str):
    """Workflow Step 2: Generate Manim code and save to a temporary file."""
    logger.start_operation("code_gen", "Generating scene code")
    
    if not project_config.get("video_plan"):
        logger.end_operation("code_gen", "No video plan found")
        return

    plan = project_config["video_plan"]
    scenes_count = len(plan.get('scenes', []))
    logger.quick_log("info", f"Generating code for {scenes_count} scenes")
    
    plan_str = json.dumps(project_config["video_plan"], indent=2)

    code_gen_prompt = f"""
    Generate a single Python file containing all Manim scene classes based on this video plan:
    {plan_str}

    Requirements:
    1. Create ONE Python file with all necessary Manim scene classes.
    2. Use modern Manim syntax (e.g., `from manim import *`).
    3. Each class name must exactly match a `scene_name` from the plan.
    4. Implement the `construct` method for each scene based on its `description`.
    5. Set the background color for every scene to WHITE.
    6. Add comments to explain the code.
    7. At the end, add: `SCENE_ORDER = ["IntroScene", "MainConceptScene", ...]`
    8. Return ONLY the Python code, inside a single markdown block.
    
    CRITICAL CONSTRAINTS:
    - NEVER use ImageMobject, SVGMobject, or any file references.
    - NEVER use MathTex, Tex, or any LaTeX-dependent objects (LaTeX not installed).
    - Use Text() for all text including mathematical formulas.
    - Create all visuals using Manim's built-in shapes and functions ONLY.
    - ONLY use these colors: RED, BLUE, GREEN, YELLOW, WHITE, BLACK, GRAY, ORANGE, PURPLE, PINK.
    - To position objects, call .to_edge(), .move_to(), or .shift() DIRECTLY on the Mobject.
    - For math formulas, use Text() with simple strings like "x² + y² = z²".
    """
    
    response_text = generate_with_ai(gemini_client, code_gen_prompt, is_json=False)
    if not response_text:
        logger.end_operation("code_gen", "No AI response received")
        return

    if "```python" in response_text:
        all_scenes_code = response_text.split("```python")[1].split("```")[0].strip()
    else:
        all_scenes_code = response_text.strip()
        
    all_scenes_code = fix_invalid_colors(all_scenes_code)

    with tempfile.NamedTemporaryFile(mode='w+', suffix='.py', delete=False, dir=temp_dir_name) as temp_code_file:
        temp_code_file.write(all_scenes_code)
        temp_code_file_path = temp_code_file.name

    for scene_name in project_config["scenes"]:
        project_config["scenes"][scene_name]["code_file"] = temp_code_file_path

    code_size = len(all_scenes_code)
    logger.end_operation("code_gen", "Scene code generated", f"{code_size} chars → {temp_code_file_path}")

def regenerate_single_scene_code(project_config: dict, gemini_client, temp_dir_name: str, scene_name: str) -> bool:
    """Regenerate code for a single scene when debugging fails."""
    logger.start_operation("regen_scene", f"Regenerating code for {scene_name}")
    
    if not project_config.get("video_plan"):
        logger.end_operation("regen_scene", "No video plan found")
        return False

    # Find the specific scene in the plan
    scenes = project_config["video_plan"].get("scenes", [])
    target_scene = None
    for scene in scenes:
        if scene.get("scene_name") == scene_name:
            target_scene = scene
            break
    
    if not target_scene:
        logger.end_operation("regen_scene", f"Scene {scene_name} not found in plan")
        return False

    scene_prompt = f"""
    Generate a single Manim scene class for this specific scene:
    
    Scene Name: {scene_name}
    Description: {target_scene.get('description', '')}
    Narration: {target_scene.get('narration', '')}

    Requirements:
    1. Create ONE Python file with the complete scene class.
    2. Use modern Manim syntax (e.g., `from manim import *`).
    3. Class name must be exactly: {scene_name}
    4. Implement the `construct` method based on the description.
    5. Set background color to WHITE.
    6. Add detailed comments.
    7. Return ONLY the Python code, inside a single markdown block.
    
    CRITICAL CONSTRAINTS:
    - NEVER use ImageMobject, SVGMobject, or any file references.
    - NEVER use MathTex, Tex, or any LaTeX-dependent objects (LaTeX not installed).
    - Use Text() for all text including mathematical formulas.
    - Create all visuals using Manim's built-in shapes and functions ONLY.
    - ONLY use these colors: RED, BLUE, GREEN, YELLOW, WHITE, BLACK, GRAY, ORANGE, PURPLE, PINK.
    - To position objects, call .to_edge(), .move_to(), or .shift() DIRECTLY on the Mobject.
    - For math formulas, use Text() with simple strings like "x² + y² = z²".
    - Make the animation clear and educational.
    """
    
    response_text = generate_with_ai(gemini_client, scene_prompt, is_json=False)
    if not response_text:
        logger.end_operation("regen_scene", "No AI response received")
        return False

    if "```python" in response_text:
        scene_code = response_text.split("```python")[1].split("```")[0].strip()
    else:
        scene_code = response_text.strip()
        
    scene_code = fix_invalid_colors(scene_code)

    # Create a new temporary file for this regenerated scene
    with tempfile.NamedTemporaryFile(mode='w+', suffix='.py', delete=False, dir=temp_dir_name) as temp_code_file:
        temp_code_file.write(scene_code)
        temp_code_file_path = temp_code_file.name

    # Update the project config for this specific scene
    if scene_name in project_config["scenes"]:
        project_config["scenes"][scene_name]["code_file"] = temp_code_file_path
        project_config["scenes"][scene_name]["status"] = "code_regenerated"

    code_size = len(scene_code)
    logger.end_operation("regen_scene", f"Scene code regenerated", f"{code_size} chars → {temp_code_file_path}")
    return True