orchestrator.py

"""
NeuroAnim Orchestrator

This script coordinates the entire STEM animation generation pipeline:
1. Concept Planning
2. Code Generation
3. Rendering
4. Vision-based Analysis
5. Audio Generation
6. Final Merging

It uses the MCP servers (renderer and creative) to accomplish these tasks.
"""

import ast
import asyncio
import json
import logging
import os
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional

import aiofiles
from dotenv import load_dotenv
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from utils.tts import TTSGenerator

load_dotenv()
# Set up logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class NeuroAnimOrchestrator:
    """Main orchestrator for NeuroAnim pipeline."""

    def __init__(
        self, hf_api_key: Optional[str] = None, elevenlabs_api_key: Optional[str] = None
    ):
        self.hf_api_key = hf_api_key or os.getenv("HUGGINGFACE_API_KEY")
        self.elevenlabs_api_key = elevenlabs_api_key or os.getenv("ELEVENLABS_API_KEY")
        self.renderer_session: Optional[ClientSession] = None
        self.creative_session: Optional[ClientSession] = None

        # Initialize TTS generator
        self.tts_generator = TTSGenerator(
            elevenlabs_api_key=self.elevenlabs_api_key,
            hf_api_key=self.hf_api_key,
            fallback_enabled=True,
        )

        # Context managers for MCP client connections
        self._renderer_cm = None
        self._creative_cm = None
        self._renderer_streams = None
        self._creative_streams = None

        # Working directories
        self.work_dir: Optional[Path] = None
        self.output_dir: Optional[Path] = None

    async def initialize(self):
        """Initialize MCP server connections."""
        # Set up working directories
        self.work_dir = Path(tempfile.mkdtemp(prefix="neuroanim_work_"))
        self.output_dir = Path("outputs")
        self.output_dir.mkdir(exist_ok=True)

        logger.info(f"Working directory: {self.work_dir}")
        logger.info(f"Output directory: {self.output_dir}")

        # Initialize renderer server
        # stdio_client is an async context manager, must use async with
        renderer_params = StdioServerParameters(
            command="python", args=["mcp_servers/renderer.py"]
        )

        self._renderer_cm = stdio_client(renderer_params)
        self._renderer_streams = await self._renderer_cm.__aenter__()
        read_stream, write_stream = self._renderer_streams
        self.renderer_session = ClientSession(read_stream, write_stream)
        # Start background receive loop for the client session
        await self.renderer_session.__aenter__()
        await self.renderer_session.initialize()
        logger.info("Renderer MCP server connected")

        # Initialize creative server
        creative_params = StdioServerParameters(
            command="python",
            args=["mcp_servers/creative.py"],
            env={"HUGGINGFACE_API_KEY": self.hf_api_key} if self.hf_api_key else None,
        )

        self._creative_cm = stdio_client(creative_params)
        self._creative_streams = await self._creative_cm.__aenter__()
        read_stream, write_stream = self._creative_streams
        self.creative_session = ClientSession(read_stream, write_stream)
        # Start background receive loop for the client session
        await self.creative_session.__aenter__()
        await self.creative_session.initialize()
        logger.info("Creative MCP server connected")

    async def cleanup(self):
        """Clean up resources."""
        import shutil

        # Close sessions first
        if self.renderer_session:
            try:
                await self.renderer_session.__aexit__(None, None, None)
            except (Exception, asyncio.CancelledError) as e:
                logger.debug(f"Error closing renderer session: {e}")

        if self.creative_session:
            try:
                await self.creative_session.__aexit__(None, None, None)
            except (Exception, asyncio.CancelledError) as e:
                logger.debug(f"Error closing creative session: {e}")

        # Then close the stdio_client context managers with timeout
        if self._renderer_cm:
            try:
                async with asyncio.timeout(2):  # 2 second timeout
                    await self._renderer_cm.__aexit__(None, None, None)
            except (Exception, asyncio.CancelledError, TimeoutError) as e:
                logger.debug(f"Error closing renderer context manager: {e}")

        if self._creative_cm:
            try:
                async with asyncio.timeout(2):  # 2 second timeout
                    await self._creative_cm.__aexit__(None, None, None)
            except (Exception, asyncio.CancelledError, TimeoutError) as e:
                logger.debug(f"Error closing creative context manager: {e}")

        # Clean up working directory
        if self.work_dir and self.work_dir.exists():
            try:
                shutil.rmtree(self.work_dir)
                logger.info(f"Cleaned up working directory: {self.work_dir}")
            except Exception as e:
                logger.warning(f"Failed to clean up working directory: {e}")

    async def call_tool(
        self, session: ClientSession, tool_name: str, arguments: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Call a tool on an MCP server."""
        result = await session.call_tool(tool_name, arguments)

        if hasattr(result, "content") and result.content:
            content = result.content[0]
            if hasattr(content, "text"):
                return {
                    "text": content.text,
                    "isError": getattr(result, "isError", False),
                }

        return {"text": str(result), "isError": False}

    async def generate_animation(
        self,
        topic: str,
        target_audience: str = "general",
        animation_length_minutes: float = 2.0,
        output_filename: str = "animation.mp4",
        quality: str = "medium",
        progress_callback: Optional[callable] = None,
    ) -> Dict[str, Any]:
        """Complete animation generation pipeline."""

        try:
            logger.info(f"Starting animation generation for: {topic}")

            # Step 1: Concept Planning
            logger.info("Step 1: Planning concept...")
            if progress_callback:
                progress_callback("Planning concept", 0.1)
            concept_result = await self.call_tool(
                self.creative_session,
                "plan_concept",
                {
                    "topic": topic,
                    "target_audience": target_audience,
                    "animation_length_minutes": animation_length_minutes,
                },
            )

            if concept_result["isError"]:
                raise Exception(f"Concept planning failed: {concept_result['text']}")

            concept_plan = concept_result["text"]
            logger.info("Concept planning completed")

            # Step 2: Generate Narration
            logger.info("Step 2: Generating narration...")
            if progress_callback:
                progress_callback("Generating narration script", 0.25)
            narration_result = await self.call_tool(
                self.creative_session,
                "generate_narration",
                {
                    "concept": topic,
                    "scene_description": concept_plan,
                    "target_audience": target_audience,
                    "duration_seconds": int(animation_length_minutes * 60),
                },
            )

            if narration_result["isError"]:
                raise Exception(
                    f"Narration generation failed: {narration_result['text']}"
                )

            # Clean narration text - remove title/prefix before TTS
            narration_text = self._clean_narration_text(narration_result["text"])
            logger.info("Narration generation completed")
            logger.info(f"Narration preview: {narration_text[:100]}...")

            # Step 3: Generate Manim Code with retry logic
            logger.info("Step 3: Generating Manim code...")
            if progress_callback:
                progress_callback("Creating Manim animation code", 0.40)
            target_duration_seconds = int(animation_length_minutes * 60)
            manim_code = await self._generate_and_validate_code(
                topic=topic,
                concept_plan=concept_plan,
                duration_seconds=target_duration_seconds,
                max_retries=3,
            )
            logger.info("Manim code generation completed and validated")

            # Step 4: Write Manim File
            logger.info("Step 4: Writing Manim file...")
            manim_file = self.work_dir / "animation.py"
            write_result = await self.call_tool(
                self.renderer_session,
                "write_manim_file",
                {"filepath": str(manim_file), "code": manim_code},
            )

            if write_result["isError"]:
                raise Exception(f"File writing failed: {write_result['text']}")

            # Extract scene name from code
            scene_name = self._extract_scene_name(manim_code)
            logger.info(f"Scene name detected: {scene_name}")

            # Step 5: Render Animation with retry on runtime errors
            logger.info("Step 5: Rendering animation...")
            if progress_callback:
                progress_callback("Rendering animation video", 0.55)
            max_render_retries = 5
            video_file = None
            
            for render_attempt in range(max_render_retries):
                render_result = await self.call_tool(
                    self.renderer_session,
                    "render_manim_animation",
                    {
                        "scene_name": scene_name,
                        "file_path": str(manim_file),
                        "output_dir": str(self.work_dir),
                        "quality": quality,  # Use the quality parameter
                        "format": "mp4",
                        "frame_rate": 30,
                    },
                )

                if not render_result["isError"]:
                    # Success! Find the rendered file
                    video_file = self._find_output_file(self.work_dir, scene_name, "mp4")
                    if video_file:
                        # Check video duration
                        try:
                            actual_duration = self._get_video_duration(video_file)
                            logger.info(f"Rendered video duration: {actual_duration:.2f}s (Target: {target_duration_seconds}s)")
                            
                            if actual_duration < target_duration_seconds * 0.5:
                                logger.warning(f"Video is too short ({actual_duration:.2f}s < {target_duration_seconds * 0.5}s). Forcing retry...")
                                error_text = (
                                    f"The generated animation was TOO SHORT ({actual_duration:.1f}s). "
                                    f"The target duration is {target_duration_seconds}s. "
                                    "You MUST make the animation longer by adding more `self.wait()` calls "
                                    "and ensuring animations play slower (use run_time parameter)."
                                )
                                # Fall through to error handling logic below
                            else:
                                break
                        except Exception as e:
                            logger.warning(f"Could not verify video duration: {e}")
                            break
                    else:
                        logger.warning("Render succeeded but could not find output file")
                        if render_attempt < max_render_retries - 1:
                            continue
                
                # Rendering failed - check if it's a runtime error we can fix
                error_text = render_result["text"]
                logger.warning(f"Render attempt {render_attempt + 1} failed: {error_text[:200]}...")
                
                # Check if this is a Manim runtime error (not a "no scene" error)
                if render_attempt < max_render_retries - 1 and (
                    "TypeError" in error_text 
                    or "AttributeError" in error_text 
                    or "ValueError" in error_text
                    or "KeyError" in error_text
                ):
                    logger.info(f"Detected runtime error in Manim code. Regenerating code (attempt {render_attempt + 2}/{max_render_retries})...")
                    
                    # Regenerate code with error feedback
                    runtime_error_msg = f"Runtime Error during Manim rendering:\n{error_text}\n\nPlease fix the code to be compatible with Manim version 0.19.0."
                    manim_code = await self._generate_and_validate_code(
                        topic=topic,
                        concept_plan=concept_plan,
                        duration_seconds=target_duration_seconds,
                        max_retries=3,  # Allow retries for syntax errors during fix
                        previous_error=runtime_error_msg,
                        previous_code=manim_code,
                    )
                    
                    # Write the new code
                    write_result = await self.call_tool(
                        self.renderer_session,
                        "write_manim_file",
                        {"filepath": str(manim_file), "code": manim_code},
                    )
                    
                    if write_result["isError"]:
                        raise Exception(f"File writing failed: {write_result['text']}")
                    
                    # Extract scene name from new code
                    scene_name = self._extract_scene_name(manim_code)
                    logger.info(f"Regenerated code with scene: {scene_name}")
                    
                    # Loop will retry rendering with new code
                    continue
                else:
                    # Not a runtime error or out of retries
                    raise Exception(f"Rendering failed: {error_text}")
            
            if not video_file:
                raise Exception("Could not find rendered video file after all attempts")

            logger.info(f"Animation rendered: {video_file}")

            # Step 6: Generate Speech Audio
            logger.info("Step 6: Generating speech audio...")
            if progress_callback:
                progress_callback("Generating audio narration", 0.75)
            audio_file = self.work_dir / "narration.mp3"

            # Use TTS generator with automatic fallback
            try:
                tts_result = await self.tts_generator.generate_speech(
                    text=narration_text, output_path=audio_file, voice="rachel"
                )
                logger.info(
                    f"Audio generated with {tts_result['provider']}: {audio_file}"
                )

                # Validate audio file
                validation = self.tts_generator.validate_audio_file(audio_file)
                if not validation["valid"]:
                    logger.warning(
                        f"Audio validation warning: {validation.get('error', 'Unknown issue')}"
                    )
                    logger.info("Audio file may have issues but continuing...")
                else:
                    logger.info(
                        f"Audio validated: {validation.get('duration', 'N/A')}s, {validation.get('size', 0)} bytes"
                    )

            except Exception as e:
                logger.error(f"TTS generation failed: {e}")
                raise Exception(f"Speech generation failed: {str(e)}")

            # Step 7: Merge Video and Audio
            logger.info("Step 7: Merging video and audio...")
            if progress_callback:
                progress_callback("Merging video and audio", 0.90)
            final_output = self.output_dir / output_filename
            merge_result = await self.call_tool(
                self.renderer_session,
                "merge_video_audio",
                {
                    "video_file": str(video_file),
                    "audio_file": str(audio_file),
                    "output_file": str(final_output),
                },
            )

            if merge_result["isError"]:
                raise Exception(f"Merging failed: {merge_result['text']}")

            # Step 8: Generate Quiz
            logger.info("Step 8: Generating quiz...")
            if progress_callback:
                progress_callback("Creating quiz questions", 0.95)
            quiz_result = await self.call_tool(
                self.creative_session,
                "generate_quiz",
                {
                    "concept": topic,
                    "difficulty": "medium",
                    "num_questions": 3,
                    "question_types": ["multiple_choice"],
                },
            )

            quiz_content = (
                quiz_result["text"]
                if not quiz_result["isError"]
                else "Quiz generation failed"
            )

            # Return results
            results = {
                "success": True,
                "topic": topic,
                "target_audience": target_audience,
                "concept_plan": concept_plan,
                "narration": narration_text,
                "manim_code": manim_code,
                "output_file": str(final_output),
                "quiz": quiz_content,
                "work_dir": str(self.work_dir),
            }

            logger.info(f"Animation generation completed successfully: {final_output}")
            return results

        except Exception as e:
            logger.error(f"Animation generation failed: {str(e)}")
            return {
                "success": False,
                "error": str(e),
                "work_dir": str(self.work_dir) if self.work_dir else None,
            }

    def _clean_narration_text(self, text: str) -> str:
        """
        Clean narration text by removing title prefixes and formatting artifacts.

        The creative server returns text with prefixes like "Narration Script:\n\n"
        which should not be sent to TTS.
        """
        # Remove common prefixes
        prefixes_to_remove = [
            "Narration Script:",
            "Script:",
            "Narration:",
            "Text:",
        ]

        cleaned = text.strip()

        # Remove any of the prefixes (case-insensitive)
        for prefix in prefixes_to_remove:
            if cleaned.lower().startswith(prefix.lower()):
                cleaned = cleaned[len(prefix) :].strip()
                break

        # Remove leading newlines and whitespace
        cleaned = cleaned.lstrip("\n").strip()

        # Remove any markdown code block markers
        if cleaned.startswith("```"):
            lines = cleaned.split("\n")
            # Remove first line (opening ```)
            if len(lines) > 1:
                lines = lines[1:]
            # Remove last line if it's closing ```
            if lines and lines[-1].strip() == "```":
                lines = lines[:-1]
            cleaned = "\n".join(lines).strip()

        return cleaned

    def _extract_python_code(self, text: str) -> str:
        """Extract Python code from markdown response."""
        # Look for code blocks
        if "```python" in text:
            start = text.find("```python") + 9
            end = text.find("```", start)
            if end == -1:
                end = len(text)
            return text[start:end].strip()
        elif "```" in text:
            start = text.find("```") + 3
            end = text.find("```", start)
            if end == -1:
                end = len(text)
            return text[start:end].strip()
        else:
            return text.strip()

    async def _generate_and_validate_code(
        self,
        topic: str,
        concept_plan: str,
        duration_seconds: int = 60,
        max_retries: int = 3,
        previous_error: Optional[str] = None,
        previous_code: Optional[str] = None,
    ) -> str:
        """Generate Manim code with retry logic for syntax errors."""
        for attempt in range(max_retries):
            try:
                logger.info(f"Code generation attempt {attempt + 1}/{max_retries}")

                # Build arguments for code generation
                arguments = {
                    "concept": topic,
                    "scene_description": concept_plan,
                    "visual_elements": ["text", "shapes", "animations"],
                    "duration_seconds": duration_seconds,
                }

                # If this is a retry, include error feedback
                if previous_error:
                    if previous_code:
                        arguments["previous_code"] = previous_code
                    arguments["error_message"] = previous_error
                    logger.info(
                        f"Retrying with error feedback: {previous_error[:100]}..."
                    )

                # Generate code
                code_result = await self.call_tool(
                    self.creative_session, "generate_manim_code", arguments
                )

                if code_result["isError"]:
                    if attempt < max_retries - 1:
                        logger.warning(
                            f"Code generation failed, retrying: {code_result['text']}"
                        )
                        previous_error = code_result["text"]
                        # Keep previous_code if we had it, for better context in retry
                        continue
                    else:
                        raise Exception(
                            f"Code generation failed: {code_result['text']}"
                        )

                # Extract Python code from response
                manim_code = self._extract_python_code(code_result["text"])

                # Validate Python syntax
                syntax_errors = self._validate_python_syntax(manim_code)
                if syntax_errors:
                    if attempt < max_retries - 1:
                        logger.warning(
                            f"Syntax error detected, retrying: {syntax_errors}"
                        )
                        previous_error = f"Syntax Error:\n{syntax_errors}"
                        previous_code = manim_code
                        continue
                    else:
                        raise Exception(
                            f"Generated code has syntax errors after {max_retries} attempts:\n{syntax_errors}"
                        )

                # Validate that code contains a Scene class
                has_scene = self._validate_has_scene_class(manim_code)
                if not has_scene:
                    if attempt < max_retries - 1:
                        logger.warning(
                            "No Scene class found in generated code, retrying..."
                        )
                        previous_error = (
                            "Error: The generated code does not contain any Scene class. "
                            "Please ensure you create a class that inherits from manim.Scene, "
                            "manim.MovingCameraScene, or manim.ThreeDScene."
                        )
                        previous_code = manim_code
                        continue
                    else:
                        raise Exception(
                            f"Generated code does not contain a Scene class after {max_retries} attempts"
                        )

                # Success!
                logger.info(f"Valid code generated on attempt {attempt + 1}")
                return manim_code

            except Exception as e:
                if attempt < max_retries - 1:
                    logger.warning(f"Attempt {attempt + 1} failed: {str(e)}")
                    previous_error = str(e)
                    continue
                else:
                    raise

        raise Exception("Failed to generate valid code after all retries")

    def _validate_python_syntax(self, code: str) -> Optional[str]:
        """Validate Python code syntax. Returns error message if invalid, None if valid."""
        try:
            ast.parse(code)
            return None
        except SyntaxError as e:
            # Build detailed error message with context
            error_msg = f"Line {e.lineno}: {e.msg}"

            # Show surrounding context (3 lines before and after)
            if e.lineno is not None:
                code_lines = code.split("\n")
                start_line = max(0, e.lineno - 4)  # 3 lines before
                end_line = min(len(code_lines), e.lineno + 2)  # 2 lines after

                error_msg += "\n\nContext:"
                for i in range(start_line, end_line):
                    line_num = i + 1
                    prefix = ">>> " if line_num == e.lineno else "    "
                    error_msg += f"\n{prefix}{line_num:3d} | {code_lines[i]}"

                    # Add pointer for error line
                    if line_num == e.lineno and e.offset:
                        error_msg += f"\n    {' ' * 4}{' ' * (e.offset - 1)}^"

            return error_msg
        except Exception as e:
            return f"Unexpected error during syntax validation: {str(e)}"

    def _validate_has_scene_class(self, code: str) -> bool:
        """Check if code contains at least one Scene class."""
        import re
        
        # Check for Scene class inheritance
        scene_patterns = [
            r"class\s+\w+\s*\(\s*Scene\s*\)",
            r"class\s+\w+\s*\(\s*MovingCameraScene\s*\)",
            r"class\s+\w+\s*\(\s*ThreeDScene\s*\)",
            r"class\s+\w+\s*\(\s*\w*Scene\s*\)",
        ]
        
        for pattern in scene_patterns:
            if re.search(pattern, code):
                return True
        
        # Also check using AST parsing as a backup
        try:
            tree = ast.parse(code)
            for node in ast.walk(tree):
                if isinstance(node, ast.ClassDef):
                    # Check if any base class contains "Scene"
                    for base in node.bases:
                        if isinstance(base, ast.Name) and "Scene" in base.id:
                            return True
        except Exception:
            pass
        
        return False

    def _extract_scene_name(self, code: str) -> str:
        """Extract scene class name from Manim code."""
        import re

        # Try multiple patterns to find Scene class
        patterns = [
            r"class\s+(\w+)\s*\(\s*Scene\s*\)",  # class Name(Scene)
            r"class\s+(\w+)\s*\(\s*MovingCameraScene\s*\)",  # class Name(MovingCameraScene)
            r"class\s+(\w+)\s*\(\s*ThreeDScene\s*\)",  # class Name(ThreeDScene)
            r"class\s+(\w+)\s*\(\s*\w*Scene\s*\)",  # class Name(AnyScene)
        ]
        
        for pattern in patterns:
            match = re.search(pattern, code)
            if match:
                scene_name = match.group(1)
                logger.info(f"Found scene class: {scene_name}")
                return scene_name
        
        # If no scene found, look for any class definition and warn
        any_class = re.search(r"class\s+(\w+)\s*\(", code)
        if any_class:
            class_name = any_class.group(1)
            logger.warning(
                f"Could not find Scene class, using first class found: {class_name}"
            )
            return class_name
        
        # Last resort - parse the AST to find classes
        try:
            tree = ast.parse(code)
            for node in ast.walk(tree):
                if isinstance(node, ast.ClassDef):
                    logger.warning(
                        f"Using first class from AST parsing: {node.name}"
                    )
                    return node.name
        except Exception as e:
            logger.error(f"Failed to parse code AST: {e}")
        
        # Absolute fallback
        logger.error("No scene class found in code! This will likely cause rendering to fail.")
        return "Scene"  # fallback

    def _find_output_file(
        self, directory: Path, scene_name: str, extension: str
    ) -> Optional[Path]:
        """Find output file with given scene name and extension."""
        for file in directory.glob(f"{scene_name}*.{extension}"):
            return file
        return None


async def main():
    """Main function for running the orchestrator."""
    import argparse

    parser = argparse.ArgumentParser(description="NeuroAnim STEM Animation Generator")
    parser.add_argument("topic", help="STEM topic for the animation")
    parser.add_argument(
        "--audience",
        choices=["elementary", "middle_school", "high_school", "college", "general"],
        default="general",
        help="Target audience",
    )
    parser.add_argument(
        "--duration", type=float, default=2.0, help="Animation duration in minutes"
    )
    parser.add_argument("--output", default="animation.mp4", help="Output filename")
    parser.add_argument(
        "--api-key", help="Hugging Face API key (or set HUGGINGFACE_API_KEY env var)"
    )
    parser.add_argument(
        "--elevenlabs-key",
        help="ElevenLabs API key (or set ELEVENLABS_API_KEY env var)",
    )

    args = parser.parse_args()

    # Initialize and run orchestrator
    orchestrator = NeuroAnimOrchestrator(
        hf_api_key=args.api_key, elevenlabs_api_key=args.elevenlabs_key
    )

    try:
        await orchestrator.initialize()

        results = await orchestrator.generate_animation(
            topic=args.topic,
            target_audience=args.audience,
            animation_length_minutes=args.duration,
            output_filename=args.output,
        )

        if results["success"]:
            print("\n🎉 Animation Generated Successfully!")
            print(f"📹 Output file: {results['output_file']}")
            print(f"🎯 Topic: {results['topic']}")
            print(f"👥 Audience: {results['target_audience']}")
            print(f"\n📝 Concept Plan:")
            print(
                results["concept_plan"][:500] + "..."
                if len(results["concept_plan"]) > 500
                else results["concept_plan"]
            )
            print(f"\n🎭 Narration:")
            print(
                results["narration"][:300] + "..."
                if len(results["narration"]) > 300
                else results["narration"]
            )
            print(f"\n📚 Quiz Questions:")
            print(results["quiz"])
        else:
            print(f"\n❌ Animation Generation Failed: {results['error']}")

    except KeyboardInterrupt:
        print("\n⚠️  Process interrupted by user")
    except Exception as e:
        print(f"\n💥 Unexpected error: {str(e)}")
    finally:
        await orchestrator.cleanup()


if __name__ == "__main__":
    asyncio.run(main())