#!/usr/bin/env python3
"""
Theorem Explanation Agent - Hugging Face Spaces App
Generates educational videos using Gemini 2.0 Flash and Manim
"""

import os
import sys
import asyncio
import time
import random
from typing import Dict, Any, Tuple, Optional
from pathlib import Path
import gradio as gr

# Environment setup
DEMO_MODE = os.getenv("DEMO_MODE", "false").lower() == "true"
video_generator = None
CAN_IMPORT_DEPENDENCIES = True
GRADIO_OUTPUT_DIR = "gradio_outputs"

def setup_environment():
    """Setup environment for HF Spaces."""
    print("🚀 Setting up Theorem Explanation Agent...")
    
    # Create output directory
    os.makedirs(GRADIO_OUTPUT_DIR, exist_ok=True)
    
    gemini_keys = os.getenv("GEMINI_API_KEY", "")
    if gemini_keys:
        key_count = len([k.strip() for k in gemini_keys.split(',') if k.strip()])
        print(f"✅ Found {key_count} Gemini API key(s)")
        return True
    else:
        print("⚠️ No Gemini API keys found - running in demo mode")
        return False

def initialize_video_generator():
    """Initialize video generator with proper dependencies."""
    global video_generator, CAN_IMPORT_DEPENDENCIES
    
    try:
        if DEMO_MODE:
            return "⚠️ Demo mode enabled - No video generation"
        
        gemini_keys = os.getenv("GEMINI_API_KEY", "")
        if not gemini_keys:
            return "⚠️ No API keys found - Set GEMINI_API_KEY environment variable"
        
        # Import dependencies
        try:
            from generate_video import VideoGenerator
            from mllm_tools.litellm import LiteLLMWrapper
            print("✅ Successfully imported video generation dependencies")
        except ImportError as e:
            CAN_IMPORT_DEPENDENCIES = False
            print(f"❌ Import error: {e}")
            return f"⚠️ Missing dependencies: {str(e)}"
        
        # Initialize models with comma-separated API key support
        planner_model = LiteLLMWrapper(
            model_name="gemini/gemini-2.0-flash-exp",
            temperature=0.7,
            print_cost=True,
            verbose=False,
            use_langfuse=False
        )
        
        # Initialize video generator
        video_generator = VideoGenerator(
            planner_model=planner_model,  
            helper_model=planner_model,
            scene_model=planner_model,
            output_dir=GRADIO_OUTPUT_DIR,
            use_rag=False,
            use_context_learning=False,
            use_visual_fix_code=False,
            verbose=True
        )
        
        return "✅ Video generator initialized successfully"
        
    except Exception as e:
        CAN_IMPORT_DEPENDENCIES = False
        print(f"❌ Error initializing video generator: {e}")
        return f"❌ Initialization failed: {str(e)}"

def simulate_video_generation(topic: str, context: str, max_scenes: int, progress_callback=None):
    """Simulate video generation for demo mode."""
    stages = [
        ("🔍 Analyzing topic", 15),
        ("📝 Planning scenes", 30),
        ("🎬 Generating content", 50),
        ("✨ Creating animations", 75),
        ("🎥 Rendering video", 90),
        ("✅ Finalizing", 100)
    ]
    
    results = []
    for stage, progress in stages:
        if progress_callback:
            progress_callback(progress, stage)
        time.sleep(random.uniform(0.5, 1.0))
        results.append(f"• {stage}")
    
    return {
        "success": True,
        "message": f"Demo simulation completed for: {topic}",
        "scenes_created": max_scenes,
        "processing_steps": results,
        "demo_note": "This is a simulation - set GEMINI_API_KEY and DEMO_MODE=false for real generation"
    }

async def generate_video_async(topic: str, context: str, max_scenes: int, progress_callback=None):
    """Generate video asynchronously using the actual VideoGenerator."""
    global video_generator
    
    if not topic.strip():
        return {"success": False, "error": "Please enter a topic"}
    
    try:
        if DEMO_MODE or not CAN_IMPORT_DEPENDENCIES or video_generator is None:
            return simulate_video_generation(topic, context, max_scenes, progress_callback)
        
        if progress_callback:
            progress_callback(10, "🚀 Starting video generation...")
        
        # Use the actual video generation pipeline
        result = await video_generator.generate_video_pipeline(
            topic=topic,
            description=context or f"Educational video about {topic}",
            max_retries=3,
            only_plan=False,
            specific_scenes=list(range(1, max_scenes + 1)) if max_scenes > 0 else None
        )
        
        if progress_callback:
            progress_callback(100, "✅ Video generation completed!")
        
        # Check for generated video files
        file_prefix = topic.lower().replace(' ', '_')
        file_prefix = ''.join(c for c in file_prefix if c.isalnum() or c == '_')
        
        output_folder = os.path.join(GRADIO_OUTPUT_DIR, file_prefix)
        video_files = []
        
        if os.path.exists(output_folder):
            # Look for combined video
            combined_video = os.path.join(output_folder, f"{file_prefix}_combined.mp4")
            if os.path.exists(combined_video):
                video_files.append(combined_video)
            
            # Look for individual scene videos
            for i in range(1, max_scenes + 1):
                scene_video = os.path.join(output_folder, f"scene{i}", f"{file_prefix}_scene{i}.mp4")
                if os.path.exists(scene_video):
                    video_files.append(scene_video)
        
        return {
            "success": True, 
            "message": f"Video generated successfully for: {topic}",
            "video_files": video_files,
            "output_folder": output_folder,
            "result": result
        }
        
    except Exception as e:
        print(f"❌ Error in video generation: {e}")
        return {"success": False, "error": str(e)}

def generate_video_gradio(topic: str, context: str, max_scenes: int, progress=gr.Progress()) -> Tuple[str, str, Optional[str]]:
    """Main Gradio function that handles video generation and returns results."""
    def progress_callback(percent, message):
        progress(percent / 100, desc=message)
    
    # Create new event loop for this generation
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    
    try:
        result = loop.run_until_complete(
            generate_video_async(topic, context, max_scenes, progress_callback)
        )
    finally:
        loop.close()
    
    if result["success"]:
        output = f"""# 🎓 Video Generation Complete!

**Topic:** {topic}
**Context:** {context if context else "None"}
**Scenes:** {max_scenes}

## ✅ Result
{result["message"]}
"""
        
        # Add processing steps if available
        if "processing_steps" in result:
            output += "\n## 🔄 Processing Steps\n"
            for step in result["processing_steps"]:
                output += f"{step}\n"
        
        # Add demo note if in demo mode
        if "demo_note" in result:
            output += f"\n⚠️ **{result['demo_note']}**"
        
        # Add video file information
        video_path = None
        if "video_files" in result and result["video_files"]:
            output += f"\n## 🎥 Generated Videos\n"
            for video_file in result["video_files"]:
                output += f"• {os.path.basename(video_file)}\n"
            video_path = result["video_files"][0]  # Return first video for display
        elif "output_folder" in result:
            output += f"\n📁 **Output folder:** {result['output_folder']}\n"
        
        status = "🎮 Demo completed" if DEMO_MODE else "✅ Generation completed"
        return output, status, video_path
    
    else:
        error_output = f"""# ❌ Video Generation Failed

**Error:** {result.get("error", "Unknown error")}

## 💡 Troubleshooting Tips
1. **Check API Keys:** Ensure GEMINI_API_KEY is set with valid keys
2. **Topic Clarity:** Use specific, educational topics
3. **Dependencies:** Make sure all required packages are installed
4. **Demo Mode:** Set DEMO_MODE=false for real generation

## 🔧 Environment Setup
```bash
export GEMINI_API_KEY="your-key-1,your-key-2,your-key-3"
export DEMO_MODE=false
```
"""
        return error_output, "❌ Generation failed", None

def get_examples():
    """Educational example topics."""
    return [
        ["Pythagorean Theorem", "Mathematical proof with geometric visualization"],
        ["Newton's Second Law", "F=ma with real-world examples and demonstrations"],
        ["Derivatives in Calculus", "Rate of change with graphical interpretation"],
        ["Photosynthesis Process", "Cellular process with chemical equations"],
        ["Wave-Particle Duality", "Quantum physics concept with experiments"],
        ["Quadratic Formula", "Step-by-step derivation and applications"],
        ["DNA Replication", "Biological process with molecular details"],
        ["Ohm's Law", "Electrical relationship with circuit examples"]
    ]

# Initialize the system
has_api_keys = setup_environment()
init_status = initialize_video_generator()

# Create Gradio interface
with gr.Blocks(
    title="🎓 Theorem Explanation Agent",
    theme=gr.themes.Soft(),
    css="footer {visibility: hidden}"
) as demo:
    
    gr.HTML("""
    <div style="text-align: center; padding: 25px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 15px; color: white; margin-bottom: 25px; box-shadow: 0 4px 6px rgba(0,0,0,0.1);">
        <h1 style="margin: 0; font-size: 2.5em;">🎓 Theorem Explanation Agent</h1>
        <p style="margin: 10px 0 0 0; font-size: 1.2em; opacity: 0.9;">Generate Educational Videos with AI</p>
        <p style="margin: 5px 0 0 0; font-size: 0.9em; opacity: 0.8;">Powered by Gemini 2.0 Flash & Manim</p>
    </div>
    """)
    
    # Status and setup information
    with gr.Row():
        with gr.Column():
            gr.HTML(f"""
            <div style="background: {'#d4edda' if has_api_keys else '#fff3cd'}; padding: 15px; border-radius: 10px; margin-bottom: 15px; border-left: 4px solid {'#28a745' if has_api_keys else '#ffc107'};">
                <h4 style="margin: 0 0 8px 0;">🔐 API Setup Status</h4>
                <p style="margin: 0;"><strong>Status:</strong> {"✅ API keys configured" if has_api_keys else "⚠️ No API keys found"}</p>
                <p style="margin: 5px 0 0 0; font-size: 0.9em;">{"Ready for video generation" if has_api_keys else "Running in demo mode"}</p>
            </div>
            """)
            
        with gr.Column():
            system_status = gr.Textbox(
                label="🔧 System Status",
                value=init_status,
                interactive=False,
                lines=2
            )
    
    # API Configuration Help
    if not has_api_keys or DEMO_MODE:
        gr.HTML("""
        <div style="background: #f8f9fa; padding: 20px; border-radius: 10px; margin: 15px 0; border: 1px solid #dee2e6;">
            <h4 style="color: #495057; margin-top: 0;">🚀 Enable Full Functionality</h4>
            <p style="margin-bottom: 15px;">To generate actual videos instead of simulations:</p>
            <div style="background: #e9ecef; padding: 15px; border-radius: 5px; font-family: monospace;">
                <strong>Single API Key:</strong><br>
                <code>GEMINI_API_KEY=your-gemini-api-key</code><br><br>
                <strong>Multiple Keys (Recommended):</strong><br>
                <code>GEMINI_API_KEY=key1,key2,key3,key4</code><br><br>
                <strong>Disable Demo Mode:</strong><br>
                <code>DEMO_MODE=false</code>
            </div>
            <p style="margin-top: 15px; font-size: 0.9em; color: #6c757d;">
                Multiple API keys enable automatic failover and load distribution across different billing accounts.
            </p>
        </div>
        """)
    
    # Main interface
    with gr.Row():
        with gr.Column(scale=2):
            topic_input = gr.Textbox(
                label="📚 Educational Topic",
                placeholder="e.g., Pythagorean Theorem, Newton's Laws, Derivatives...",
                lines=1
            )
            
            context_input = gr.Textbox(
                label="📝 Additional Context (Optional)",
                placeholder="Specify focus areas, target audience, or particular aspects to emphasize...",
                lines=3
            )
            
            max_scenes_slider = gr.Slider(
                label="🎬 Maximum Scenes",
                minimum=1,
                maximum=6,
                value=3,
                step=1,
                info="More scenes = longer videos but more API usage"
            )
            
            generate_btn = gr.Button("🚀 Generate Educational Video", variant="primary", size="lg")
            
        with gr.Column(scale=1):
            gr.HTML("""
            <div style="background: #f8f9fa; padding: 20px; border-radius: 10px; height: fit-content;">
                <h4 style="color: #495057; margin-top: 0;">💡 Tips for Best Results</h4>
                <ul style="color: #6c757d; font-size: 0.9em; line-height: 1.6;">
                    <li><strong>Be Specific:</strong> "Pythagorean Theorem proof" vs "Math"</li>
                    <li><strong>Educational Focus:</strong> Topics work best for teaching</li>
                    <li><strong>Context Helps:</strong> Specify audience or emphasis</li>
                    <li><strong>Start Small:</strong> Try 2-3 scenes first</li>
                </ul>
            </div>
            """)
    
    # Examples
    examples = gr.Examples(
        examples=get_examples(),
        inputs=[topic_input, context_input],
        label="📖 Example Topics"
    )
    
    # Output section
    with gr.Row():
        with gr.Column(scale=2):
            output_display = gr.Markdown(
                value="👋 **Ready to generate!** Enter an educational topic above and click 'Generate Educational Video' to begin.",
                label="📋 Generation Results"
            )
        
        with gr.Column(scale=1):
            video_output = gr.Video(
                label="🎥 Generated Video",
                visible=True
            )
    
    # Wire up the interface
    generate_btn.click(
        fn=generate_video_gradio,
        inputs=[topic_input, context_input, max_scenes_slider],
        outputs=[output_display, system_status, video_output],
        show_progress=True
    )

# Launch configuration
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )