import gradio as gr
import subprocess
import os
import re
from pathlib import Path
import tempfile

try:
    from google import genai
    from google.genai import types
    GENAI_AVAILABLE = True
except ImportError:
    GENAI_AVAILABLE = False

GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")

# Video templates for different content types
TEMPLATES = {
    "Problem Solver": {
        "structure": "Title → Problem Statement → Step-by-Step Solution → Final Answer",
        "style": "Clean, focused, educational",
        "duration": "Focus on clear problem-solving steps"
    },
    "Concept Explainer": {
        "structure": "Title → What is it? → Visual Example → Key Points → Summary",
        "style": "Visual, intuitive, smooth transitions",
        "duration": "Build understanding gradually"
    },
    "Formula Breakdown": {
        "structure": "Title → Formula Display → What Each Part Means → Example → Application",
        "style": "Clear labeling, color-coded parts",
        "duration": "Thorough explanation of components"
    },
    "Quick Trick": {
        "structure": "Title → Normal Method → Shortcut Method → Why It Works → Practice",
        "style": "Fast-paced, attention-grabbing",
        "duration": "Efficient and memorable"
    }
}

def fix_all_known_errors(code):
    """Fix ALL known Manim errors"""
    
    # Fix colors to base 10 only
    color_map = {
        'ORANGE_A': 'ORANGE', 'ORANGE_B': 'ORANGE', 'ORANGE_C': 'ORANGE',
        'GRAY_A': 'WHITE', 'GRAY_B': 'WHITE', 'GREY': 'WHITE', 'GRAY': 'WHITE',
        'PURPLE_A': 'PURPLE', 'RED_A': 'RED', 'BLUE_A': 'BLUE', 'GREEN_A': 'GREEN',
        'BROWN': 'ORANGE', 'GOLD': 'YELLOW', 'CYAN': 'TEAL', 'MAGENTA': 'PINK',
    }
    for old, new in color_map.items():
        code = re.sub(rf'\b{old}\b', new, code)
    
    # Fix invalid objects
    code = re.sub(r'Checkmark\([^)]*\)', 'Circle(radius=0.3, color=GREEN, fill_opacity=1)', code)
    code = re.sub(r'SVGMobject\([^)]+\)', 'Circle(radius=0.8, color=PINK, fill_opacity=0.8)', code)
    code = re.sub(r'ImageMobject\([^)]+\)', 'Square(side_length=2, color=BLUE, fill_opacity=0.5)', code)
    
    # Remove problematic methods
    code = re.sub(r'\.add_prefix\([^)]+\)', '', code)
    code = re.sub(r'\.add_suffix\([^)]+\)', '', code)
    code = re.sub(r'\.get_x_axis_label\([^)]+\)', '', code)
    code = re.sub(r'\.get_y_axis_label\([^)]+\)', '', code)
    
    # Remove Table usage - too error-prone
    if 'Table(' in code:
        code = re.sub(r'(\s+)(.*Table\(.*)', r'\1# \2  # Removed', code)
        code = re.sub(r'(\s+)(.*\.get_entries\(.*)', r'\1# \2', code)
    
    # STRICT limits
    code = re.sub(r'font_size=(\d+)', lambda m: f'font_size={min(int(m.group(1)), 32)}', code)
    code = re.sub(r'\.shift\(UP\s*\*\s*\d+\.?\d*\)', '.shift(UP*1)', code)
    code = re.sub(r'\.shift\(DOWN\s*\*\s*\d+\.?\d*\)', '.shift(DOWN*1)', code)
    code = re.sub(r'\.shift\(LEFT\s*\*\s*\d+\.?\d*\)', '.shift(LEFT*2)', code)
    code = re.sub(r'\.shift\(RIGHT\s*\*\s*\d+\.?\d*\)', '.shift(RIGHT*2)', code)
    code = re.sub(r'\.to_edge\((UP|DOWN|LEFT|RIGHT)\)(?!\s*,)', r'.to_edge(\1, buff=0.8)', code)
    
    # Remove unsupported parameters
    code = re.sub(r',\s*weight=\w+', '', code)  # Remove weight parameter
    code = re.sub(r'weight=\w+,\s*', '', code)  # Remove weight at start
    
    # Add color=BLACK to Text/MathTex without color
    def add_black(match):
        text = match.group(0)
        if 'color=' not in text:
            text = text[:-1] + ', color=BLACK)'
        return text
    
    code = re.sub(r'Text\([^)]+\)', add_black, code)
    code = re.sub(r'MathTex\([^)]+\)', add_black, code)
    
    return code

def generate_code_with_gemini(prompt, template, duration, quality):
    """Generate Manim code using Gemini API"""
    
    if not GENAI_AVAILABLE or not GEMINI_API_KEY:
        return None, "❌ Setup GEMINI_API_KEY in Settings → Secrets"
    
    try:
        client = genai.Client(api_key=GEMINI_API_KEY)
        
        # Duration guidance
        duration_guide = {
            "Short (30-60s)": "Keep it concise - 3-4 key points max, quick transitions",
            "Medium (2-5 min)": "Detailed explanation - 5-8 steps, smooth pacing",
            "Long (5-10 min)": "Comprehensive - 10+ steps, thorough breakdown, multiple examples"
        }
        
        template_info = TEMPLATES[template]
        duration_info = duration_guide[duration]
        
        full_prompt = f"""You are creating a YOUTUBE MATH VIDEO for government exam preparation (SSC, Banking, Railways).

TEMPLATE: {template}
STRUCTURE: {template_info['structure']}
STYLE: {template_info['style']} (inspired by 3Blue1Brown)
TARGET DURATION: {duration} - {duration_info}

USER REQUEST: {prompt}

CRITICAL RULES:

1. YOUTUBE QUALITY:
   - Professional, clean, smooth animations
   - Clear visual hierarchy (title → content → conclusion)
   - Engaging pacing with proper timing
   - Everything centered and visible

2. STRUCTURE (Follow {template} template):
   - Start with BOLD title card (2 seconds)
   - Present problem/concept clearly
   - Step-by-step visual explanation
   - Highlight key insights/shortcuts
   - End with summary or practice prompt

3. COLORS & VISIBILITY:
   - Background: WHITE
   - ALL Text/MathTex: color=BLACK (MANDATORY!)
   - Use vibrant accent colors: BLUE (trust), GREEN (correct), RED (important), ORANGE (highlight)
   - Color-code different elements (formulas=BLUE, steps=GREEN, answers=RED)

4. PREVENT OVERLAP (CRITICAL):
   - ALWAYS FadeOut before new content
   - Pattern:
   ```python
   step1 = Text("Step 1: ...", font_size=28, color=BLACK).to_edge(UP, buff=0.8)
   self.play(Write(step1), run_time=1)
   self.wait(2)  # Let viewers read
   self.play(FadeOut(step1))
   self.wait(0.3)  # Small pause
   
   step2 = Text("Step 2: ...", font_size=28, color=BLACK).to_edge(UP, buff=0.8)
   self.play(Write(step2), run_time=1)
   self.wait(2)
   self.play(FadeOut(step2))
   ```

5. TIMING (YouTube pacing):
   - Title card: 2 seconds
   - Each explanation step: 2-3 seconds
   - Complex calculations: 3-4 seconds
   - Final answer highlight: 2 seconds
   - Total should match {duration}

6. VISUAL ELEMENTS:
   - Use boxes/rectangles to group related info
   - Arrows to show relationships
   - Underlines for emphasis
   - Fade transitions between sections
   - Scale/rotate for dynamic interest

7. STAY IN FRAME:
   - Font: MAX 32 (title), 28 (content), 24 (details)
   - Safe zone: X[-2.5, 2.5], Y[-1.5, 1.5]
   - Use .to_edge() with buff=0.8
   - Center main content with .move_to(ORIGIN)

8. NEVER use:
   - Table, DecimalTable (use VGroup instead)
   - get_x_axis_label/get_y_axis_label
   - Checkmark, SVGMobject, ImageMobject
   - Complex grids

9. QUALITY SETTINGS:
   - Use run_time for smooth animations (1-2 seconds)
   - Add fill_opacity=0.8 for shapes
   - Use stroke_width=2 for emphasis

10. CODE STRUCTURE:
```python
from manim import *

class MathExplainer(Scene):
    def construct(self):
        self.camera.background_color = WHITE
        
        # Title Card (2 sec)
        title = Text("Topic Title", font_size=32, color=BLACK)
        self.play(Write(title), run_time=1)
        self.wait(2)
        self.play(FadeOut(title))
        
        # Main Content (follow template structure)
        # ... your explanation with proper FadeOuts
        
        # Conclusion
        conclusion = Text("Key Takeaway", font_size=28, color=GREEN)
        self.play(Write(conclusion), run_time=1)
        self.wait(2)
        self.play(FadeOut(conclusion))
```

Generate professional, YouTube-ready code following the {template} template for: {prompt}"""
        
        response = client.models.generate_content(
            model='gemini-2.5-flash-lite',
            contents=full_prompt
        )
        code = response.text
        
        return code, None
        
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

def render_video_from_code(code, quality):
    """Render video from code"""
    if not code or not code.strip():
        return "❌ No code!", None
    
    try:
        temp_dir = Path(tempfile.mkdtemp(prefix="manim_"))
        code = fix_all_known_errors(code)
        
        match = re.search(r'class\s+(\w+)\s*\(Scene\)', code)
        if not match:
            return "❌ No Scene class!", None
        
        class_name = match.group(1)
        code_file = temp_dir / "animation.py"
        code_file.write_text(code)
        
        quality_map = {'Low (fast)': '-ql', 'Medium': '-qm', 'High (1080p)': '-qh'}
        quality_flag = quality_map.get(quality, '-qm')
        
        cmd = f"manim {code_file} {class_name} {quality_flag} --disable_caching --media_dir {temp_dir}/media"
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=300)
        
        if result.returncode == 0:
            videos = list((temp_dir / "media").rglob("*.mp4"))
            if videos:
                return "✅ Rendered!", str(videos[0])
        return f"❌ Failed:\n{result.stderr[-500:]}", None
    except Exception as e:
        return f"❌ {e}", None

def render_video(prompt, template, duration, quality):
    """Generate and render video with auto-retry"""
    if not prompt:
        yield "❌ Enter a prompt!", None, None
        return
    
    MAX_ATTEMPTS = 3
    
    for attempt in range(1, MAX_ATTEMPTS + 1):
        try:
            temp_dir = Path(tempfile.mkdtemp(prefix="manim_"))
            
            if attempt == 1:
                yield f"🤖 Generating {template} video ({duration})...", None, None
            else:
                yield f"🔄 Attempt {attempt}/{MAX_ATTEMPTS} - Simplifying code...", None, None
            
            # Adjust prompt complexity based on attempt
            if attempt == 1:
                current_prompt = prompt
            elif attempt == 2:
                current_prompt = f"KEEP IT VERY SIMPLE. {prompt}. Use only basic shapes and text, no complex arrangements."
            else:
                current_prompt = f"EXTREMELY SIMPLE VERSION: {prompt}. Only 2-3 simple elements, no arrays, no loops, no complex positioning."
            
            code, error = generate_code_with_gemini(current_prompt, template, duration, quality)
            
            if error:
                if attempt < MAX_ATTEMPTS:
                    continue
                yield error, None, None
                return
            
            if "```python" in code:
                code = code.split("```python")[1].split("```")[0]
            elif "```" in code:
                code = code.split("```")[1].split("```")[0]
            code = code.strip()
            
            code = fix_all_known_errors(code)
            
            # Extra safety fixes for common index errors
            # Remove array/list indexing that might fail
            code = re.sub(r'\[(\d+)\]\.', '[0].', code)  # Replace [N] with [0]
            code = re.sub(r'for\s+\w+\s+in\s+range\([^)]+\):', '# Loop removed', code)  # Comment out loops
            
            match = re.search(r'class\s+(\w+)\s*\(Scene\)', code)
            if not match:
                if attempt < MAX_ATTEMPTS:
                    continue
                yield "❌ No Scene class!", code, None
                return
            
            class_name = match.group(1)
            code_file = temp_dir / "animation.py"
            code_file.write_text(code)
            
            yield f"🎬 Rendering (attempt {attempt})...", code, None
            
            quality_map = {'Low (fast)': '-ql', 'Medium': '-qm', 'High (1080p)': '-qh'}
            quality_flag = quality_map.get(quality, '-qm')
            
            cmd = f"manim {code_file} {class_name} {quality_flag} --disable_caching --media_dir {temp_dir}/media"
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=300)
            
            if result.returncode == 0:
                videos = list((temp_dir / "media").rglob("*.mp4"))
                if videos:
                    if attempt > 1:
                        yield f"✅ Success on attempt {attempt}! (Simplified version) 🎉", code, str(videos[0])
                    else:
                        yield "✅ YouTube-ready video rendered! 🎉", code, str(videos[0])
                    return
                else:
                    if attempt < MAX_ATTEMPTS:
                        continue
                    yield "❌ Video not found", code, None
                    return
            else:
                # Rendering failed, try again with simpler code
                if attempt < MAX_ATTEMPTS:
                    error_snippet = result.stderr[-200:] if result.stderr else "Unknown error"
                    print(f"Attempt {attempt} failed: {error_snippet}")
                    continue
                else:
                    yield f"❌ All {MAX_ATTEMPTS} attempts failed. Last error:\n{result.stderr[-500:]}", code, None
                    return
                
        except Exception as e:
            if attempt < MAX_ATTEMPTS:
                print(f"Attempt {attempt} crashed: {e}")
                continue
            yield f"❌ Error after {MAX_ATTEMPTS} attempts: {e}", None, None
            return

# Gradio Interface
with gr.Blocks(title="YouTube Math Explainer") as demo:
    gr.Markdown("""
    # 🎬 YouTube Math Video Generator (Govt Exam Prep)
    
    **Professional Manim animations for your YouTube channel!**  
    Perfect for SSC, Banking, Railways exam preparation videos.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="📝 Video Topic",
                placeholder="e.g., Solve percentage problem: If 20% of 500 is increased by 40%, find the result",
                lines=4
            )
            
            template_input = gr.Dropdown(
                choices=list(TEMPLATES.keys()),
                value="Problem Solver",
                label="🎨 Video Template",
                info="Choose content structure"
            )
            
            duration_input = gr.Radio(
                choices=["Short (30-60s)", "Medium (2-5 min)", "Long (5-10 min)"],
                value="Short (30-60s)",
                label="⏱️ Target Duration"
            )
            
            quality_input = gr.Radio(
                choices=["Low (fast)", "Medium", "High (1080p)"],
                value="Medium",
                label="🎥 Quality"
            )
            
            with gr.Row():
                generate_btn = gr.Button("🚀 Generate Video", variant="primary", size="lg")
                render_btn = gr.Button("🔄 Render Edited", variant="secondary")
        
        with gr.Column(scale=1):
            status_output = gr.Textbox(label="📊 Status", lines=3)
            video_output = gr.Video(label="🎥 Your Video")
    
    code_output = gr.Code(
        label="💻 Generated Code (Editable)",
        language="python",
        lines=15,
        interactive=True
    )
    
    gr.Markdown("""
    ### 📚 Templates Explained:
    
    - **Problem Solver**: Step-by-step solution for math problems
    - **Concept Explainer**: Visual explanation of mathematical concepts
    - **Formula Breakdown**: Detailed formula explanation with examples
    - **Quick Trick**: Shortcuts and time-saving techniques
    
    ### 🎯 Tips for YouTube Success:
    - Use **Problem Solver** for numerical questions
    - Use **Quick Trick** for shortcuts (great for Shorts!)
    - Use **Concept Explainer** for theory topics
    - Start with **Short videos** to test, then go longer
    - **High (1080p)** quality for main channel, **Low** for testing
    
    ### ✨ Features:
    ✅ 3Blue1Brown inspired style  
    ✅ Professional animations  
    ✅ No text overlap  
    ✅ Everything stays in frame  
    ✅ Exam-focused content  
    ✅ Edit & re-render support
    """)
    
    generate_btn.click(
        fn=render_video,
        inputs=[prompt_input, template_input, duration_input, quality_input],
        outputs=[status_output, code_output, video_output]
    )
    
    render_btn.click(
        fn=render_video_from_code,
        inputs=[code_output, quality_input],
        outputs=[status_output, video_output]
    )
    
    gr.Examples(
        examples=[
            ["Solve: If 25% of a number is 75, find the number", "Problem Solver", "Short (30-60s)", "Medium"],
            ["Explain Pythagoras theorem with visual proof", "Concept Explainer", "Medium (2-5 min)", "Medium"],
            ["Quick trick: Calculate 15% of any number mentally", "Quick Trick", "Short (30-60s)", "Medium"],
            ["Break down the quadratic formula: x = (-b ± √(b²-4ac))/2a", "Formula Breakdown", "Medium (2-5 min)", "Medium"],
        ],
        inputs=[prompt_input, template_input, duration_input, quality_input]
    )

if __name__ == "__main__":
    demo.launch()