neuroanim/agents/nodes.py

"""
LangGraph Agent Nodes for NeuroAnim Pipeline

This module contains all the node functions used in the LangGraph workflow.
Each node represents a step in the animation generation pipeline and
communicates with the Manim MCP server to perform its task.
"""

import ast
import json
import logging
import re
import tempfile
import time
from pathlib import Path
from typing import Any, Dict

from mcp import ClientSession

from neuroanim.graph.state import AnimationState
from utils.tts import TTSGenerator

logger = logging.getLogger(__name__)


class AnimationNodes:
    """Container for all animation pipeline nodes."""

    def __init__(
        self,
        mcp_session: ClientSession,
        tts_generator: TTSGenerator,
        work_dir: Path,
        output_dir: Path,
    ):
        """
        Initialize the animation nodes.

        Args:
            mcp_session: MCP client session for tool calls
            tts_generator: TTS generator instance
            work_dir: Working directory for temporary files
            output_dir: Output directory for final files
        """
        self.mcp_session = mcp_session
        self.tts_generator = tts_generator
        self.work_dir = work_dir
        self.output_dir = output_dir

    async def call_mcp_tool(
        self, tool_name: str, arguments: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Call an MCP tool and return the result.

        Args:
            tool_name: Name of the tool to call
            arguments: Arguments to pass to the tool

        Returns:
            Dictionary with 'text' and 'isError' keys
        """
        try:
            result = await self.mcp_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}

        except Exception as e:
            logger.error(f"MCP tool call failed for {tool_name}: {e}")
            return {"text": f"Tool call failed: {str(e)}", "isError": True}

    async def initialize_node(self, state: AnimationState) -> AnimationState:
        """
        Initialize the pipeline state and working directories.

        Args:
            state: Current animation state

        Returns:
            Updated state with initialized paths and metadata
        """
        logger.info("🚀 Initializing animation pipeline")

        state["start_time"] = time.time()
        state["work_dir"] = str(self.work_dir)
        state["output_dir"] = str(self.output_dir)
        state["current_step"] = "initialization"
        state["completed_steps"].append("initialization")

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

        return state

    async def plan_concept_node(self, state: AnimationState) -> AnimationState:
        """
        Plan the animation concept using AI.

        Args:
            state: Current animation state

        Returns:
            Updated state with concept plan
        """
        logger.info("📋 Planning concept...")
        state["current_step"] = "concept_planning"

        try:
            result = await self.call_mcp_tool(
                "plan_concept",
                {
                    "topic": state["topic"],
                    "target_audience": state["target_audience"],
                    "animation_length_minutes": state["animation_length_minutes"],
                },
            )

            if result["isError"]:
                state["errors"].append(f"Concept planning failed: {result['text']}")
                return state

            concept_plan = result["text"]
            state["concept_plan"] = concept_plan

            # Try to parse JSON from the concept plan
            try:
                # Extract JSON if it's embedded in the response
                json_match = re.search(r"\{.*\}", concept_plan, re.DOTALL)
                if json_match:
                    plan_data = json.loads(json_match.group())
                    state["learning_objectives"] = plan_data.get(
                        "learning_objectives", []
                    )
                    state["visual_metaphors"] = plan_data.get("visual_metaphors", [])
                    state["scene_flow"] = plan_data.get("scene_flow", [])
            except json.JSONDecodeError:
                logger.warning("Could not parse concept plan as JSON")

            state["completed_steps"].append("concept_planning")
            logger.info("✅ Concept planning completed")

        except Exception as e:
            logger.error(f"Concept planning failed: {e}")
            state["errors"].append(f"Concept planning error: {str(e)}")

        return state

    async def generate_narration_node(self, state: AnimationState) -> AnimationState:
        """
        Generate narration script for the animation.

        Args:
            state: Current animation state

        Returns:
            Updated state with narration text
        """
        logger.info("🎙️  Generating narration...")
        state["current_step"] = "narration_generation"

        try:
            duration_seconds = int(state["animation_length_minutes"] * 60)

            result = await self.call_mcp_tool(
                "generate_narration",
                {
                    "concept": state["topic"],
                    "scene_description": state.get("concept_plan", ""),
                    "target_audience": state["target_audience"],
                    "duration_seconds": duration_seconds,
                },
            )

            if result["isError"]:
                state["errors"].append(f"Narration generation failed: {result['text']}")
                return state

            # Extract narration text from response
            narration_text = result["text"]
            if "Narration Script:" in narration_text:
                narration_text = narration_text.split("Narration Script:")[-1].strip()

            state["narration_text"] = narration_text
            state["narration_duration"] = duration_seconds
            state["completed_steps"].append("narration_generation")

            logger.info("✅ Narration generation completed")

        except Exception as e:
            logger.error(f"Narration generation failed: {e}")
            state["errors"].append(f"Narration generation error: {str(e)}")

        return state

    async def generate_code_node(self, state: AnimationState) -> AnimationState:
        """
        Generate Manim code for the animation.

        Args:
            state: Current animation state

        Returns:
            Updated state with generated code
        """
        logger.info("💻 Generating Manim code...")
        state["current_step"] = "code_generation"

        try:
            # Check if this is a retry
            previous_code = None
            error_message = None
            if state["code_generation_attempts"] > 0:
                previous_code = state.get("manim_code")
                if state.get("previous_code_errors"):
                    error_message = state["previous_code_errors"][-1]

            state["code_generation_attempts"] += 1

            arguments = {
                "concept": state["topic"],
                "scene_description": state.get("concept_plan", ""),
                "visual_elements": ["text", "shapes", "animations"],
            }

            if previous_code and error_message:
                arguments["previous_code"] = previous_code
                arguments["error_message"] = error_message
                logger.info(
                    f"Retrying code generation (attempt {state['code_generation_attempts']})"
                )

            result = await self.call_mcp_tool("generate_manim_code", arguments)

            if result["isError"]:
                state["errors"].append(f"Code generation failed: {result['text']}")
                return state

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

            # Validate syntax
            validation_error = self._validate_python_syntax(manim_code)
            if validation_error:
                logger.warning(f"Code validation failed: {validation_error}")
                if not state.get("previous_code_errors"):
                    state["previous_code_errors"] = []
                state["previous_code_errors"].append(validation_error)
                state["warnings"].append(f"Code validation issue: {validation_error}")
            else:
                logger.info("✅ Code validation passed")

            state["manim_code"] = manim_code
            state["scene_name"] = self._extract_scene_name(manim_code)
            state["completed_steps"].append("code_generation")

            logger.info(f"✅ Code generation completed (scene: {state['scene_name']})")

        except Exception as e:
            logger.error(f"Code generation failed: {e}")
            state["errors"].append(f"Code generation error: {str(e)}")

        return state

    async def write_file_node(self, state: AnimationState) -> AnimationState:
        """
        Write the Manim code to a file.

        Args:
            state: Current animation state

        Returns:
            Updated state with file path
        """
        logger.info("📝 Writing Manim file...")
        state["current_step"] = "file_writing"

        try:
            manim_file = Path(state["work_dir"]) / "animation.py"
            state["manim_file_path"] = str(manim_file)

            result = await self.call_mcp_tool(
                "write_manim_file",
                {"filepath": str(manim_file), "code": state["manim_code"]},
            )

            if result["isError"]:
                state["errors"].append(f"File writing failed: {result['text']}")
                return state

            state["completed_steps"].append("file_writing")
            logger.info(f"✅ Manim file written to {manim_file}")

        except Exception as e:
            logger.error(f"File writing failed: {e}")
            state["errors"].append(f"File writing error: {str(e)}")

        return state

    async def render_animation_node(self, state: AnimationState) -> AnimationState:
        """
        Render the Manim animation.

        Args:
            state: Current animation state

        Returns:
            Updated state with rendered video path
        """
        logger.info("🎬 Rendering animation...")
        state["current_step"] = "rendering"

        try:
            result = await self.call_mcp_tool(
                "render_manim_animation",
                {
                    "scene_name": state["scene_name"],
                    "file_path": state["manim_file_path"],
                    "output_dir": state["work_dir"],
                    "quality": state["rendering_quality"],
                    "format": state["rendering_format"],
                    "frame_rate": state["frame_rate"],
                },
            )

            if result["isError"]:
                state["errors"].append(f"Rendering failed: {result['text']}")
                return state

            # Find the rendered video file
            video_file = (
                Path(state["work_dir"])
                / f"{state['scene_name']}.{state['rendering_format']}"
            )
            if not video_file.exists():
                state["errors"].append(f"Rendered video not found at {video_file}")
                return state

            state["video_file_path"] = str(video_file)
            state["completed_steps"].append("rendering")

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

        except Exception as e:
            logger.error(f"Rendering failed: {e}")
            state["errors"].append(f"Rendering error: {str(e)}")

        return state

    async def generate_audio_node(self, state: AnimationState) -> AnimationState:
        """
        Generate speech audio from narration text.

        Args:
            state: Current animation state

        Returns:
            Updated state with audio file path
        """
        logger.info("🔊 Generating speech audio...")
        state["current_step"] = "audio_generation"

        try:
            audio_file = Path(state["work_dir"]) / "narration.mp3"
            state["audio_file_path"] = str(audio_file)

            # Use TTS generator with automatic fallback
            tts_result = await self.tts_generator.generate_speech(
                text=state["narration_text"], output_path=audio_file, voice="rachel"
            )

            logger.info(f"Audio generated with {tts_result['provider']}")

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

            state["completed_steps"].append("audio_generation")
            logger.info(f"✅ Audio generated: {audio_file}")

        except Exception as e:
            logger.error(f"Audio generation failed: {e}")
            state["errors"].append(f"Audio generation error: {str(e)}")

        return state

    async def merge_video_audio_node(self, state: AnimationState) -> AnimationState:
        """
        Merge video and audio into final output.

        Args:
            state: Current animation state

        Returns:
            Updated state with final output path
        """
        logger.info("🎞️  Merging video and audio...")
        state["current_step"] = "video_audio_merge"

        try:
            final_output = Path(state["output_dir"]) / state["output_filename"]
            state["final_output_path"] = str(final_output)

            result = await self.call_mcp_tool(
                "merge_video_audio",
                {
                    "video_file": state["video_file_path"],
                    "audio_file": state["audio_file_path"],
                    "output_file": str(final_output),
                },
            )

            if result["isError"]:
                state["errors"].append(f"Video/audio merge failed: {result['text']}")
                return state

            state["completed_steps"].append("video_audio_merge")
            logger.info(f"✅ Video and audio merged: {final_output}")

        except Exception as e:
            logger.error(f"Video/audio merge failed: {e}")
            state["errors"].append(f"Merge error: {str(e)}")

        return state

    async def generate_quiz_node(self, state: AnimationState) -> AnimationState:
        """
        Generate quiz questions for the topic.

        Args:
            state: Current animation state

        Returns:
            Updated state with quiz content
        """
        logger.info("❓ Generating quiz...")
        state["current_step"] = "quiz_generation"

        try:
            result = await self.call_mcp_tool(
                "generate_quiz",
                {
                    "concept": state["topic"],
                    "difficulty": "medium",
                    "num_questions": 3,
                    "question_types": ["multiple_choice"],
                },
            )

            if result["isError"]:
                state["warnings"].append(f"Quiz generation failed: {result['text']}")
                state["quiz_content"] = "Quiz generation failed"
            else:
                state["quiz_content"] = result["text"]
                # Try to parse quiz questions
                try:
                    json_match = re.search(r"\[.*\]", result["text"], re.DOTALL)
                    if json_match:
                        state["quiz_questions"] = json.loads(json_match.group())
                except json.JSONDecodeError:
                    logger.warning("Could not parse quiz as JSON")

            state["completed_steps"].append("quiz_generation")
            logger.info("✅ Quiz generation completed")

        except Exception as e:
            logger.error(f"Quiz generation failed: {e}")
            state["warnings"].append(f"Quiz generation error: {str(e)}")
            state["quiz_content"] = "Quiz generation failed"

        return state

    async def finalize_node(self, state: AnimationState) -> AnimationState:
        """
        Finalize the pipeline and compute metadata.

        Args:
            state: Current animation state

        Returns:
            Final state with metadata
        """
        logger.info("🏁 Finalizing pipeline...")
        state["current_step"] = "finalization"

        state["end_time"] = time.time()
        state["total_duration"] = state["end_time"] - state["start_time"]

        # Check if pipeline succeeded
        if not state["errors"] and state.get("final_output_path"):
            state["success"] = True
            logger.info(
                f"✅ Pipeline completed successfully in {state['total_duration']:.2f}s"
            )
        else:
            state["success"] = False
            logger.error(f"❌ Pipeline failed with {len(state['errors'])} error(s)")

        state["completed_steps"].append("finalization")

        return state

    # Helper methods

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

    def _extract_scene_name(self, code: str) -> str:
        """Extract the scene class name from Manim code."""
        try:
            tree = ast.parse(code)
            for node in ast.walk(tree):
                if isinstance(node, ast.ClassDef):
                    # Check if it inherits from Scene or MovingCameraScene
                    for base in node.bases:
                        if isinstance(base, ast.Name) and base.id in [
                            "Scene",
                            "MovingCameraScene",
                            "ThreeDScene",
                        ]:
                            return node.name
        except SyntaxError:
            pass

        # Fallback: use regex
        match = re.search(r"class\s+(\w+)\s*\(.*Scene.*\):", code)
        if match:
            return match.group(1)

        return "GenScene"

    def _validate_python_syntax(self, code: str) -> str | None:
        """
        Validate Python code syntax.

        Returns:
            Error message if validation fails, None if valid
        """
        try:
            ast.parse(code)
            return None
        except SyntaxError as e:
            return f"Syntax error at line {e.lineno}: {e.msg}"
        except Exception as e:
            return f"Validation error: {str(e)}"