refactor_agent.py

"""
Refactor Agent - AI-powered code refactoring using Gemini API

This module provides automated code refactoring capabilities by taking review
findings from ReviewAgent and applying fixes to the source code. It uses
Google's Gemini API to generate specific line-by-line refactoring suggestions.

The agent implements safety features:
- Backs up original code before making changes
- Auto-applies only safe fixes (style, best_practice categories)
- Requires human approval for bug/security fixes

Usage:
    agent = RefactorAgent(api_key="your-gemini-api-key")
    suggestions = agent.generate_fixes(code, findings)
    fixed_code = agent.auto_refactor(code, findings)
"""

import os
import json
import re
import shutil
from dataclasses import dataclass, field
from typing import Optional, List, Dict
from enum import Enum
from datetime import datetime

# Google Gemini API client
from google import genai


class RefactorCategory(Enum):
    """
    Categories for refactoring suggestions.
    
    These match the categories from ReviewAgent but add an APPROVED status
    for tracking which fixes have been human-approved.
    """
    BUG = "bug"
    SECURITY = "security"
    STYLE = "style"
    PERFORMANCE = "performance"
    BEST_PRACTICE = "best_practice"
    APPROVED = "approved"  # Human-approved fixes


class ApprovalStatus(Enum):
    """
    Approval status for applying fixes.
    
    AUTO_SAFE: Can be applied automatically (style, best_practice)
    REQUIRES_APPROVAL: Needs human approval (bug, security)
    APPROVED: Has been approved by human
    REJECTED: Rejected by human
    """
    AUTO_SAFE = "auto_safe"
    REQUIRES_APPROVAL = "requires_approval"
    APPROVED = "approved"
    REJECTED = "rejected"


@dataclass
class RefactorSuggestion:
    """
    Represents a single refactoring suggestion.
    
    This dataclass holds all information about a proposed code change,
    including what to change, what to change it to, and why.
    
    Attributes:
        line: Line number where the change should be applied (1-indexed)
        original_code: The original code at that line
        new_code: The new code to replace it with
        explanation: Human-readable explanation of why this change helps
        category: Type of issue being fixed (bug, security, style, etc.)
        approval_status: Current approval status of this fix
    """
    line: int
    original_code: str
    new_code: str
    explanation: str
    category: RefactorCategory
    approval_status: ApprovalStatus = ApprovalStatus.REQUIRES_APPROVAL
    
    def can_auto_apply(self) -> bool:
        """
        Check if this suggestion can be applied automatically.
        
        Only style and best_practice fixes are auto-applied.
        Bug and security fixes require human approval.
        
        Returns:
            True if the fix can be applied without human approval
        """
        return self.category in [
            RefactorCategory.STYLE, 
            RefactorCategory.BEST_PRACTICE
        ]
    
    def to_dict(self) -> dict:
        """
        Convert suggestion to dictionary format for serialization.
        
        Returns:
            Dictionary representation of this suggestion
        """
        return {
            "line": self.line,
            "original_code": self.original_code,
            "new_code": self.new_code,
            "explanation": self.explanation,
            "category": self.category.value,
            "approval_status": self.approval_status.value
        }


@dataclass
class RefactorResult:
    """
    Container for refactoring results.
    
    Tracks the original code, applied suggestions, and the resulting
    fixed code, plus any issues that were found.
    
    Attributes:
        original_code: The original code before refactoring
        fixed_code: The code after applying suggestions
        applied_suggestions: List of suggestions that were applied
        rejected_suggestions: List of suggestions that were rejected
        backup_path: Path to the backup file (if created)
    """
    original_code: str
    fixed_code: str
    applied_suggestions: List[RefactorSuggestion] = field(default_factory=list)
    rejected_suggestions: List[RefactorSuggestion] = field(default_factory=list)
    backup_path: Optional[str] = None
    
    def get_approval_summary(self) -> str:
        """
        Get a human-readable summary of what was approved/rejected.
        
        Returns:
            Summary string describing the refactoring results
        """
        applied = len(self.applied_suggestions)
        rejected = len(self.rejected_suggestions)
        
        if applied == 0 and rejected == 0:
            return "No changes made - no suggestions provided."
        
        parts = []
        if applied > 0:
            parts.append(f"{applied} fix(es) applied")
        if rejected > 0:
            parts.append(f"{rejected} fix(es) rejected")
        
        return ", ".join(parts)


class RefactorAgent:
    """
    AI Agent that applies code fixes based on review findings.
    
    This agent takes the output from ReviewAgent and uses Gemini to generate
    specific line-by-line refactoring suggestions. It implements safety
    controls to prevent automatic application of risky fixes.
    
    Attributes:
        api_key: Google Gemini API key for authentication
        backup_dir: Directory to store code backups
        
    Example:
        >>> agent = RefactorAgent(api_key="AIza...")
        >>> findings = review_agent.review_file("example.py")
        >>> result = agent.auto_refactor(code, findings.findings)
        >>> print(result.fixed_code)
    """
    
    # System prompt that guides Gemini's refactoring approach
    SYSTEM_PROMPT = """You are an expert code refactoring assistant. Your job is to take code review 
findings and generate specific, line-by-line fixes.

For each finding, you must provide:
- line: The exact line number to modify (1-indexed)
- original_code: The exact original code at that line (must match exactly!)
- new_code: The exact new code to replace it with
- explanation: Brief explanation of why this fix helps
- category: The type of fix (bug, security, style, performance, best_practice)

IMPORTANT:
1. Match original_code EXACTLY - include all whitespace, quotes, and characters
2. Provide complete line replacements - don't partial replace
3. For multi-line issues, target the most specific line possible
4. If original_code doesn't match any line exactly, note that in explanation

Respond in JSON format with an array of suggestions. Each suggestion is:
{"line": 10, "original_code": "x = y", "new_code": "result = y", "explanation": "...", "category": "style"}

Only respond with the JSON, no other text."""

    def __init__(self, api_key: str, backup_dir: str = ".refactor_backups"):
        """
        Initialize the RefactorAgent with a Gemini API key.
        
        Args:
            api_key: Google Gemini API key. Can be provided directly or via
                    GEMINI_API_KEY environment variable if not provided.
            backup_dir: Directory to store code backups (default: .refactor_backups)
            
        Raises:
            ValueError: If no valid API key is provided
        """
        # Use provided key or fall back to environment variable
        self.api_key = api_key or os.environ.get("GEMINI_API_KEY")
        
        if not self.api_key:
            raise ValueError(
                "Gemini API key required. Provide as parameter or set GEMINI_API_KEY env var."
            )
        
        # Initialize the Gemini client
        self.client = genai.Client(api_key=self.api_key)
        
        # Set up backup directory
        self.backup_dir = backup_dir
        self._ensure_backup_dir()
    
    def _ensure_backup_dir(self):
        """Create backup directory if it doesn't exist."""
        os.makedirs(self.backup_dir, exist_ok=True)
    
    def _create_backup(self, code: str, file_path: Optional[str] = None) -> str:
        """
        Create a backup of the original code.
        
        Args:
            code: The code content to back up
            file_path: Optional file path to name the backup after
            
        Returns:
            Path to the backup file
        """
        # Generate timestamp-based filename
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        if file_path:
            # Use original filename with timestamp
            base_name = os.path.basename(file_path)
            name, ext = os.path.splitext(base_name)
            backup_name = f"{name}_{timestamp}{ext}"
        else:
            backup_name = f"code_{timestamp}.py"
        
        backup_path = os.path.join(self.backup_dir, backup_name)
        
        # Write backup file
        with open(backup_path, 'w', encoding='utf-8') as f:
            f.write(code)
        
        return backup_path
    
    def generate_fixes(self, code: str, findings: List, file_path: Optional[str] = None) -> List[RefactorSuggestion]:
        """
        Generate specific line-by-line fixes from review findings.
        
        This method takes the original code and review findings from ReviewAgent,
        then uses Gemini to generate specific refactoring suggestions that can
        be applied to the code.
        
        Args:
            code: The original source code to fix
            findings: List of ReviewFinding objects from ReviewAgent
            file_path: Optional path to the source file for context
            
        Returns:
            List of RefactorSuggestion objects representing proposed fixes
            
        Raises:
            RuntimeError: If API call fails
        """
        # Convert findings to a format Gemini can understand
        findings_text = self._format_findings_for_prompt(findings)
        
        # Build the prompt
        prompt = f"""{self.SYSTEM_PROMPT}

Original code to refactor:
```{self._detect_language(file_path) if file_path else 'python'}
{code}
```

Review findings to address:
{findings_text}

Generate specific refactoring suggestions. Make sure original_code EXACTLY matches
the lines in the original code above.

Respond with a JSON array of suggestions."""
        
        try:
            # Call Gemini API to generate fixes
            response = self.client.models.generate_content(
                model="gemini-2.0-flash",
                contents=prompt
            )
            
            # Parse the response
            suggestions_data = self._parse_json_response(response.text)
            
            # Convert to RefactorSuggestion objects
            suggestions = []
            for item in suggestions_data:
                # Parse the category
                category_raw = item.get("category", "best_practice").lower()
                category_map = {
                    "bugs": RefactorCategory.BUG,
                    "security": RefactorCategory.SECURITY,
                    "styles": RefactorCategory.STYLE,
                    "performance": RefactorCategory.PERFORMANCE,
                    "best_practices": RefactorCategory.BEST_PRACTICE,
                    "bug": RefactorCategory.BUG,
                    "security": RefactorCategory.SECURITY,
                    "style": RefactorCategory.STYLE,
                    "performance": RefactorCategory.PERFORMANCE,
                    "best_practice": RefactorCategory.BEST_PRACTICE,
                }
                category = category_map.get(category_raw, RefactorCategory.BEST_PRACTICE)
                
                # Determine approval status based on category
                if category in [RefactorCategory.STYLE, RefactorCategory.BEST_PRACTICE]:
                    approval_status = ApprovalStatus.AUTO_SAFE
                else:
                    approval_status = ApprovalStatus.REQUIRES_APPROVAL
                
                suggestion = RefactorSuggestion(
                    line=item.get("line", 0),
                    original_code=item.get("original_code", ""),
                    new_code=item.get("new_code", ""),
                    explanation=item.get("explanation", ""),
                    category=category,
                    approval_status=approval_status
                )
                suggestions.append(suggestion)
            
            return suggestions
            
        except Exception as e:
            raise RuntimeError(f"Failed to generate fixes: {str(e)}") from e
    
    def apply_fixes(self, code: str, suggestions: List[RefactorSuggestion], 
                    auto_approve_safe: bool = True) -> RefactorResult:
        """
        Apply the suggested fixes to the code.
        
        This method applies fixes to the code, respecting approval status.
        By default, style and best_practice fixes are auto-approved while
        bug and security fixes require explicit approval.
        
        Args:
            code: The original source code
            suggestions: List of RefactorSuggestion objects to apply
            auto_approve_safe: If True, automatically approve style/best_practice fixes
            
        Returns:
            RefactorResult containing the fixed code and metadata
        """
        # Create backup first
        backup_path = self._create_backup(code)
        
        # Split code into lines
        lines = code.split('\n')
        
        applied = []
        rejected = []
        
        # Sort suggestions by line number (descending to apply from bottom up)
        sorted_suggestions = sorted(suggestions, key=lambda s: s.line, reverse=True)
        
        for suggestion in sorted_suggestions:
            # Check if we should apply this fix
            should_apply = False
            
            if auto_approve_safe and suggestion.can_auto_apply():
                # Auto-apply style and best_practice fixes
                should_apply = True
                suggestion.approval_status = ApprovalStatus.AUTO_SAFE
            elif suggestion.approval_status == ApprovalStatus.APPROVED:
                # Already explicitly approved
                should_apply = True
            
            if should_apply:
                # Verify the line matches
                if self._verify_and_apply_fix(lines, suggestion):
                    applied.append(suggestion)
                else:
                    # Line didn't match - reject
                    suggestion.approval_status = ApprovalStatus.REJECTED
                    rejected.append(suggestion)
            else:
                # Not approved - reject
                rejected.append(suggestion)
        
        # Reconstruct the code
        fixed_code = '\n'.join(lines)
        
        return RefactorResult(
            original_code=code,
            fixed_code=fixed_code,
            applied_suggestions=applied,
            rejected_suggestions=rejected,
            backup_path=backup_path
        )
    
    def _verify_and_apply_fix(self, lines: List[str], suggestion: RefactorSuggestion) -> bool:
        """
        Verify that the original code matches and apply the fix.
        
        Args:
            lines: List of code lines (modified in place)
            suggestion: The suggestion to apply
            
        Returns:
            True if fix was applied, False if verification failed
        """
        # Convert to 0-indexed for list access
        line_index = suggestion.line - 1
        
        # Check if line exists
        if line_index < 0 or line_index >= len(lines):
            return False
        
        # Normalize both strings for comparison (strip whitespace)
        original_normalized = lines[line_index].strip()
        suggestion_normalized = suggestion.original_code.strip()
        
        # Try exact match first, then normalized match
        if lines[line_index] != suggestion.original_code:
            if original_normalized != suggestion_normalized:
                # Try to find a close match
                return False
            
            # Use the existing line content if normalized matches
            suggestion.original_code = lines[line_index]
        
        # Apply the fix
        lines[line_index] = suggestion.new_code
        return True
    
    def preview_changes(self, code: str, suggestions: List[RefactorSuggestion]) -> Dict:
        """
        Show what changes would be made without actually applying them.
        
        This is useful for reviewing what will change before committing
        to the refactoring.
        
        Args:
            code: The original source code
            suggestions: List of RefactorSuggestion objects to preview
            
        Returns:
            Dictionary containing:
            - suggestions_by_category: Suggestions grouped by category
            - auto_safe_count: Number of fixes that can be auto-applied
            - approval_required_count: Number requiring human approval
            - line_changes: Preview of each change
        """
        suggestions_by_category = {
            "style": [],
            "best_practice": [],
            "bug": [],
            "security": [],
            "performance": []
        }
        
        for suggestion in suggestions:
            cat = suggestion.category.value
            if cat in suggestions_by_category:
                suggestions_by_category[cat].append(suggestion)
        
        auto_safe = sum(1 for s in suggestions if s.can_auto_apply())
        approval_required = len(suggestions) - auto_safe
        
        # Generate line-by-line preview
        line_changes = []
        for suggestion in sorted(suggestions, key=lambda s: s.line):
            line_changes.append({
                "line": suggestion.line,
                "original": suggestion.original_code,
                "new": suggestion.new_code,
                "explanation": suggestion.explanation,
                "category": suggestion.category.value,
                "auto_apply": suggestion.can_auto_apply()
            })
        
        return {
            "suggestions_by_category": suggestions_by_category,
            "auto_safe_count": auto_safe,
            "approval_required_count": approval_required,
            "line_changes": line_changes
        }
    
    def auto_refactor(self, code: str, findings: List, 
                     file_path: Optional[str] = None,
                     auto_approve_safe: bool = True) -> RefactorResult:
        """
        Full refactoring pipeline: generate fixes and apply them.
        
        This is the main entry point for automated refactoring. It:
        1. Generates specific line-by-line fixes from findings
        2. Creates a backup of the original code
        3. Applies fixes (auto-applying safe ones, requiring approval for others)
        
        Args:
            code: The original source code
            findings: List of ReviewFinding objects from ReviewAgent
            file_path: Optional path for backup naming
            auto_approve_safe: If True, automatically apply style/best_practice fixes
            
        Returns:
            RefactorResult with the fixed code and metadata
            
        Example:
            >>> agent = RefactorAgent(api_key="AIza...")
            >>> findings = review_agent.review_file("example.py").findings
            >>> result = agent.auto_refactor(code, findings)
            >>> print(result.fixed_code)  # The refactored code
            >>> print(result.get_approval_summary())  # What was applied
        """
        # Step 1: Generate fixes from findings
        suggestions = self.generate_fixes(code, findings, file_path)
        
        if not suggestions:
            # No suggestions - return original code
            return RefactorResult(
                original_code=code,
                fixed_code=code,
                backup_path=self._create_backup(code, file_path)
            )
        
        # Step 2: Apply fixes
        result = self.apply_fixes(code, suggestions, auto_approve_safe)
        
        return result
    
    def approve_suggestion(self, suggestions: List[RefactorSuggestion], 
                          line_numbers: List[int]) -> List[RefactorSuggestion]:
        """
        Manually approve specific suggestions by line number.
        
        This allows human override for bug/security fixes that were
        initially marked as requiring approval.
        
        Args:
            suggestions: List of RefactorSuggestion objects
            line_numbers: List of line numbers to approve
            
        Returns:
            Updated list of suggestions with approval status changed
        """
        for suggestion in suggestions:
            if suggestion.line in line_numbers:
                suggestion.approval_status = ApprovalStatus.APPROVED
        return suggestions
    
    def reject_suggestion(self, suggestions: List[RefactorSuggestion],
                         line_numbers: List[int]) -> List[RefactorSuggestion]:
        """
        Manually reject specific suggestions by line number.
        
        Args:
            suggestions: List of RefactorSuggestion objects
            line_numbers: List of line numbers to reject
            
        Returns:
            Updated list of suggestions with approval status changed
        """
        for suggestion in suggestions:
            if suggestion.line in line_numbers:
                suggestion.approval_status = ApprovalStatus.REJECTED
        return suggestions
    
    def _format_findings_for_prompt(self, findings: List) -> str:
        """
        Format review findings for the Gemini prompt.
        
        Args:
            findings: List of ReviewFinding objects
            
        Returns:
            Formatted string for the prompt
        """
        formatted = []
        for i, finding in enumerate(findings, 1):
            line_info = f"Line {finding.line_number}" if finding.line_number else "Whole file"
            formatted.append(
                f"{i}. [{finding.category.value.upper()}] {line_info}\n"
                f"   Issue: {finding.description}\n"
                f"   Suggested fix: {finding.suggested_fix}"
            )
        return '\n'.join(formatted) if formatted else "No specific findings provided."
    
    def _detect_language(self, file_path: Optional[str]) -> str:
        """
        Detect programming language from file extension.
        
        Args:
            file_path: Path to the file
            
        Returns:
            Language string (defaults to 'python')
        """
        if not file_path:
            return 'python'
        
        extension_map = {
            '.py': 'python',
            '.js': 'javascript',
            '.ts': 'typescript',
            '.jsx': 'javascript',
            '.tsx': 'typescript',
            '.java': 'java',
            '.c': 'c',
            '.cpp': 'cpp',
            '.go': 'go',
            '.rs': 'rust',
            '.rb': 'ruby',
            '.php': 'php',
            '.swift': 'swift',
            '.kt': 'kotlin',
        }
        
        ext = os.path.splitext(file_path)[1].lower()
        return extension_map.get(ext, 'python')
    
    def _parse_json_response(self, response_text: str) -> List[dict]:
        """
        Parse JSON from Gemini's response text with robust error handling.
        
        Handles:
        - Markdown code blocks (```json ... ```)
        - Extra text before/after the JSON
        - Multiple JSON arrays in response
        - Invalid JSON gracefully
        
        Args:
            response_text: Raw response from Gemini API
            
        Returns:
            List of suggestion dictionaries (empty list on failure)
        """
        if not response_text:
            print("Warning: Empty response from Gemini API")
            return []
        
        # Strategy 1: Strip markdown code blocks and extract JSON
        # Try ```json ... ``` first
        json_str = None
        for pattern in [
            r'```json\s*(\[[\s\S]*?\])\s*```',  # ```json [...] ```
            r'```\s*(\[[\s\S]*?\])\s*```',       # ``` [...] ```
        ]:
            match = re.search(pattern, response_text)
            if match:
                json_str = match.group(1)
                break
        
        # Strategy 2: If no code block, look for JSON array anywhere in text
        if not json_str:
            # Find the first [ and last ] to extract the array
            start = response_text.find('[')
            end = response_text.rfind(']')
            if start != -1 and end != -1 and end > start:
                json_str = response_text[start:end+1]
        
        # Strategy 3: Try the whole response stripped
        if not json_str:
            json_str = response_text.strip()
        
        # Now try to parse with multiple attempts
        parsing_strategies = [
            # Direct parse
            lambda s: json.loads(s),
            # Strip whitespace
            lambda s: json.loads(s.strip()),
            # Remove common markdown artifacts
            lambda s: json.loads(re.sub(r'^```[\w]*|```$', '', s, flags=re.MULTILINE).strip()),
            # Remove leading/trailing non-JSON text
            lambda s: json.loads(s[s.find('['):s.rfind(']')+1] if '[' in s and ']' in s else s),
        ]
        
        for strategy in parsing_strategies:
            try:
                data = strategy(json_str)
                if isinstance(data, list):
                    return data
                elif isinstance(data, dict) and 'suggestions' in data:
                    return data['suggestions']
                elif isinstance(data, dict):
                    # Wrap single object in list
                    return [data]
            except (json.JSONDecodeError, ValueError, AttributeError) as e:
                continue
        
        # Last resort: try to find any valid JSON array in the original response
        print(f"Warning: All parsing strategies failed, trying fallback search...")
        try:
            # More aggressive extraction - find any [...] pattern
            array_matches = re.findall(r'\[[\s\S]*\]', response_text)
            for match in array_matches:
                try:
                    data = json.loads(match)
                    if isinstance(data, list) and len(data) > 0:
                        # Check if it looks like our expected format
                        if isinstance(data[0], dict) and 'line' in data[0]:
                            return data
                except json.JSONDecodeError:
                    continue
        except Exception:
            pass
        
        print(f"Warning: Failed to parse JSON response after all strategies")
        return []


# =============================================================================
# Demo / Test Code
# =============================================================================

def demo():
    """
    Demonstrate the RefactorAgent with sample code.
    
    This demo shows how to use the RefactorAgent to refactor code
    based on review findings.
    """
    print("=" * 70)
    print("RefactorAgent Demo")
    print("=" * 70)
    
    # Sample Python code with issues (same as review_agent demo)
    sample_code = '''
import os
import sqlite3

def get_user_data(user_id):
    """Get user data from database."""
    conn = sqlite3.connect("users.db")
    cursor = conn.cursor()
    
    # SQL Injection vulnerability!
    query = f"SELECT * FROM users WHERE id = {user_id}"
    cursor.execute(query)
    
    return cursor.fetchone()

def process_data(data):
    """Process some data."""
    # Using bare except - bad practice
    try:
        result = data / 0  # Will always fail
    except:
        pass  # Silently ignoring errors
    
    return result

def inefficient_loop(items):
    """Process items inefficiently."""
    result = []
    for item in items:
        # Creating new list each iteration
        temp = list(items)
        result.append(temp)
    return result

class BadClass:
    # No docstring
    def __init__(self):
        self.x = 1
    
    def do_something(self):
        # Using single letter variable
        for i in range(100):
            x = i * 2
        return x
'''
    
    print("\nOriginal code:")
    print("-" * 40)
    print(sample_code)
    print("-" * 40)
    
    # Create mock findings (simulating ReviewAgent output)
    # In real usage, you'd get these from ReviewAgent
    from review_agent import ReviewFinding, Severity, Category
    
    mock_findings = [
        ReviewFinding(
            severity=Severity.CRITICAL,
            category=Category.SECURITY,
            line_number=12,
            description="SQL Injection vulnerability - user input directly concatenated",
            suggested_fix="Use parameterized queries"
        ),
        ReviewFinding(
            severity=Severity.WARNING,
            category=Category.BEST_PRACTICE,
            line_number=21,
            description="Bare except clause catches all exceptions",
            suggested_fix="Catch specific exceptions or log the error"
        ),
        ReviewFinding(
            severity=Severity.WARNING,
            category=Category.PERFORMANCE,
            line_number=27,
            description="Inefficient loop - creating new list each iteration",
            suggested_fix="Remove temp list creation"
        ),
        ReviewFinding(
            severity=Severity.INFO,
            category=Category.STYLE,
            line_number=33,
            description="Class is missing a docstring",
            suggested_fix="Add a docstring to document the class"
        ),
    ]
    
    # Check for API key
    api_key = os.environ.get("GEMINI_API_KEY")
    
    if not api_key:
        print("\nWARNING: GEMINI_API_KEY not set in environment.")
        print("Running in simulation mode with pre-defined suggestions...\n")
        
        # Create simulated suggestions for demo
        suggestions = [
            RefactorSuggestion(
                line=12,
                original_code='    query = f"SELECT * FROM users WHERE id = {user_id}"',
                new_code='    query = "SELECT * FROM users WHERE id = ?"',
                explanation="Use parameterized query to prevent SQL injection",
                category=RefactorCategory.SECURITY,
                approval_status=ApprovalStatus.REQUIRES_APPROVAL
            ),
            RefactorSuggestion(
                line=21,
                original_code='    except:',
                new_code='    except Exception as e:',
                explanation="Catch specific exception type and store the error",
                category=RefactorCategory.BEST_PRACTICE,
                approval_status=ApprovalStatus.AUTO_SAFE
            ),
            RefactorSuggestion(
                line=22,
                original_code='        pass  # Silently ignoring errors',
                new_code='        print(f"Error processing data: {e}")',
                explanation="Log errors instead of silently ignoring them",
                category=RefactorCategory.BEST_PRACTICE,
                approval_status=ApprovalStatus.AUTO_SAFE
            ),
            RefactorSuggestion(
                line=29,
                original_code='        temp = list(items)',
                new_code='        # Removed unnecessary list creation',
                explanation="Eliminate redundant list copy for better performance",
                category=RefactorCategory.PERFORMANCE,
                approval_status=ApprovalStatus.REQUIRES_APPROVAL
            ),
            RefactorSuggestion(
                line=34,
                original_code='class BadClass:',
                new_code='class BadClass:\n    """A class demonstrating refactoring needs."""',
                explanation="Add docstring to document class purpose",
                category=RefactorCategory.STYLE,
                approval_status=ApprovalStatus.AUTO_SAFE
            ),
        ]
        
        # Show preview
        print("\n--- PREVIEW CHANGES ---")
        preview = RefactorAgent._preview_changes.__get__(object, object)
        
        # Manual preview since we don't have an agent instance
        print("\nLine-by-line changes:")
        for s in sorted(suggestions, key=lambda x: x.line):
            auto = "✓ Auto-safe" if s.can_auto_apply() else "⚠ Requires approval"
            print(f"\nLine {s.line} [{s.category.value}] {auto}")
            print(f"  Original: {s.original_code[:50]}...")
            print(f"  New:      {s.new_code[:50]}...")
            print(f"  Why:      {s.explanation}")
        
        # Apply fixes (auto-approving safe ones)
        print("\n--- APPLYING FIXES ---")
        
        # Create a mock RefactorAgent instance for applying
        class MockAgent:
            def __init__(self):
                self.backup_dir = ".refactor_backups"
                self._ensure_backup_dir()
            
            def _ensure_backup_dir(self):
                os.makedirs(self.backup_dir, exist_ok=True)
            
            def _create_backup(self, code, file_path=None):
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                backup_path = os.path.join(self.backup_dir, f"demo_{timestamp}.py")
                with open(backup_path, 'w') as f:
                    f.write(code)
                return backup_path
        
        agent = MockAgent()
        backup_path = agent._create_backup(sample_code)
        
        # Apply fixes manually for demo
        lines = sample_code.split('\n')
        applied = []
        rejected = []
        
        for s in sorted(suggestions, key=lambda x: x.line, reverse=True):
            if s.can_auto_apply():
                idx = s.line - 1
                if 0 <= idx < len(lines):
                    lines[idx] = s.new_code
                    applied.append(s)
                    print(f"✓ Applied line {s.line} ({s.category.value}): {s.explanation[:40]}...")
            else:
                rejected.append(s)
                print(f"○ Skipped line {s.line} ({s.category.value}) - requires approval")
        
        fixed_code = '\n'.join(lines)
        
        result = RefactorResult(
            original_code=sample_code,
            fixed_code=fixed_code,
            applied_suggestions=applied,
            rejected_suggestions=rejected,
            backup_path=backup_path
        )
        
    else:
        print("\nAPI key found. Running actual refactoring with Gemini...\n")
        
        # Create agent and run full pipeline
        agent = RefactorAgent(api_key=api_key)
        result = agent.auto_refactor(sample_code, mock_findings, auto_approve_safe=True)
    
    # Display results
    print("\n" + "=" * 70)
    print("REFACTOR RESULT")
    print("=" * 70)
    
    print(f"\n{result.get_approval_summary()}")
    print(f"\nBackup saved to: {result.backup_path}")
    
    print("\n--- FIXED CODE ---")
    print(result.fixed_code)
    
    if result.rejected_suggestions:
        print("\n--- REQUIRES APPROVAL ---")
        for s in result.rejected_suggestions:
            print(f"  Line {s.line} ({s.category.value}): {s.explanation}")
    
    print("\n" + "=" * 70)
    print("Demo complete!")
    print("=" * 70)


if __name__ == "__main__":
    demo()