"""
Code Post-Processor for Manim

Fixes common issues in AI-generated Manim code, particularly around
SurroundingRectangle with indexed MathTex elements.
"""

import re
from typing import List


def fix_surrounding_rectangles(code: str) -> str:
    """
    Fix or remove SurroundingRectangle calls that use indexed access to MathTex.
    
    Problem: AI often generates code like:
        SurroundingRectangle(equation[0][5], ...)
    
    This doesn't work reliably because MathTex indexing is unpredictable.
    
    Solution: Comment out these problematic lines with an explanation.
    
    Args:
        code: Raw generated Manim code
    
    Returns:
        Fixed code with problematic SurroundingRectangle calls commented out
    """
    lines = code.split('\n')
    fixed_lines = []
    skip_next = False  # Initialize before the loop
    
    for line in lines:
        # Pattern: SurroundingRectangle with indexed access like equation[0][5]
        # Matches: SurroundingRectangle(something[X][Y], ...) or SurroundingRectangle(something[X:Y], ...)
        if 'SurroundingRectangle(' in line:
            # Check if it has indexed or sliced access
            # Pattern: variable[number] or variable[number:number] inside SurroundingRectangle()
            if re.search(r'SurroundingRectangle\([^,\)]*\[\d+\]', line) or \
               re.search(r'SurroundingRectangle\([^,\)]*\[\d+:\d*\]', line):
                # Replace with empty VGroup to prevent NameError in subsequent animations
                # while avoiding the rendering crash
                indent = len(line) - len(line.lstrip())
                # Extract variable name if present (e.g. "box = SurroundingRectangle(...)")
                var_name = ""
                if "=" in line:
                    var_name = line.split("=")[0].strip()
                    replacement = ' ' * indent + f'{var_name} = VGroup()  # Replaced invalid SurroundingRectangle'
                else:
                    # If no assignment, just comment it out as it won't be referenced
                    replacement = ' ' * indent + '# ' + line.lstrip() + '  # Auto-disabled: indexed SurroundingRectangle'
                
                fixed_lines.append(replacement)
            else:
                fixed_lines.append(line)
        else:
            fixed_lines.append(line)
    
    return '\n'.join(fixed_lines)


def remove_problematic_indexing(code: str) -> str:
    """
    More aggressive approach: Remove entire blocks that use indexed MathTex highlighting.
    
    This removes the entire voiceover block if it contains indexed SurroundingRectangle.
    """
    # For now, use the commenting approach which is safer
    return fix_surrounding_rectangles(code)


def fix_undefined_colors(code: str) -> str:
    """
    Fix undefined color constants by replacing them with valid Manim colors.
    
    Common issues:
    - ORANGE_A, ORANGE_B, etc. -> ORANGE
    - RED_A, RED_B, etc. -> RED
    - Similar patterns for other colors
    
    Args:
        code: Raw generated code
    
    Returns:
        Code with undefined colors replaced
    """
    # Color mappings: undefined variants -> standard colors
    color_replacements = {
        # Orange variants
        r'\bORANGE_[A-Z]\b': 'ORANGE',
        # Red variants
        r'\bRED_[A-Z]\b': 'RED',
        # Blue variants
        r'\bBLUE_[A-Z]\b': 'BLUE',
        # Green variants
        r'\bGREEN_[A-Z]\b': 'GREEN',
        # Yellow variants
        r'\bYELLOW_[A-Z]\b': 'YELLOW',
        # Purple variants
        r'\bPURPLE_[A-Z]\b': 'PURPLE',
        # Pink variants
        r'\bPINK_[A-Z]\b': 'PINK',
        # Teal variants
        r'\bTEAL_[A-Z]\b': 'TEAL',
        # Gray variants
        r'\bGRAY_[A-Z]\b': 'GRAY',
    }
    
    for pattern, replacement in color_replacements.items():
        code = re.sub(pattern, replacement, code)
    
    return code


def clean_duplicate_imports(code: str) -> str:
    """
    Remove duplicate and incorrect imports.
    Keep only the correct imports for our architecture.
    """
    lines = code.split('\n')
    cleaned_lines = []
    seen_imports = set()
    
    # Imports to remove (old/wrong patterns)
    bad_imports = [
        'from manim_voiceover import VoiceoverScene',
        'from manimator.services import ElevenLabsService',
        'from manim_voiceover.services.gtts import GTTSService',
        'from manim_voiceover.services.elevenlabs import ElevenLabsService',
    ]
    
    for line in lines:
        stripped = line.strip()
        
        # Skip bad imports
        if any(bad_import in line for bad_import in bad_imports):
            continue
        
        # Track and deduplicate good imports
        if stripped.startswith('from ') or stripped.startswith('import '):
            if stripped not in seen_imports:
                seen_imports.add(stripped)
                cleaned_lines.append(line)
        else:
            cleaned_lines.append(line)
    
    return '\n'.join(cleaned_lines)


def fix_invalid_manim_parameters(code: str) -> str:
    """
    Fix invalid Manim parameters that cause runtime errors.
    
    Common issues:
    - rounded_corners parameter in Rectangle (doesn't exist)
    - corner_radius parameter in Rectangle (doesn't exist in older Manim versions)
    - Invalid parameter names
    """
    # Remove corner_radius from Rectangle/RoundedRectangle calls (line by line)
    lines = code.split('\n')
    fixed_lines = []
    
    for line in lines:
        # Skip lines that are just the parameter
        if 'corner_radius' in line and '=' in line:
            # Check if this is a standalone parameter line
            stripped = line.strip()
            if stripped.startswith('corner_radius=') or stripped.startswith('rounded_corners='):
                # Skip this line entirely
                continue
        
        # Remove inline corner_radius/rounded_corners parameters
        line = re.sub(r',\s*corner_radius\s*=\s*[^,)]+', '', line)
        line = re.sub(r',\s*rounded_corners\s*=\s*[^,)]+', '', line)
        
        # Remove invalid scale_tips parameter (often hallucinated for scale() method)
        if 'scale_tips' in line and '=' in line:
            stripped = line.strip()
            if stripped.startswith('scale_tips='):
                continue
                
        # Handle inline scale_tips removal (both with and without leading comma)
        line = re.sub(r',\s*scale_tips\s*=\s*[^,)]+', '', line)
        line = re.sub(r'\bscale_tips\s*=\s*[^,)]+\s*,?', '', line)
        
        fixed_lines.append(line)
    
    return '\n'.join(fixed_lines)


def fix_visual_leaks(code: str) -> str:
    """
    Prevent 'visual memory leaks' where reassigning a variable leaves the old Mobject on screen.
    Also enforces 'Title Exclusivity' to prevent multiple titles from stacking.
    
    Problem 1 (Variable Reuse):
        text = Text("A")
        self.play(Write(text))
        text = Text("B") # "A" is still on screen!
        
    Problem 2 (Title Stacking):
        title1 = Text("Intro")
        self.play(Write(title1))
        title2 = Text("Chapter 1") # "Intro" is still on screen!
        
    Solution:
        Inject `self.remove(var)` checks.
    """
    lines = code.split('\n')
    fixed_lines = []
    
    # Regex to detect Mobject assignments
    # Matches: var = ClassName(...)
    # We broaden this to catch any Capitalized class instantiation to be safer
    assignment_pattern = re.compile(r'^\s*([a-zA-Z_]\w*)\s*=\s*([A-Z]\w*)\(')
    
    # Track variables that have been assigned
    assigned_vars = set()
    # Track variables that look like titles
    active_titles = []
    
    for line in lines:
        match = assignment_pattern.match(line)
        if match:
            var_name = match.group(1)
            class_name = match.group(2)
            
            indent = len(line) - len(line.lstrip())
            indent_str = ' ' * indent
            
            # 1. Handle Variable Reuse (Same variable name)
            # We inject this for EVERY assignment to handle loops correctly.
            # In a loop, the line `t = Text(...)` appears once but runs multiple times.
            # By injecting the check, we ensure the previous iteration's object is removed.
            removal = f"{indent_str}if '{var_name}' in locals(): self.remove({var_name})"
            fixed_lines.append(removal)
            
            # 2. Handle Title Exclusivity (Different variable names, but both are titles)
            # Check if this looks like a title (contains 'title', 'header', 'heading')
            # But exclude 'subtitle' or 'sub_title'
            is_title = re.search(r'title|header|heading', var_name, re.IGNORECASE) and \
                       not re.search(r'sub', var_name, re.IGNORECASE)
            
            if is_title:
                # If we are creating a new title, remove ALL previous titles
                for old_title in active_titles:
                    # Don't remove if it's the same variable (handled above)
                    if old_title != var_name:
                        removal = f"{indent_str}if '{old_title}' in locals(): self.remove({old_title})"
                        fixed_lines.append(removal)
                
                # Add this to active titles (if not already there)
                if var_name not in active_titles:
                    active_titles.append(var_name)
            
            assigned_vars.add(var_name)
            fixed_lines.append(line)
        else:
            fixed_lines.append(line)
            
    return '\n'.join(fixed_lines)


def post_process_code(code: str) -> str:
    """
    Main entry point for code post-processing.
    
    Applies all fixes to AI-generated Manim code.
    
    Args:
        code: Raw generated code
    
    Returns:
        Cleaned and fixed code
    """
    # Check if we need to add header (before making changes)
    has_undefined_colors = bool(re.search(r'\b(ORANGE|RED|BLUE|GREEN|YELLOW|PURPLE|PINK|TEAL|GRAY)_[A-Z]\b', code))
    
    # Apply fixes in order
    code = clean_duplicate_imports(code)
    code = fix_undefined_colors(code)
    code = fix_invalid_manim_parameters(code)
    code = fix_invalid_manim_parameters(code)
    code = fix_surrounding_rectangles(code)
    code = fix_visual_leaks(code)
    
    # Add header comment explaining post-processing
    header = """# NOTE: This code has been automatically post-processed to fix common issues.
# Indexed SurroundingRectangle calls have been disabled as they don't reliably
# highlight the intended equation parts in MathTex objects.
# Undefined color constants have been replaced with standard Manim colors.
# Invalid Manim parameters have been removed or corrected.

"""
    
    # Only add header if we actually made changes
    if '# Auto-disabled:' in code or has_undefined_colors or 'rounded_corners' in code:
        code = header + code
    
    return code


def validate_code_structure(code: str) -> List[str]:
    """
    Validate the generated code for common issues.
    
    Returns:
        List of warning messages (empty if no issues)
    """
    warnings = []
    
    # Check for common issues
    if 'SurroundingRectangle(' in code:
        if re.search(r'SurroundingRectangle\([^,\)]*\[\d+\]', code):
            warnings.append("Code contains indexed SurroundingRectangle calls (will be auto-fixed)")
    
    if 'from manim_voiceover.services.gtts import GTTSService' in code:
        warnings.append("Code still uses deprecated GTTSService (should use ElevenLabsService)")
    
    return warnings