import gradio as gr
import subprocess
import tempfile
import os
import shutil
from pathlib import Path
import time

# Default example code
DEFAULT_CODE = '''from manim import *

class SimpleAnimation(Scene):
    def construct(self):
        # Create a simple text animation
        title = Text("Welcome to Manimator!", font_size=48, color=BLUE)
        subtitle = Text("Mathematical Animation Renderer", font_size=24)
        subtitle.next_to(title, DOWN)
        
        # Create geometric shapes
        circle = Circle(radius=1, color=RED)
        square = Square(side_length=2, color=GREEN)
        triangle = Triangle(color=YELLOW)
        
        # Position shapes
        circle.shift(LEFT * 3)
        square.shift(RIGHT * 3)
        triangle.shift(DOWN * 2)
        
        # Animate!
        self.play(Write(title), run_time=2)
        self.wait(0.5)
        self.play(Write(subtitle), run_time=1.5)
        self.wait(1)
        
        # Show shapes
        self.play(
            Create(circle), 
            Create(square), 
            Create(triangle),
            run_time=2
        )
        self.wait(1)
        
        # Transform shapes
        self.play(
            circle.animate.set_color(PURPLE),
            square.animate.rotate(PI/4),
            triangle.animate.scale(1.5),
            run_time=2
        )
        self.wait(1)
        
        # Group and move
        everything = Group(title, subtitle, circle, square, triangle)
        self.play(everything.animate.shift(UP * 0.5), run_time=1.5)
        self.wait(2)
'''

EXAMPLE_CODES = {
    "Simple Animation": DEFAULT_CODE,
    
    "Moving Shapes": '''from manim import *

class MovingShapes(Scene):
    def construct(self):
        # Create shapes
        circle = Circle(color=BLUE)
        square = Square(color=RED).shift(RIGHT * 2)
        
        # Animations
        self.play(Create(circle))
        self.play(circle.animate.shift(RIGHT * 2))
        self.play(Transform(circle, square))
        self.play(square.animate.scale(2).rotate(PI/4))
        self.wait()
''',
    
    "Text Animation": '''from manim import *

class TextAnimation(Scene):
    def construct(self):
        # Create text
        text1 = Text("This is Manim", font_size=72)
        text2 = Text("Create Amazing Animations!", font_size=48, color=YELLOW)
        text2.next_to(text1, DOWN)
        
        # Animate
        self.play(Write(text1))
        self.wait()
        self.play(text1.animate.scale(0.7).shift(UP))
        self.play(FadeIn(text2, shift=UP))
        self.wait(2)
''',
    
    "Graph Example": '''from manim import *

class GraphExample(Scene):
    def construct(self):
        # Create axes
        axes = Axes(
            x_range=[-3, 3, 1],
            y_range=[-3, 3, 1],
            axis_config={"color": BLUE}
        )
        
        # Create graph
        graph = axes.plot(lambda x: x**2, color=RED)
        graph_label = axes.get_graph_label(graph, label='f(x) = x^2')
        
        # Animate
        self.play(Create(axes))
        self.play(Create(graph), Write(graph_label))
        self.wait(2)
''',
    
    "Dot Movement": '''from manim import *

class DotMovement(Scene):
    def construct(self):
        # Create dots
        dot = Dot(color=RED)
        
        # Create path
        path = VMobject()
        path.set_points_as_corners([
            LEFT * 2,
            UP * 2,
            RIGHT * 2,
            DOWN * 2,
            LEFT * 2
        ])
        path.set_color(BLUE)
        
        # Animate
        self.play(Create(path))
        self.play(MoveAlongPath(dot, path), run_time=4)
        self.wait()
''',

    "Math Formulas (Text)": '''from manim import *

class MathFormulas(Scene):
    def construct(self):
        # Math using Text() - works without LaTeX!
        title = Text("Famous Formulas", font_size=48, color=BLUE)
        title.to_edge(UP)
        
        # Einstein's equation
        einstein = Text("E = mc²", font_size=60, color=YELLOW)
        
        # Pythagorean theorem
        pythagoras = Text("a² + b² = c²", font_size=60, color=GREEN)
        pythagoras.next_to(einstein, DOWN, buff=1)
        
        # Euler's identity
        euler = Text("e^(iπ) + 1 = 0", font_size=60, color=RED)
        euler.next_to(pythagoras, DOWN, buff=1)
        
        # Animate
        self.play(Write(title))
        self.wait(0.5)
        self.play(Write(einstein), run_time=2)
        self.wait(1)
        self.play(Write(pythagoras), run_time=2)
        self.wait(1)
        self.play(Write(euler), run_time=2)
        self.wait(2)
''',

    "3D Shapes": '''from manim import *

class ThreeDShapes(Scene):
    def construct(self):
        # Note: 3D requires ThreeDScene for rotation
        # This shows basic 3D objects in 2D view
        
        title = Text("3D Objects", font_size=48)
        title.to_edge(UP)
        
        # Create 3D shapes (shown in 2D projection)
        sphere = Sphere(radius=1, color=BLUE).shift(LEFT * 3)
        cube = Cube(side_length=2, color=RED)
        torus = Torus(color=GREEN).shift(RIGHT * 3).scale(0.7)
        
        # Animate
        self.play(Write(title))
        self.wait(0.5)
        self.play(Create(sphere), Create(cube), Create(torus), run_time=3)
        self.wait(2)
''',
}

def extract_scene_class(code):
    """Extract the Scene class name from code"""
    lines = code.split('\n')
    for line in lines:
        line = line.strip()
        if line.startswith('class ') and '(Scene)' in line:
            class_name = line.split('class ')[1].split('(')[0].strip()
            return class_name
    return None

def render_manim(code, quality, progress=gr.Progress()):
    """Render Manim animation from code"""
    
    if not code.strip():
        return None, "❌ Error: Please enter some code!"
    
    # Create temp directory
    temp_dir = tempfile.mkdtemp(prefix="gradio_manim_")
    
    try:
        progress(0.1, desc="Saving code...")
        
        # Save code to file
        temp_file = os.path.join(temp_dir, "animation.py")
        with open(temp_file, 'w', encoding='utf-8') as f:
            f.write(code)
        
        # Extract scene class
        scene_class = extract_scene_class(code)
        if not scene_class:
            return None, "❌ Error: No Scene class found in code!\n\nMake sure your code has a class that inherits from Scene, like:\nclass MyAnimation(Scene):"
        
        progress(0.2, desc=f"Found scene: {scene_class}")
        
        # Prepare command
        quality_map = {
            "Low (Fast)": "l",
            "Medium": "m", 
            "High": "h",
            "Production": "p"
        }
        quality_flag = quality_map.get(quality, "l")
        
        output_dir = os.path.join(temp_dir, "media")
        
        # Use Python -m manim for cross-platform compatibility
        import sys
        cmd = [
            sys.executable, "-m", "manim", "render",
            temp_file, scene_class,
            "-q", quality_flag,
            "--media_dir", output_dir,
            "--disable_caching"
        ]
        
        progress(0.3, desc="Starting render...")
        
        print(f"Running command: {' '.join(cmd)}")
        
        # Run manim
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=300,  # 5 minute timeout
            cwd=temp_dir  # Set working directory
        )
        
        print(f"Cmd run complete. Return code: {result.returncode}")
        
        progress(0.8, desc="Processing output...")
        
        # Combine stdout and stderr for output
        output_text = result.stdout + "\n" + result.stderr
        
        if result.returncode != 0:
            error_msg = f"❌ Rendering failed!\n\n{output_text}"
            return None, error_msg
        
        progress(0.9, desc="Finding video file...")
        
        # Find the generated video
        video_file = None
        videos_dir = os.path.join(output_dir, "videos", "animation")
        
        if os.path.exists(videos_dir):
            for root, dirs, files in os.walk(videos_dir):
                for file in files:
                    if file.endswith('.mp4'):
                        video_file = os.path.join(root, file)
                        break
                if video_file:
                    break
        
        if not video_file or not os.path.exists(video_file):
            # Check for static image
            images_dir = os.path.join(output_dir, "images", "animation")
            if os.path.exists(images_dir):
                for root, dirs, files in os.walk(images_dir):
                    for file in files:
                        if file.endswith(('.png', '.jpg')):
                            return None, f"⚠️ Generated static image instead of video.\n\nYour scene has no animations!\nAdd self.play(...) commands to create animated videos.\n\nOutput:\n{output_text}"
            
            return None, f"❌ Could not find output file.\n\n{output_text}"
        
        progress(1.0, desc="Done!")
        
        # Success message
        file_size = os.path.getsize(video_file) / (1024 * 1024)  # MB
        success_msg = f"✅ Animation rendered successfully!\n\n📊 Details:\n• Scene: {scene_class}\n• Quality: {quality}\n• File size: {file_size:.2f} MB"
        
        return video_file, success_msg
        
    except subprocess.TimeoutExpired:
        return None, "❌ Error: Rendering took too long (>5 minutes). Try reducing animation length or quality."
    except Exception as e:
        return None, f"❌ Error: {str(e)}"
    finally:
        # Cleanup is handled by Gradio, but we keep temp dir for video access
        pass

def load_example(example_name):
    """Load example code"""
    return EXAMPLE_CODES.get(example_name, DEFAULT_CODE)

def clear_all():
    """Clear all fields"""
    return "", None, "Ready to render!"

# Create Gradio interface
with gr.Blocks(title="Manim Animation Editor", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎬 Manim Animation Editor
    
    Create mathematical animations using Manim! Write your code, click render, and watch your animation come to life.
    
    **📚 Quick Tips:**
    - Use `self.play(...)` to create animations (required for videos!)
    - Use `self.wait()` to pause
    - Use `Text()` for formulas (works everywhere!)
    - ⚠️ `MathTex()` requires LaTeX (may not work on hosted versions)
    - Check the examples below to get started
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            # Code editor
            code_input = gr.Code(
                value=DEFAULT_CODE,
                language="python",
                label="📝 Manim Code",
                lines=25
            )
            
            with gr.Row():
                # Quality selector
                quality_input = gr.Radio(
                    choices=["Low (Fast)", "Medium", "High", "Production"],
                    value="Low (Fast)",
                    label="🎥 Quality",
                    info="Higher quality = longer render time"
                )
                
            with gr.Row():
                # Action buttons
                render_btn = gr.Button("▶️ Render Animation", variant="primary", size="lg")
                clear_btn = gr.Button("🗑️ Clear", size="lg")
            
            # Examples
            gr.Markdown("### 📚 Example Animations")
            example_dropdown = gr.Dropdown(
                choices=list(EXAMPLE_CODES.keys()),
                label="Load Example",
                value="Simple Animation"
            )
            load_example_btn = gr.Button("Load Example Code")
            
        with gr.Column(scale=2):
            # Video output
            video_output = gr.Video(
                label="🎬 Animation Output",
                height=400
            )
            
            # Status/output text
            output_text = gr.Textbox(
                label="📋 Status & Console Output",
                value="Ready to render! Write your code and click 'Render Animation'.",
                lines=15,
                max_lines=15
            )
    
    # Info section
    with gr.Accordion("ℹ️ Help & Documentation", open=False):
        gr.Markdown("""
        ### Getting Started
        
        **Basic Structure:**
        ```python
        from manim import *
        
        class MyAnimation(Scene):
            def construct(self):
                # Create objects
                text = Text("Hello!")
                
                # Animate them
                self.play(Write(text))
                self.wait()
        ```
        
        **Common Commands:**
        - `Write(obj)` - Write text or formula
        - `Create(obj)` - Draw shapes
        - `FadeIn(obj)` - Fade in object
        - `Transform(obj1, obj2)` - Transform one object to another
        - `obj.animate.move_to(...)` - Move object
        - `self.wait(seconds)` - Pause for seconds
        
        **Common Objects:**
        - `Text("...")` - Text
        - `Circle()` - Circle
        - `Square()` - Square
        - `Line(start, end)` - Line
        - `Dot()` - Dot/Point
        
        **Colors:**
        - `RED`, `BLUE`, `GREEN`, `YELLOW`, `PURPLE`, `ORANGE`, `WHITE`, `BLACK`
        
        **Positions:**
        - `UP`, `DOWN`, `LEFT`, `RIGHT`
        - `ORIGIN` - Center (0, 0, 0)
        
        ### Requirements
        - Manim Community must be installed: `pip install manim`
        - For math formulas (MathTex), LaTeX must be installed
        
        ### Troubleshooting
        - **"No Scene class found"** - Make sure you have `class YourName(Scene):`
        - **"Static image generated"** - Add `self.play(...)` commands for animations
        - **LaTeX errors** - Use `Text()` instead of `MathTex()` or install LaTeX
        """)
    
    # Event handlers
    render_btn.click(
        fn=render_manim,
        inputs=[code_input, quality_input],
        outputs=[video_output, output_text]
    )
    
    load_example_btn.click(
        fn=load_example,
        inputs=[example_dropdown],
        outputs=[code_input]
    )
    
    clear_btn.click(
        fn=clear_all,
        outputs=[code_input, video_output, output_text]
    )
    
    # Footer
    gr.Markdown("""
    ---
    💡 **Pro Tip:** Start with examples, then modify them to learn!
    
    📖 [Manim Documentation](https://docs.manim.community/) | 
    🎓 [Manim Tutorial](https://docs.manim.community/en/stable/tutorials.html)
    """)

# Launch the app
if __name__ == "__main__":
    print("🚀 Starting Manim Animation Editor...")
    print("📱 Access from mobile: Use the public URL (with share=True)")
    print("💻 Local access: http://localhost:7860")
    
    demo.launch(
        # share=False,  # Set to True for public URL
        # debug=True
    )