main.py

"""
Main runner for the Coding Visualizer
Provides easy interface to run different problem visualizations
Supports both manual problems and dynamic LeetCode generation
"""

import subprocess
import sys
import os
from pathlib import Path
import tempfile

# Add src to path for imports
sys.path.append(str(Path(__file__).parent / "src"))

class CodingVisualizerRunner:
    """Main runner for coding problem visualizations"""
    
    def __init__(self):
        self.base_dir = Path(__file__).parent
        self.problems_dir = self.base_dir / "src" / "problems"
        self.leetcode_repo_path = self.base_dir / "leetcode"
        
    def is_leetcode_problem(self, problem_name: str) -> bool:
        """Check if this is a LeetCode problem request"""
        return problem_name.startswith("leetcode_") or problem_name.isdigit()
    
    def generate_leetcode_visualization(self, problem_identifier: str) -> bool:
        """Generate visualization for a LeetCode problem"""
        try:
            from leetcode.manager import LeetCodeVisualizationManager
            
            # Extract problem ID
            if problem_identifier.startswith("leetcode_"):
                problem_id = int(problem_identifier.replace("leetcode_", ""))
            else:
                problem_id = int(problem_identifier)
            
            # Check if LeetCode repo exists
            if not self.leetcode_repo_path.exists():
                print(f"❌ LeetCode repository not found at: {self.leetcode_repo_path}")
                print("   Please ensure the LeetCode solutions repository is cloned in the project directory")
                return False
            
            # Generate the visualization
            manager = LeetCodeVisualizationManager(self.leetcode_repo_path, self.problems_dir)
            scene_file = manager.generate_visualization_for_problem(problem_id)
            
            if scene_file:
                print(f"✅ Generated LeetCode visualization: {scene_file}")
                return True
            else:
                print(f"❌ Failed to generate visualization for LeetCode problem {problem_id}")
                return False
                
        except ImportError:
            print("❌ LeetCode visualization module not found")
            return False
        except ValueError:
            print(f"❌ Invalid problem ID: {problem_identifier}")
            return False
        except Exception as e:
            print(f"❌ Error generating LeetCode visualization: {e}")
            return False
        
    def list_available_problems(self):
        """List all available problem visualizations"""
        problems = []
        for file in self.problems_dir.glob("*.py"):
            if file.name != "__init__.py":
                problem_name = file.stem.replace("_", " ").title()
                problems.append((file.stem, problem_name))
        
        # Add LeetCode problems option
        if self.leetcode_repo_path.exists():
            problems.append(("leetcode", "🎯 Generate LeetCode Problem"))
        
        return problems
    
    def combine_videos(self, problem_name, quality="medium_quality"):
        """Combine all scene videos into one complete video"""
        # Map quality to directory name
        quality_dir_map = {
            "low_quality": "480p15",
            "medium_quality": "720p30", 
            "high_quality": "1080p60"
        }
        quality_dir = quality_dir_map.get(quality, "720p30")
        
        # Video directory path
        video_dir = self.base_dir / "media" / "videos" / problem_name / quality_dir
        
        if not video_dir.exists():
            print(f"❌ Video directory not found: {video_dir}")
            return False
        
        # Define the order of scenes for combining
        scene_order = [
            f"{problem_name.replace('_', ' ').title().replace(' ', '')}ProblemStatement.mp4",
            f"{problem_name.replace('_', ' ').title().replace(' ', '')}Visualization.mp4", 
            f"{problem_name.replace('_', ' ').title().replace(' ', '')}CodeSolution.mp4"
        ]
        
        # For LeetCode problems, we need to handle the naming differently
        if problem_name.startswith("leetcode_"):
            # Extract problem info and create proper names
            problem_id = problem_name.replace("leetcode_", "")
            
            # Check what files actually exist
            actual_files = list(video_dir.glob("*.mp4"))
            actual_names = [f.name for f in actual_files]
            
            print(f"Available video files: {actual_names}")
            
            # Define proper chronological order
            def get_scene_priority(filename):
                """Return priority for scene ordering (lower = earlier)"""
                name_lower = filename.lower()
                if 'problemstatement' in name_lower:
                    return 1  # First: Problem Statement
                elif 'visualization' in name_lower:
                    return 2  # Second: Visualization 
                elif 'codesolution' in name_lower:
                    return 3  # Third: Code Solution
                elif 'thankyou' in name_lower or 'conclusion' in name_lower:
                    return 4  # Fourth: Thank you/Conclusion
                else:
                    return 5  # Any other scenes last
            
            # Sort by priority, then alphabetically for same priority
            scene_order = sorted(actual_names, key=lambda x: (get_scene_priority(x), x))
        
        # Find existing video files in the correct order
        video_files = []
        for scene_file in scene_order:
            video_path = video_dir / scene_file
            if video_path.exists():
                video_files.append(str(video_path))
                print(f"✅ Found: {scene_file}")
            else:
                print(f"⚠️  Missing: {scene_file}")
        
        if not video_files:
            print("❌ No video files found to combine!")
            return False
        
        if len(video_files) == 1:
            print("ℹ️  Only one video file found, no combination needed.")
            return True
        
        # Create output filename
        output_file = video_dir / f"Complete{problem_name.replace('_', ' ').title().replace(' ', '')}Presentation.mp4"
        
        try:
            # Create temporary file list for ffmpeg
            with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
                for video_file in video_files:
                    f.write(f"file '{video_file}'\n")
                filelist_path = f.name
            
            # Use ffmpeg to combine videos
            cmd = [
                "ffmpeg", "-f", "concat", "-safe", "0", 
                "-i", filelist_path, 
                "-c", "copy", 
                str(output_file),
                "-y"  # Overwrite output file if exists
            ]
            
            print(f"🎬 Combining videos into: {output_file.name}")
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            # Clean up temporary file
            os.unlink(filelist_path)
            
            if result.returncode == 0:
                print(f"✅ Combined video created: {output_file}")
                return True
            else:
                print(f"❌ FFmpeg error: {result.stderr}")
                return False
                
        except FileNotFoundError:
            print("❌ FFmpeg not found! Please install FFmpeg to combine videos.")
            print("   Download from: https://ffmpeg.org/download.html")
            return False
        except Exception as e:
            print(f"❌ Error combining videos: {e}")
            return False
    
    def run_problem(self, problem_name, scene_name=None, quality="medium_quality", combine=True):
        """Run a specific problem visualization"""
        
        # Check if this is a LeetCode problem
        if self.is_leetcode_problem(problem_name):
            print(f"🎯 Detected LeetCode problem: {problem_name}")
            
            # Generate the visualization if it doesn't exist
            success = self.generate_leetcode_visualization(problem_name)
            if not success:
                return False
            
            # Update problem_name to the generated file
            if problem_name.startswith("leetcode_"):
                problem_file_name = problem_name
            else:
                problem_file_name = f"leetcode_{problem_name}"
        else:
            problem_file_name = problem_name
        
        problem_file = self.problems_dir / f"{problem_file_name}.py"
        
        if not problem_file.exists():
            print(f"❌ Problem '{problem_file_name}' not found!")
            return False
        
        try:
            # Map quality settings to Manim's format
            quality_map = {
                "low_quality": "-ql",
                "medium_quality": "-qm", 
                "high_quality": "-qh"
            }
            quality_flag = quality_map.get(quality, "-qm")
            
            if scene_name:
                cmd = [
                    sys.executable, "-m", "manim", "render",
                    str(problem_file),
                    scene_name,
                    quality_flag
                ]
            else:
                # Render all scenes automatically when no specific scene is provided
                cmd = [
                    sys.executable, "-m", "manim", "render",
                    str(problem_file),
                    "-a",  # Render all scenes
                    quality_flag
                ]
            
            print(f"🎬 Running: {' '.join(cmd)}")
            result = subprocess.run(cmd, cwd=self.base_dir)
            
            success = result.returncode == 0
            
            # If successful and no specific scene was requested, combine all videos
            if success and scene_name is None and combine:
                print("\n🎯 Combining all scene videos into one complete presentation...")
                self.combine_videos(problem_file_name, quality)
            
            return success
            
        except Exception as e:
            print(f"❌ Error running visualization: {e}")
            return False
    
    def run_interactive(self):
        """Run interactive mode to select problems"""
        print("🎬 Coding Visualizer - Interactive Mode")
        print("=" * 50)
        
        problems = self.list_available_problems()
        if not problems:
            print("❌ No problems found!")
            return
        
        print("Available problems:")
        for i, (file_name, display_name) in enumerate(problems, 1):
            print(f"{i}. {display_name}")
        
        try:
            choice = int(input("\nSelect a problem (number): ")) - 1
            if 0 <= choice < len(problems):
                file_name, display_name = problems[choice]
                
                # Special handling for LeetCode option
                if file_name == "leetcode":
                    self._handle_leetcode_interactive()
                    return
                
                print(f"\n🎯 Selected: {display_name}")
                
                # Ask for scene selection
                scenes = self._get_scenes_from_file(file_name)
                if len(scenes) > 1:
                    print("\nAvailable scenes:")
                    for i, scene in enumerate(scenes, 1):
                        print(f"{i}. {scene}")
                    print(f"{len(scenes) + 1}. All scenes")
                    
                    scene_choice = int(input("\nSelect scene (number): ")) - 1
                    if 0 <= scene_choice < len(scenes):
                        scene_name = scenes[scene_choice]
                    elif scene_choice == len(scenes):
                        scene_name = None  # Run all scenes
                    else:
                        print("❌ Invalid scene selection!")
                        return
                else:
                    scene_name = None
                
                # Ask for quality
                print("\nQuality options:")
                print("1. Low (480p)")
                print("2. Medium (720p)")
                print("3. High (1080p)")
                
                quality_choice = input("Select quality (1-3, default: 2): ").strip()
                quality_map = {
                    "1": "low_quality",
                    "2": "medium_quality", 
                    "3": "high_quality"
                }
                quality = quality_map.get(quality_choice, "medium_quality")
                
                # Run the visualization
                success = self.run_problem(file_name, scene_name, quality)
                if success:
                    print("✅ Visualization completed successfully!")
                else:
                    print("❌ Visualization failed!")
            else:
                print("❌ Invalid selection!")
                
        except (ValueError, KeyboardInterrupt):
            print("\n👋 Goodbye!")
    
    def _handle_leetcode_interactive(self):
        """Handle interactive LeetCode problem selection"""
        try:
            from leetcode.manager import LeetCodeVisualizationManager
            
            manager = LeetCodeVisualizationManager(self.leetcode_repo_path, self.problems_dir)
            manager.interactive_selection()
            
        except ImportError:
            print("❌ LeetCode visualization module not available")
        except Exception as e:
            print(f"❌ Error in LeetCode mode: {e}")
    
    def _get_scenes_from_file(self, problem_name):
        """Extract scene class names from a problem file"""
        scenes = []
        problem_file = self.problems_dir / f"{problem_name}.py"
        
        try:
            with open(problem_file, 'r') as f:
                content = f.read()
                
            # Simple parsing to find class definitions that inherit from Scene
            for line in content.split('\n'):
                if line.startswith('class ') and '(Scene)' in line:
                    class_name = line.split('class ')[1].split('(')[0].strip()
                    scenes.append(class_name)
                    
        except Exception:
            pass
        
        return scenes

def main():
    """Main entry point"""
    runner = CodingVisualizerRunner()
    
    if len(sys.argv) > 1:
        # Command line mode
        problem_name = sys.argv[1]
        scene_name = sys.argv[2] if len(sys.argv) > 2 else None
        quality = sys.argv[3] if len(sys.argv) > 3 else "medium_quality"
        combine = sys.argv[4].lower() == 'true' if len(sys.argv) > 4 else True
        
        success = runner.run_problem(problem_name, scene_name, quality, combine)
        sys.exit(0 if success else 1)
    else:
        # Interactive mode
        runner.run_interactive()

if __name__ == "__main__":
    main()