import gradio as gr
import os
import subprocess
import time
import shutil
import glob
import re

# ---------------------------------------------------------
# 1. Helper Functions
# ---------------------------------------------------------

def cleanup_media_directory():
    """Removes the old media directory to prevent clutter."""
    media_dir = 'media'
    if os.path.exists(media_dir):
        try:
            shutil.rmtree(media_dir)
        except OSError:
            pass

def make_even(n):
    return int(n) if int(n) % 2 == 0 else int(n) + 1

def get_resolution_flags(orientation, quality):
    """Calculates exact width and height based on orientation and quality."""
    
    if orientation == "Carousel Portrait (4:5)":
        res_map = {
            "Preview (360p)": (288, 360),
            "480p": (384, 480),
            "720p": (576, 720),
            "1080p": (1080, 1350),  # Exactly 1080x1350 for 4:5 carousels
            "4k": (2160, 2700)
        }
    elif orientation == "Portrait (9:16)":
        res_map = {
            "Preview (360p)": (360, 640),
            "480p": (480, 854),
            "720p": (720, 1280),
            "1080p": (1080, 1920),
            "4k": (2160, 3840)
        }
    else: # Landscape (16:9)
        res_map = {
            "Preview (360p)": (640, 360),
            "480p": (854, 480),
            "720p": (1280, 720),
            "1080p": (1920, 1080),
            "4k": (3840, 2160)
        }
        
    w, h = res_map.get(quality, (1080, 1350))
    return f"{make_even(w)},{make_even(h)}"

def get_scene_names_in_order(code_str):
    """Parses the code to find the order of Scene classes defined."""
    pattern = r"class\s+([A-Za-z0-9_]+)\s*\([^)]*\):"
    return re.findall(pattern, code_str)

def run_manim(code_str, orientation, quality):
    """
    Executes Manim to render all Scene classes as separate PNG images.
    """
    print(f"🎨 Starting Carousel Render: {orientation} @ {quality}...", flush=True)
    
    with open("scene.py", "w", encoding="utf-8") as f: 
        f.write(code_str)
    
    res_str = get_resolution_flags(orientation, quality)
    
    # -a renders all scenes in the file
    # -s instructs Manim to save the last frame of each scene as an image
    cmd = [
        "manim", 
        "--resolution", res_str, 
        "--disable_caching", 
        "--progress_bar", "none", 
        "scene.py", 
        "-a", "-s" 
    ]
    
    full_logs = ""

    try:
        process = subprocess.run(cmd, capture_output=True, timeout=60, check=False)
        stdout_log = process.stdout.decode('utf-8', 'ignore')
        stderr_log = process.stderr.decode('utf-8', 'ignore')
        full_logs = f"--- MANIM STDOUT ---\n{stdout_log}\n\n--- MANIM STDERR ---\n{stderr_log}"
        
        if process.returncode != 0:
            print(f"❌ Render Failed (Process Error). Return Code: {process.returncode}", flush=True)

    except subprocess.TimeoutExpired:
        print("⌛ Render timed out.", flush=True)
        return [], "❌ Failure: Render Timed Out.", False

    # Manim saves images into the media/images/scene folder
    media_image_base = os.path.join("media", "images", "scene")
    generated_images = []
    
    if os.path.exists(media_image_base):
        # Extract scene names in the exact order they were written in the code
        expected_scenes = get_scene_names_in_order(code_str)
        
        for scene_name in expected_scenes:
            expected_img_path = os.path.join(media_image_base, f"{scene_name}.png")
            if os.path.exists(expected_img_path):
                generated_images.append(expected_img_path)

        # Fallback: if regex missed anything, grab whatever is left in the folder
        all_pngs = glob.glob(os.path.join(media_image_base, "*.png"))
        for png in all_pngs:
            if png not in generated_images:
                generated_images.append(png)
                
        if generated_images:
            print(f"✅ Generated {len(generated_images)} images.", flush=True)
            return generated_images, f"✅ Rendering Successful. Generated {len(generated_images)} slides.\n\n{full_logs}", True
    
    return [], f"❌ Failure: No images were created.\n\n{full_logs}", False

# ---------------------------------------------------------
# 2. Main API Function
# ---------------------------------------------------------

def render_carousel_from_code(code, orientation, quality):
    try:
        cleanup_media_directory()
        
        if not code or "from manim import" not in code:
            return [], "Error: No valid Manim code provided."
        
        images, logs, success = run_manim(code, orientation, quality)
        return images, logs
    except Exception as e:
        return [], f"Rendering failed: {str(e)}"

# ---------------------------------------------------------
# 3. Gradio Interface
# ---------------------------------------------------------

# Carousel Example Code specifically tailored to look good on 4:5 1080x1350
DEFAULT_CODE = """from manim import *

class Slide1_Title(Scene):
    def construct(self):
        title = Text("Manim Carousel", font_size=64, color=YELLOW, weight=BOLD)
        subtitle = Text("Swipe to see more ->", font_size=32).next_to(title, DOWN, buff=0.5)
        
        # Self.add puts it on the screen immediately for the static PNG export
        self.add(title, subtitle)

class Slide2_Shapes(Scene):
    def construct(self):
        circle = Circle(color=BLUE, fill_opacity=0.5, radius=2)
        text = Text("High Quality Shapes", font_size=48).next_to(circle, DOWN, buff=1)
        self.add(circle, text)

class Slide3_Conclusion(Scene):
    def construct(self):
        square = Square(color=RED, fill_opacity=0.5, side_length=4)
        text = Text("Perfect 1080x1350 Output", font_size=48).next_to(square, DOWN, buff=1)
        self.add(square, text)
"""

with gr.Blocks(title="Manim Carousel Generator") as demo:
    gr.Markdown("## 🎠 Manim Carousel Image Generator")
    gr.Markdown("Define **multiple `Scene` classes** in your code. Manim will render the final frame of *each* class and display them sequentially. Select `Carousel Portrait (4:5)` and `1080p` for exact **1080x1350** dimensions.")
    
    with gr.Row():
        with gr.Column(scale=1):
            code_input = gr.Code(label="Python Code", language="python", value=DEFAULT_CODE) 
            orientation_opt = gr.Radio(
                choices=["Landscape (16:9)", "Portrait (9:16)", "Carousel Portrait (4:5)"], 
                value="Carousel Portrait (4:5)", # Defaulted to Carousel 4:5
                label="Orientation"
            )
            quality_opt = gr.Dropdown(
                choices=["Preview (360p)", "480p", "720p", "1080p", "4k"], 
                value="1080p", # Defaulted to 1080p for 1080x1350 output
                label="Quality"
            )
            render_btn = gr.Button("Generate Carousel", variant="primary")
            
        with gr.Column(scale=1):
            gallery_output = gr.Gallery(label="Carousel Output", columns=2, object_fit="contain", height="auto")
            status_output = gr.Textbox(label="Status/Logs", lines=10)
    
    render_btn.click(
        fn=render_carousel_from_code,
        inputs=[code_input, orientation_opt, quality_opt],
        outputs=[gallery_output, status_output],
        api_name="render_carousel"
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)