import re
import ast

def extract_code_blocks(text):
    """Extract code blocks from a markdown-formatted text"""
    # Pattern to match code blocks with ```python or ``` markers
    pattern = r'```(?:python)?\s*([\s\S]*?)```'
    matches = re.findall(pattern, text)
    
    # If no code blocks found, try to extract the entire text as code
    if not matches and text.strip():
        # Check if the text looks like Python code (has common imports or patterns)
        if re.search(r'import\s+\w+|def\s+\w+\(|class\s+\w+:|if\s+__name__\s*==\s*[\'"]__main__[\'"]:', text):
            return [text.strip()]
    
    return [match.strip() for match in matches]

def sanitize_code(code):
    """Remove potentially harmful operations from the code"""
    # Basic sanitization - replace known harmful functions
    harmful_patterns = [
        (r'__import__\([\'"]os[\'"]\)', 'None'),
        (r'exec\(', 'print('),
        (r'eval\(', 'print('),
        (r'open\(.*,.*[\'"]w[\'"].*\)', 'open("safe_file.txt", "r")'),
        (r'subprocess\.\w+\(', 'print('),
        (r'os\.system\(', 'print('),
        (r'os\.popen\(', 'print('),
        (r'os\.unlink\(', 'print('),
        (r'os\.remove\(', 'print('),
        (r'shutil\.rmtree\(', 'print('),
    ]
    
    sanitized_code = code
    for pattern, replacement in harmful_patterns:
        sanitized_code = re.sub(pattern, replacement, sanitized_code)
    
    return sanitized_code

def validate_gradio_code(code):
    """Validate that the code only uses Gradio and safe libraries"""
    try:
        # Parse the code into an AST
        tree = ast.parse(code)
        
        # Check imports
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for name in node.names:
                    if name.name not in ['gradio', 'numpy', 'pandas', 'matplotlib', 'PIL', 'os', 'io', 'base64', 
                                         'time', 'datetime', 'json', 'random', 'math', 'sys', 're', 'pathlib',
                                         'collections', 'typing', 'warnings']:
                        return False, f"Unauthorized import: {name.name}"
            
            elif isinstance(node, ast.ImportFrom):
                if node.module not in ['gradio', 'numpy', 'pandas', 'matplotlib', 'PIL', 'os', 'io', 'base64',
                                      'time', 'datetime', 'json', 'random', 'math', 'sys', 're', 'pathlib',
                                      'collections', 'typing', 'warnings', None]:
                    return False, f"Unauthorized import from: {node.module}"
        
        # Basic check for potentially harmful OS operations
        code_str = code.lower()
        harmful_operations = [
            'subprocess', 'system(', 'popen(', 'execve(', 'fork(', 'chmod(',
            'rmdir(', 'remove(', 'unlink(', 'rmtree(', 'shutil.rm', 'socket',
            'urllib.request', 'requests', 'http', 'ftp', 'telnet', 'eval(', 'exec('
        ]
        
        for op in harmful_operations:
            if op in code_str:
                return False, f"Potentially harmful operation detected: {op}"
        
        # Check for launch parameters
        launch_pattern = r'\.launch\s*\(([^)]*)\)'
        launch_matches = re.findall(launch_pattern, code)
        
        for match in launch_matches:
            if 'debug=' in match and 'debug=False' not in match:
                return False, "Debug mode is not allowed"
            if 'share=' in match and 'share=False' not in match and 'share=True' in match:
                return False, "Share mode must be set to False"
        
        return True, None
        
    except SyntaxError as e:
        return False, f"Syntax error in the code: {str(e)}"
    except Exception as e:
        return False, f"Error validating code: {str(e)}"