validators.py

"""
Code validation and security checking utilities
"""

import ast
from typing import Tuple, List, Optional
from config import ALLOWED_IMPORTS, BLOCKED_IMPORTS


def validate_python_syntax(code: str) -> Tuple[bool, str]:
    """
    Validate Python code syntax using AST parsing

    Args:
        code: Python code string to validate

    Returns:
        Tuple of (is_valid, error_message)
        - If valid: (True, "OK")
        - If invalid: (False, "SyntaxError: ...")
    """
    try:
        ast.parse(code)
        return True, "OK"
    except SyntaxError as e:
        error_msg = f"SyntaxError on line {e.lineno}: {e.msg}"
        if e.text:
            error_msg += f"\n  {e.text.strip()}"
        return False, error_msg
    except Exception as e:
        return False, f"Parse error: {str(e)}"


def check_for_dangerous_imports(code: str) -> Tuple[bool, str]:
    """
    Check for potentially dangerous imports that could pose security risks

    Args:
        code: Python code string to analyze

    Returns:
        Tuple of (is_safe, error_message)
        - If safe: (True, "OK")
        - If unsafe: (False, "Blocked import: ...")
    """
    try:
        tree = ast.parse(code)
    except SyntaxError:
        # Syntax errors are handled by validate_python_syntax
        return True, "OK"

    for node in ast.walk(tree):
        # Check regular imports: import os, import sys
        if isinstance(node, ast.Import):
            for alias in node.names:
                module_name = alias.name.split('.')[0]
                if module_name in BLOCKED_IMPORTS:
                    return False, f"Blocked import: '{alias.name}' - This module is restricted for security reasons"
                # Optionally enforce whitelist (commented out for flexibility)
                # if module_name not in ALLOWED_IMPORTS:
                #     return False, f"Unauthorized import: '{alias.name}' - Only {ALLOWED_IMPORTS} are allowed"

        # Check from imports: from os import system
        elif isinstance(node, ast.ImportFrom):
            if node.module:
                module_name = node.module.split('.')[0]
                if module_name in BLOCKED_IMPORTS:
                    return False, f"Blocked import: 'from {node.module}' - This module is restricted for security reasons"

        # Check for dangerous built-in functions
        elif isinstance(node, ast.Name):
            if node.id in BLOCKED_IMPORTS:
                return False, f"Blocked function: '{node.id}' - This function is restricted for security reasons"

    return True, "OK"


def extract_scene_classes(code: str) -> List[str]:
    """
    Extract all Scene class names from Manim code using AST

    Args:
        code: Python code string containing Manim scenes

    Returns:
        List of Scene class names found in the code
    """
    try:
        tree = ast.parse(code)
    except SyntaxError:
        return []

    scene_classes = []

    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            # Check if class inherits from Scene or any Scene-like class
            for base in node.bases:
                # Handle simple inheritance: class MyScene(Scene)
                if isinstance(base, ast.Name):
                    if 'Scene' in base.id:
                        scene_classes.append(node.name)
                        break
                # Handle module.Class inheritance: class MyScene(manim.Scene)
                elif isinstance(base, ast.Attribute):
                    if 'Scene' in base.attr:
                        scene_classes.append(node.name)
                        break

    return scene_classes


def validate_scene_exists(code: str, scene_name: str) -> Tuple[bool, str]:
    """
    Verify that a specific scene class exists in the code

    Args:
        code: Python code string
        scene_name: Name of the Scene class to verify

    Returns:
        Tuple of (exists, error_message)
        - If exists: (True, "OK")
        - If not found: (False, "Scene 'X' not found...")
    """
    scene_classes = extract_scene_classes(code)

    if scene_name in scene_classes:
        return True, "OK"
    else:
        available = ", ".join(scene_classes) if scene_classes else "none"
        return False, f"Scene '{scene_name}' not found in code. Available scenes: {available}"


def auto_detect_scene(code: str) -> Optional[str]:
    """
    Automatically detect the scene to render when user doesn't specify one

    Args:
        code: Python code string

    Returns:
        Scene class name if exactly one is found, None otherwise
    """
    scene_classes = extract_scene_classes(code)

    if len(scene_classes) == 1:
        return scene_classes[0]
    else:
        return None


def check_for_animation_commands(code: str) -> Tuple[bool, str]:
    """
    Check if code contains animation commands (self.play() or self.wait())
    Scenes without these will not generate video frames

    Args:
        code: Python code to check

    Returns:
        Tuple of (has_animations, error_message)
    """
    has_play = 'self.play(' in code
    has_wait = 'self.wait(' in code

    if not (has_play or has_wait):
        return False, (
            "ERROR: Your scene must include 'self.play()' or 'self.wait()' to generate video frames. "
            "Using only 'self.add()' will not create an animation. "
            "Example: self.play(Create(circle)) or self.wait(1)"
        )

    return True, "OK"


def check_for_latex_usage(code: str) -> Tuple[bool, str]:
    """
    Check if code uses LaTeX objects which may cause rendering issues

    Args:
        code: Python code to check

    Returns:
        Tuple of (has_latex, warning_message)
    """
    latex_classes = ['MathTex', 'Tex', 'TexTemplate', 'Text']
    found_latex = []

    for latex_class in latex_classes:
        if latex_class in code:
            found_latex.append(latex_class)

    if found_latex:
        warning = (
            f"Warning: Your code uses {', '.join(found_latex)} which requires LaTeX. "
            f"This may cause timeouts. Consider using geometric shapes instead (Circle, Square, Line, etc.)."
        )
        return True, warning

    return False, "OK"


def validate_code_full(code: str, scene_name: Optional[str] = None) -> Tuple[bool, str, Optional[str]]:
    """
    Perform full code validation including syntax, security, and scene checks

    Args:
        code: Python code to validate
        scene_name: Optional specific scene name to validate

    Returns:
        Tuple of (is_valid, error_message, detected_scene_name)
        - is_valid: True if all checks pass
        - error_message: Error description if validation fails, "OK" otherwise
        - detected_scene_name: Auto-detected scene name (if scene_name was None)
    """
    # Step 1: Check syntax
    is_valid, error_msg = validate_python_syntax(code)
    if not is_valid:
        return False, error_msg, None

    # Step 2: Check for dangerous imports
    is_safe, error_msg = check_for_dangerous_imports(code)
    if not is_safe:
        return False, error_msg, None

    # Step 2.5: Check for animation commands (REQUIRED)
    has_animations, anim_error = check_for_animation_commands(code)
    if not has_animations:
        return False, anim_error, None

    # Step 2.6: Check for LaTeX usage (warning only)
    has_latex, latex_warning = check_for_latex_usage(code)
    if has_latex:
        # Log warning but don't fail
        import logging
        logger = logging.getLogger("manim_studio.validators")
        logger.warning(latex_warning)

    # Step 3: Check scenes
    scene_classes = extract_scene_classes(code)

    if not scene_classes:
        return False, "No Scene classes found in code. Please define at least one class inheriting from manim.Scene", None

    # Step 4: Validate or auto-detect scene
    if scene_name:
        # User specified a scene - verify it exists
        is_valid, error_msg = validate_scene_exists(code, scene_name)
        if not is_valid:
            return False, error_msg, None
        return True, "OK", scene_name
    else:
        # Auto-detect scene
        if len(scene_classes) == 1:
            return True, "OK", scene_classes[0]
        else:
            # Multiple scenes found, user must specify
            available = ", ".join(scene_classes)
            return False, f"Multiple Scene classes found: {available}. Please specify scene_name in your request.", None