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

# Check if Manim is available
try:
    import manim
    MANIM_AVAILABLE = True
except ImportError:
    MANIM_AVAILABLE = False

class ManimAnimationGenerator:
    def __init__(self):
        self.temp_dir = None
        self.output_dir = None
    
    def setup_directories(self):
        """Setup temporary directories for Manim execution"""
        self.temp_dir = tempfile.mkdtemp()
        self.output_dir = os.path.join(self.temp_dir, "media", "videos", "480p15")
        os.makedirs(self.output_dir, exist_ok=True)
        return self.temp_dir
    
    def cleanup_directories(self):
        """Clean up temporary directories"""
        if self.temp_dir and os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)
            self.temp_dir = None
            self.output_dir = None
    
    def validate_manim_code(self, code):
        """Basic validation of Manim code"""
        required_imports = ["from manim import *", "import manim"]
        has_import = any(imp in code for imp in required_imports)
        
        if not has_import:
            return False, "Code must include 'from manim import *' or 'import manim'"
        
        if "class" not in code:
            return False, "Code must contain at least one class definition"
        
        if "Scene" not in code:
            return False, "Class must inherit from Scene or a Scene subclass"
        
        return True, "Code validation passed"
    
    def check_latex_installation(self):
        """Check if LaTeX is installed"""
        try:
            result = subprocess.run(
                ["latex", "--version"], 
                capture_output=True, 
                text=True, 
                timeout=10
            )
            return result.returncode == 0, result.stdout if result.returncode == 0 else result.stderr
        except FileNotFoundError:
            return False, "LaTeX not found"
        except Exception as e:
            return False, str(e)
    
    def install_system_dependencies(self):
        """Install required system dependencies including LaTeX"""
        try:
            # Update package list
            subprocess.check_call([
                "apt-get", "update"
            ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            
            # Install LaTeX and other dependencies
            subprocess.check_call([
                "apt-get", "install", "-y",
                "libcairo2-dev", "libpango1.0-dev", "libgirepository1.0-dev",
                "libxml2-dev", "libxslt1-dev", "libffi-dev", "libssl-dev",
                "pkg-config", "build-essential", "ffmpeg",
                "texlive-latex-base", "texlive-latex-extra", 
                "texlive-fonts-recommended", "texlive-fonts-extra",
                "texlive-xetex", "texlive-luatex", "dvipng", "dvisvgm"
            ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            
            return True, "System dependencies installed"
        except subprocess.CalledProcessError as e:
            return False, f"Failed to install system dependencies: {str(e)}"
        except Exception as e:
            return False, f"Error installing dependencies: {str(e)}"
    
    def install_manim(self):
        """Try to install Manim if not available"""
        try:
            # Try to install Manim first
            subprocess.check_call([
                sys.executable, "-m", "pip", "install", 
                "manim", "--quiet"
            ])
            
            # Try to install manimpango separately if needed
            try:
                subprocess.check_call([
                    sys.executable, "-m", "pip", "install", 
                    "manimpango", "--quiet"
                ])
            except subprocess.CalledProcessError:
                # manimpango might fail, but manim might still work
                pass
                
            return True, "Manim installed successfully"
        except subprocess.CalledProcessError as e:
            return False, f"Failed to install Manim: {str(e)}"
    
    def execute_manim_code(self, code, quality="low", format_type="gif"):
        """Execute Manim code and return the generated animation"""
        global MANIM_AVAILABLE
        
        # Check if Manim is available, try to install if not
        if not MANIM_AVAILABLE:
            success, message = self.install_manim()
            if not success:
                return None, f"❌ Manim Installation Error: {message}", ""
            MANIM_AVAILABLE = True
        
        # Check if code uses MathTex and LaTeX is available
        if "MathTex" in code or "Tex(" in code:
            latex_ok, latex_msg = self.check_latex_installation()
            if not latex_ok:
                return None, f"❌ LaTeX Error: LaTeX is required for MathTex but not installed. {latex_msg}", ""
        
        try:
            # Validate code
            is_valid, message = self.validate_manim_code(code)
            if not is_valid:
                return None, f"❌ Validation Error: {message}", ""
            
            # Setup directories
            temp_dir = self.setup_directories()
            
            # Create Python file
            python_file = os.path.join(temp_dir, "animation.py")
            with open(python_file, "w") as f:
                f.write(code)
            
            # Extract class name for Manim command
            class_name = self.extract_class_name(code)
            if not class_name:
                self.cleanup_directories()
                return None, "❌ Error: Could not find a valid Scene class in the code", ""
            
            # Quality settings
            quality_map = {
                "low": "-ql",
                "medium": "-qm", 
                "high": "-qh"
            }
            quality_flag = quality_map.get(quality, "-ql")
            
            # Format settings
            format_flag = "--format=gif" if format_type == "gif" else ""
            
            # Build Manim command
            cmd = [
                sys.executable, "-m", "manim",
                quality_flag,
                python_file,
                class_name
            ]
            
            if format_flag:
                cmd.append(format_flag)
            
            # Set environment variables to help with rendering
            env = os.environ.copy()
            env['MPLBACKEND'] = 'Agg'  # Use non-interactive backend
            
            # Execute Manim
            result = subprocess.run(
                cmd,
                cwd=temp_dir,
                capture_output=True,
                text=True,
                timeout=120,  # 2 minute timeout
                env=env
            )
            
            if result.returncode != 0:
                error_msg = f"❌ Manim execution failed:\n{result.stderr}"
                self.cleanup_directories()
                return None, error_msg, result.stdout
            
            # Find generated file
            output_file = self.find_output_file(temp_dir, class_name, format_type)
            if not output_file:
                self.cleanup_directories()
                return None, "❌ Error: Could not find generated animation file", result.stdout
            
            # Copy to a permanent location for Gradio
            permanent_file = f"/tmp/{class_name}_{int(time.time())}.{format_type}"
            shutil.copy2(output_file, permanent_file)
            
            success_msg = f"✅ Animation generated successfully!\nClass: {class_name}\nQuality: {quality}\nFormat: {format_type}"
            
            self.cleanup_directories()
            return permanent_file, success_msg, result.stdout
            
        except subprocess.TimeoutExpired:
            self.cleanup_directories()
            return None, "❌ Error: Animation generation timed out (2 minutes)", ""
        except Exception as e:
            self.cleanup_directories()
            return None, f"❌ Error: {str(e)}", ""
    
    def extract_class_name(self, code):
        """Extract the first Scene class name from the code"""
        lines = code.split('\n')
        for line in lines:
            if line.strip().startswith('class ') and 'Scene' in line:
                # Extract class name
                class_def = line.strip().split('class ')[1].split('(')[0].strip()
                return class_def
        return None
    
    def find_output_file(self, temp_dir, class_name, format_type):
        """Find the generated output file"""
        # Search recursively for the output file
        for root, dirs, files in os.walk(temp_dir):
            for file in files:
                if file.startswith(class_name) and file.endswith(f".{format_type}"):
                    return os.path.join(root, file)
        
        return None

# Initialize the generator
generator = ManimAnimationGenerator()

# Example Manim codes for users
example_codes = {
    "Simple Square": '''from manim import *

class CreateSquare(Scene):
    def construct(self):
        square = Square()
        square.set_fill(BLUE, opacity=0.5)
        square.set_stroke(WHITE, width=2)
        
        self.play(Create(square))
        self.play(square.animate.rotate(PI/4))
        self.wait()''',
    
    "Moving Circle": '''from manim import *

class MovingCircle(Scene):
    def construct(self):
        circle = Circle()
        circle.set_fill(RED, opacity=0.5)
        
        self.play(Create(circle))
        self.play(circle.animate.shift(RIGHT * 2))
        self.play(circle.animate.shift(UP * 2))
        self.play(circle.animate.shift(LEFT * 2))
        self.play(circle.animate.shift(DOWN * 2))
        self.wait()''',
    
    "Text Animation": '''from manim import *

class TextAnimation(Scene):
    def construct(self):
        text = Text("Hello Manim!", font_size=48)
        text2 = Text("Animated Math!", font_size=48)
        
        self.play(Write(text))
        self.wait()
        self.play(Transform(text, text2))
        self.wait()''',
    
    "Mathematical Formula": '''from manim import *

class MathFormula(Scene):
    def construct(self):
        formula = MathTex(r"\\\\frac{d}{dx}(x^2) = 2x")
        formula.scale(2)
        
        self.play(Write(formula))
        self.wait()
        
        formula2 = MathTex(r"\\\\int_0^1 x^2 dx = \\\\frac{1}{3}")
        formula2.scale(2)
        
        self.play(Transform(formula, formula2))
        self.wait()''',
    
    "Simple Math Text": '''from manim import *

class SimpleMath(Scene):
    def construct(self):
        # Alternative to MathTex that doesn't require LaTeX
        text1 = Text("f(x) = x²", font_size=48)
        text2 = Text("f'(x) = 2x", font_size=48)
        
        self.play(Write(text1))
        self.wait()
        self.play(Transform(text1, text2))
        self.wait()''',
    
    "Polynomial Plot": '''from manim import *
import numpy as np

class PolynomialPlot(Scene):
    def construct(self):
        # Set up the coordinate system
        axes = Axes(
            x_range=[-4, 4, 1],
            y_range=[-20, 80, 10],
            x_length=10,
            y_length=6,
            axis_config={"color": BLUE},
            x_axis_config={
                "numbers_to_include": np.arange(-4, 5, 1),
                "numbers_with_elongated_ticks": np.arange(-4, 5, 1),
            },
            y_axis_config={
                "numbers_to_include": np.arange(-20, 81, 20),
                "numbers_with_elongated_ticks": np.arange(-20, 81, 20),
            },
            tips=False,
        )
        
        # Add axis labels
        axes_labels = axes.get_axis_labels(x_label="x", y_label="f(x)")
        
        # Define the polynomial function
        def polynomial(x):
            return x**3 - 2*x**2 + 10*x + 8
        
        # Create the polynomial graph
        graph = axes.plot(polynomial, color=RED, x_range=[-4, 4])
        
        # Create the function label
        function_label = MathTex(r"f(x) = x^3 - 2x^2 + 10x + 8", color=RED)
        function_label.to_corner(UP + LEFT)
        
        # Animation sequence
        # First, show the axes
        self.play(Create(axes), Write(axes_labels), run_time=2)
        
        # Show the function label
        self.play(Write(function_label), run_time=1)
        
        # Animate the drawing of the polynomial curve
        self.play(Create(graph), run_time=6)
        
        # Hold the final scene
        self.wait(1)
        
        # Optional: Add some points of interest
        # Find and mark the y-intercept (when x=0)
        y_intercept_point = axes.coords_to_point(0, polynomial(0))
        y_intercept_dot = Dot(y_intercept_point, color=YELLOW, radius=0.08)
        y_intercept_label = MathTex(r"(0, 8)", color=YELLOW).next_to(y_intercept_dot, RIGHT)
        
        self.play(
            Create(y_intercept_dot),
            Write(y_intercept_label),
            run_time=0.5
        )
        
        # Final wait to complete 10 seconds total
        self.wait(0.5)'''
}

def check_system_requirements():
    """Check system requirements including LaTeX"""
    status = {}
    
    # Check Manim
    try:
        result = subprocess.run([
            sys.executable, "-c", "import manim; print('Manim version:', manim.__version__)"
        ], capture_output=True, text=True, timeout=10)
        
        if result.returncode == 0:
            status['manim'] = f"✅ {result.stdout.strip()}"
        else:
            status['manim'] = f"❌ Manim import failed: {result.stderr}"
    except Exception as e:
        status['manim'] = f"❌ Error checking Manim: {str(e)}"
    
    # Check LaTeX
    try:
        result = subprocess.run(
            ["latex", "--version"], 
            capture_output=True, 
            text=True, 
            timeout=10
        )
        if result.returncode == 0:
            version_line = result.stdout.split('\n')[0]
            status['latex'] = f"✅ LaTeX: {version_line}"
        else:
            status['latex'] = "❌ LaTeX not working properly"
    except FileNotFoundError:
        status['latex'] = "❌ LaTeX not installed (required for MathTex)"
    except Exception as e:
        status['latex'] = f"❌ LaTeX check error: {str(e)}"
    
    # Check FFmpeg
    try:
        result = subprocess.run(
            ["ffmpeg", "-version"], 
            capture_output=True, 
            text=True, 
            timeout=10
        )
        if result.returncode == 0:
            version_line = result.stdout.split('\n')[0]
            status['ffmpeg'] = f"✅ FFmpeg: {version_line.split(' ')[2]}"
        else:
            status['ffmpeg'] = "❌ FFmpeg not working"
    except FileNotFoundError:
        status['ffmpeg'] = "❌ FFmpeg not installed"
    except Exception as e:
        status['ffmpeg'] = f"❌ FFmpeg check error: {str(e)}"
    
    return status

def generate_animation(code, quality, format_type, progress=gr.Progress()):
    """Main function to generate animation"""
    progress(0, desc="Starting animation generation...")
    
    if not code.strip():
        return None, "❌ Please enter some Manim code", ""
    
    progress(0.2, desc="Checking system requirements...")
    
    progress(0.3, desc="Validating code...")
    result = generator.execute_manim_code(code, quality, format_type)
    
    progress(0.7, desc="Generating animation...")
    
    if result[0]:  # Success
        progress(1.0, desc="Animation generated successfully!")
        return result[0], result[1], result[2]
    else:  # Error
        return None, result[1], result[2]

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

# Create Gradio interface with custom CSS for fixed height code input
with gr.Blocks(
    title="Manim Animation Generator", 
    theme=gr.themes.Soft(),
    css="""
    /* Fix the height of the code input and add scrollbar */
    .code-input textarea {
        height: 400px !important;
        max-height: 400px !important;
        min-height: 400px !important;
        overflow-y: auto !important;
        resize: none !important;
    }
    
    /* Ensure the parent container doesn't expand */
    .code-input {
        height: 400px !important;
        max-height: 400px !important;
    }
    
    /* Style the scrollbar for better visibility */
    .code-input textarea::-webkit-scrollbar {
        width: 8px;
    }
    
    .code-input textarea::-webkit-scrollbar-track {
        background: #f1f1f1;
        border-radius: 4px;
    }
    
    .code-input textarea::-webkit-scrollbar-thumb {
        background: #888;
        border-radius: 4px;
    }
    
    .code-input textarea::-webkit-scrollbar-thumb:hover {
        background: #555;
    }
    """
) as app:
    gr.Markdown("""
    # 🎬 Manim Animation Generator
    
    Create beautiful mathematical animations using [Manim](https://www.manim.community/)!
    Enter your Python code below and watch it come to life.
    
    ## 📝 How to use:
    1. Write or select example Manim code
    2. Choose quality and format settings
    3. Click "Generate Animation"
    4. Wait for your animation to render
    
    ## 💡 Tips:
    - Your code must include `from manim import *`
    - Create a class that inherits from `Scene`
    - Use the `construct` method to define your animation
    """)
    
    # System status
    system_status = check_system_requirements()
    status_text = "\n".join([f"**{k.upper()}:** {v}" for k, v in system_status.items()])
    gr.Markdown(f"**System Status:**\n{status_text}")
    
    with gr.Row():
        with gr.Column(scale=2):
            code_input = gr.Code(
                label="Manim Code",
                language="python",
                lines=20,
                value=example_codes["Simple Square"],
                elem_classes=["code-input"]  # Add custom CSS class
            )
            
            with gr.Row():
                quality = gr.Dropdown(
                    choices=["low", "medium", "high"],
                    value="low",
                    label="Quality"
                )
                format_type = gr.Dropdown(
                    choices=["gif", "mp4"],
                    value="gif",
                    label="Format"
                )
            
            generate_btn = gr.Button("🎬 Generate Animation", variant="primary", size="lg")
            
            gr.Markdown("### 📚 Example Codes:")
            example_dropdown = gr.Dropdown(
                choices=list(example_codes.keys()),
                label="Load Example",
                value="Simple Square"
            )
            load_example_btn = gr.Button("📂 Load Example")
        
        with gr.Column(scale=2):
            output_video = gr.File(label="Generated Animation", file_types=[".gif", ".mp4"])
            status_output = gr.Textbox(label="Status", lines=5, max_lines=10)
            logs_output = gr.Textbox(label="Manim Logs", lines=8, max_lines=15, visible=False)
            
            with gr.Row():
                show_logs_btn = gr.Button("Show Logs", size="sm")
                hide_logs_btn = gr.Button("Hide Logs", size="sm")
    
    # Event handlers
    generate_btn.click(
        fn=generate_animation,
        inputs=[code_input, quality, format_type],
        outputs=[output_video, status_output, logs_output]
    )
    
    load_example_btn.click(
        fn=load_example,
        inputs=[example_dropdown],
        outputs=[code_input]
    )
    
    example_dropdown.change(
        fn=load_example,
        inputs=[example_dropdown],
        outputs=[code_input]
    )
    
    show_logs_btn.click(
        fn=lambda: gr.update(visible=True),
        outputs=[logs_output]
    )
    
    hide_logs_btn.click(
        fn=lambda: gr.update(visible=False),
        outputs=[logs_output]
    )
    
    gr.Markdown("""
    ## 🔧 Troubleshooting:
    - **Timeout errors**: Try simpler animations or lower quality
    - **Import errors**: Make sure to include `from manim import *`
    - **Class errors**: Your class must inherit from `Scene`
    - **No output**: Check that your `construct` method has animation commands
    
    ## 🎯 Common Manim Objects:
    - **Shapes**: `Circle()`, `Square()`, `Triangle()`, `Rectangle()`
    - **Text**: `Text("Hello")`, `MathTex(r"x^2")` (requires LaTeX)
    - **Animations**: `Create()`, `Write()`, `Transform()`, `FadeIn()`, `FadeOut()`
    - **Colors**: `RED`, `BLUE`, `GREEN`, `YELLOW`, `WHITE`, `BLACK`
    
    ---
    Made with ❤️ using [Manim](https://www.manim.community/) and [Gradio](https://gradio.app/)
    """)

if __name__ == "__main__":
    app.launch(debug=True, share=True)