import re
import sys

def enhance_logging(input_file, output_file):
    """Add comprehensive logging to the file"""
    
    with open(input_file, 'r') as f:
        content = f.read()
    
    lines = content.split('\n')
    new_lines = []
    
    # Track function and class context
    current_function = None
    current_class = None
    line_number = 0
    in_docstring = False
    docstring_char = None
    
    for line in lines:
        line_number += 1
        stripped = line.strip()
        
        # Handle docstrings
        if '"""' in line or "'''" in line:
            if not in_docstring:
                in_docstring = True
                docstring_char = '"""' if '"""' in line else "'''"
            elif docstring_char in line:
                in_docstring = False
                docstring_char = None
            new_lines.append(line)
            continue
            
        if in_docstring:
            new_lines.append(line)
            continue
        
        # Skip empty lines and comments
        if not stripped or stripped.startswith('#'):
            new_lines.append(line)
            continue
            
        # Skip import statements
        if stripped.startswith(('import ', 'from ')):
            new_lines.append(line)
            continue
            
        # Skip lines that already have logging
        if 'logger.' in stripped or 'logging.' in stripped:
            new_lines.append(line)
            continue
        
        # Track class definitions
        class_match = re.match(r'(\s*)class\s+(\w+)', line)
        if class_match:
            current_class = class_match.group(2)
            new_lines.append(line)
            continue
        
        # Track function definitions
        func_match = re.match(r'(\s*)def\s+(\w+)', line)
        if func_match:
            current_function = func_match.group(2)
            indent = func_match.group(1)
            new_lines.append(line)
            # Add function entry logging
            new_lines.append(f'{indent}    logger.info(f"[FUNC_ENTER] {current_function}() called")')
            continue
        
        # Get indentation
        indent = len(line) - len(line.lstrip())
        indent_str = ' ' * indent
        
        # Patterns for lines that need logging
        patterns_to_log = [
            # Variable assignments
            (r'^(\s*)(\w+)\s*=\s*(.+)$', 'assignment'),
            # Attribute assignments
            (r'^(\s*)(self\.\w+)\s*=\s*(.+)$', 'attribute_assignment'),
            # Method calls (not in assignments)
            (r'^(\s*).*\.(\w+)\([^)]*\)(?!\s*=)', 'method_call'),
            # Function calls
            (r'^(\s*)(\w+)\([^)]*\)(?!\s*=)', 'function_call'),
            # Return statements
            (r'^(\s*)return\s+(.+)$', 'return'),
            # Raise statements
            (r'^(\s*)raise\s+(.+)$', 'raise'),
            # Conditional statements
            (r'^(\s*)if\s+(.+):$', 'if_statement'),
            (r'^(\s*)elif\s+(.+):$', 'elif_statement'),
            (r'^(\s*)else\s*:$', 'else_statement'),
            # Loop statements
            (r'^(\s*)for\s+(.+):$', 'for_loop'),
            (r'^(\s*)while\s+(.+):$', 'while_loop'),
            # Exception handling
            (r'^(\s*)try\s*:$', 'try_block'),
            (r'^(\s*)except\s*(.*):', 'except_block'),
            (r'^(\s*)finally\s*:$', 'finally_block'),
        ]
        
        logged = False
        for pattern, log_type in patterns_to_log:
            match = re.match(pattern, line)
            if match:
                # Add before logging
                context = f"{current_class}.{current_function}" if current_class and current_function else current_function or "global"
                
                if log_type == 'assignment':
                    var_name = match.group(2)
                    value = match.group(3)[:30] + "..." if len(match.group(3)) > 30 else match.group(3)
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Assigning {var_name} = {value}")')
                    new_lines.append(line)
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Assignment completed: {var_name}")')
                    
                elif log_type == 'attribute_assignment':
                    attr_name = match.group(2)
                    value = match.group(3)[:30] + "..." if len(match.group(3)) > 30 else match.group(3)
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Setting attribute {attr_name} = {value}")')
                    new_lines.append(line)
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Attribute set: {attr_name}")')
                    
                elif log_type in ['method_call', 'function_call']:
                    method_name = match.group(2) if log_type == 'method_call' else match.group(2)
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Calling {method_name}()")')
                    new_lines.append(line)
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Call completed: {method_name}()")')
                    
                elif log_type == 'return':
                    return_value = match.group(2)[:30] + "..." if len(match.group(2)) > 30 else match.group(2)
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Returning: {return_value}")')
                    new_lines.append(line)
                    new_lines.append(f'{indent_str}logger.info(f"[FUNC_EXIT] {current_function}() returning")')
                    
                elif log_type == 'raise':
                    exception = match.group(2)[:30] + "..." if len(match.group(2)) > 30 else match.group(2)
                    new_lines.append(f'{indent_str}logger.error(f"[{context}] L{line_number}: Raising exception: {exception}")')
                    new_lines.append(line)
                    
                elif log_type == 'if_statement':
                    condition = match.group(2)[:30] + "..." if len(match.group(2)) > 30 else match.group(2)
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Evaluating condition: {condition}")')
                    new_lines.append(line)
                    
                elif log_type in ['for_loop', 'while_loop']:
                    loop_condition = match.group(2)[:30] + "..." if len(match.group(2)) > 30 else match.group(2)
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Starting {log_type}: {loop_condition}")')
                    new_lines.append(line)
                    
                elif log_type == 'try_block':
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Entering try block")')
                    new_lines.append(line)
                    
                elif log_type == 'except_block':
                    exception_type = match.group(2) if match.group(2) else "any"
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Entering except block: {exception_type}")')
                    new_lines.append(line)
                    
                elif log_type == 'finally_block':
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Entering finally block")')
                    new_lines.append(line)
                    
                else:
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Executing: {stripped[:50]}...")')
                    new_lines.append(line)
                    new_lines.append(f'{indent_str}logger.debug(f"[{context}] L{line_number}: Completed")')
                
                logged = True
                break
        
        if not logged:
            new_lines.append(line)
    
    # Write to output file
    with open(output_file, 'w') as f:
        f.write('\n'.join(new_lines))
    
    print(f"Enhanced logging added to {output_file}")

if __name__ == "__main__":
    enhance_logging('services/google-search/main.py', 'services/google-search/main.py')