manimator/api/animation_generation.py

import os
import litellm
from fastapi import HTTPException
import logging

from ..utils.system_prompts import get_system_prompt
from ..utils.code_postprocessor import post_process_code
from ..utils.code_validator import CodeValidator
from ..utils.code_fixer import CodeFixer
from ..inputs.processor import InputProcessor
from ..utils.content_preprocessor import preprocess_long_content, get_script_mode_prompt_for_long_content

logger = logging.getLogger(__name__)


def _generate_with_reflexion(prompt: str, category: str, job_id: str = "unknown") -> str:
    """
    Generate code using the Reflexion Agent.
    
    Args:
        prompt: Processed user prompt
        category: Animation category
        job_id: Job identifier for logging
        
    Returns:
        Generated and refined code
    """
    from ..agents.reflexion_agent import ReflexionAgent
    
    agent = ReflexionAgent()
    code, state = agent.generate_with_reflection(
        goal=prompt,
        category=category,
        job_id=job_id
    )
    
    logger.info(f"🔄 Reflexion complete: {state.iteration} iterations, {state.total_issues_found} issues found, {state.total_issues_fixed} fixed")
    
    return code


def _generate_legacy(prompt: str, category: str, max_attempts: int = 3) -> str:
    """
    Legacy code generation without Reflexion.
    Used as fallback when Reflexion is disabled or fails.
    
    Args:
        prompt: Processed user prompt
        category: Animation category
        max_attempts: Maximum generation attempts
        
    Returns:
        Generated code
    """
    validator = CodeValidator()
    fixer = CodeFixer()
    
    primary_model = os.getenv("CODE_GEN_MODEL")
    fallback_model = os.getenv("CODE_GEN_FALLBACK_MODEL", primary_model)
    
    for attempt in range(max_attempts):
        try:
            # Use fallback model on retry attempts
            model = fallback_model if attempt > 0 else primary_model
            
            # Get dynamic system prompt based on category
            system_prompt = get_system_prompt(category)
            
            # Detect if input is a ready-made script (long content) vs short prompt
            word_count = len(prompt.split())
            is_script_mode = word_count > 200
            
            # For very long content, preprocess into sections
            processed_prompt = prompt
            section_count = 0
            if word_count > 1000:
                processed_prompt, section_count = preprocess_long_content(prompt)
            
            if section_count > 0:
                # Very long content - use sectioned prompt
                logger.info(f"📝 LONG DOCUMENT MODE (Legacy): {word_count} words -> {section_count} sections")
                user_content = get_script_mode_prompt_for_long_content(processed_prompt, section_count)
            elif is_script_mode:
                logger.info(f"📝 SCRIPT MODE (Legacy): Input has {word_count} words - treating as ready-made script")
                user_content = f"""# 🎬 SCRIPT MODE - ANIMATE THE USER'S CONTENT

## IMPORTANT: The user has provided their COMPLETE script/content below.
This is NOT a topic to research - this IS the exact narration/content they want animated.

## YOUR TASK:
1. **Use the content below AS the voiceover text** - split it into logical sections
2. **Create beautiful animations that MATCH each section** of their content
3. **Do NOT rewrite, summarize, or generate new information** - animate THEIR words
4. **Every paragraph/section should become a voiceover block** with matching visuals

## USER'S SCRIPT TO ANIMATE:
---
{prompt}
---

NOTE!!!:
1. NO BLANK SCREENS: Keep the screen populated. If a voiceover is playing, show something.
2. NO OVERLAPS: Ensure text and objects do not overlap. Use `next_to` and `arrange`.
3. CLEAN TRANSITIONS: Fade out old content before showing new content, but don't leave the screen empty for long.
4. VARIED ANIMATIONS: Use a mix of Write, FadeIn, GrowFromCenter, etc.
5. STAY ON SCREEN: Ensure all text and objects are within the screen boundaries. Use .scale_to_fit_width(config.frame_width - 1) for large groups."""
            else:
                logger.info(f"📝 GENERATION MODE (Legacy): Input has {word_count} words - LLM will generate content")
                user_content = f"Create a video about:\n\n{prompt}\n\n NOTE!!!:\n1. NO BLANK SCREENS: Keep the screen populated. If a voiceover is playing, show something.\n2. NO OVERLAPS: Ensure text and objects do not overlap. Use `next_to` and `arrange`.\n3. CLEAN TRANSITIONS: Fade out old content before showing new content, but don't leave the screen empty for long.\n4. VARIED ANIMATIONS: Use a mix of Write, FadeIn, GrowFromCenter, etc.\n5. STAY ON SCREEN: Ensure all text and objects are within the screen boundaries. Use .scale_to_fit_width(config.frame_width - 1) for large groups."
            
            messages = [
                {
                    "role": "system",
                    "content": system_prompt,
                },
                {
                    "role": "user",
                    "content": user_content,
                },
            ]
            
            logger.info(f"Generating code (attempt {attempt + 1}/{max_attempts}) with model {model}")
            
            # Only set max_tokens for long documents
            kwargs = {
                "model": model,
                "messages": messages,
                "num_retries": 2
            }
            
            if section_count > 0:
                kwargs["max_tokens"] = 12000
            
            response = litellm.completion(**kwargs)
            
            raw_code = response.choices[0].message.content
            
            # Extract code if wrapped in markdown (handle various formats)
            import re
            
            # Try different markdown patterns
            code_patterns = [
                r'```python\n(.*?)```',     # Standard: ```python ... ```
                r'````python\n(.*?)````',   # Quad backticks 
                r'```py\n(.*?)```',         # ```py
                r'```\n(.*?)```',           # Just backticks without language
            ]
            
            for pattern in code_patterns:
                match = re.search(pattern, raw_code, re.DOTALL)
                if match:
                    raw_code = match.group(1).strip()
                    break
            
            # If still has backticks, try to clean up
            if raw_code.startswith('```'):
                lines = raw_code.split('\n')
                # Remove first line if it's just ```python or similar
                if lines[0].strip().startswith('```'):
                    lines = lines[1:]
                # Remove last line if it's just ```
                if lines and lines[-1].strip() == '```':
                    lines = lines[:-1]
                raw_code = '\n'.join(lines)
            
            # Post-process the code to fix common issues
            processed_code = post_process_code(raw_code)
            
            # Validate code
            is_valid, errors = validator.validate(processed_code)
            
            if is_valid:
                logger.info("Code validation passed")
                return processed_code
            
            # Try to auto-fix
            logger.warning(f"Code validation failed with {len(errors)} errors, attempting auto-fix")
            fixed_code, is_fixed, remaining_errors = fixer.fix_and_validate(processed_code, max_attempts=2)
            
            if is_fixed:
                logger.info("Code auto-fixed successfully")
                return fixed_code
            
            # If last attempt, return best code we have
            if attempt == max_attempts - 1:
                error_msg = f"Code generation failed after {max_attempts} attempts. Errors: {remaining_errors}"
                logger.error(error_msg)
                raise HTTPException(
                    status_code=500,
                    detail=error_msg
                )
            
            logger.info(f"Retrying code generation (attempt {attempt + 2}/{max_attempts})")
            
        except HTTPException:
            raise
        except Exception as e:
            logger.error(f"Error in code generation attempt {attempt + 1}: {str(e)}")
            if attempt == max_attempts - 1:
                raise HTTPException(
                    status_code=500, 
                    detail=f"Failed to generate animation response after {max_attempts} attempts: {str(e)}"
                )
    
    # Should not reach here, but just in case
    raise HTTPException(
        status_code=500,
        detail="Failed to generate valid animation code after all attempts"
    )


def generate_animation_response(
    input_data: str,
    input_type: str = "text",
    category: str = "mathematical",
    max_attempts: int = 3,
    job_id: str = "unknown"
) -> str:
    """Generate Manim animation code from input with validation and auto-fixing.
    
    Uses Reflexion Agent when enabled (REFLEXION_ENABLED=true) for improved
    code quality through self-critique and iterative improvement.

    Args:
        input_data (str): User's input (text, URL, or PDF path)
        input_type (str): Type of input ('text', 'url', 'pdf')
        category (str): Animation category (tech_system, product_startup, mathematical)
        max_attempts (int): Maximum generation attempts (for legacy mode)
        job_id (str): Job identifier for logging

    Returns:
        str: Generated Manim animation code (validated and post-processed)

    Raises:
        HTTPException: If code generation fails after all attempts
    """
    # Process input to get the actual prompt text
    try:
        prompt = InputProcessor.process(input_type, input_data)
    except Exception as e:
        logger.error(f"Input processing failed: {e}")
        raise HTTPException(status_code=400, detail=f"Input processing failed: {str(e)}")
    
    # Check if Reflexion is enabled
    reflexion_enabled = os.getenv("REFLEXION_ENABLED", "true").lower() == "true"
    
    if reflexion_enabled:
        logger.info("🔄 Using Reflexion Agent for code generation")
        try:
            code = _generate_with_reflexion(prompt, category, job_id)
            
            # Final validation
            validator = CodeValidator()
            fixer = CodeFixer()
            
            is_valid, errors = validator.validate(code)
            if not is_valid:
                logger.warning(f"Reflexion code has {len(errors)} validation errors, attempting auto-fix")
                code, is_fixed, _ = fixer.fix_and_validate(code, max_attempts=2)
            
            return code
            
        except Exception as e:
            logger.error(f"Reflexion failed, falling back to legacy: {e}")
            # Fall through to legacy generation
    else:
        logger.info("📝 Using legacy code generation (Reflexion disabled)")
    
    # Legacy generation (fallback or when Reflexion is disabled)
    return _generate_legacy(prompt, category, max_attempts)