# AI-powered fix suggestion generator using Hugging Face models

from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch
from typing import Dict, List, Optional

class FixSuggestionGenerator:
    def __init__(self):
        print("Loading fix suggestion models...")
        
        # Load CodeBERT for code understanding
        self.code_tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
        
        # Load a smaller model for text generation (for fix suggestions)
        # Using a smaller model for speed
        try:
            self.fix_generator = pipeline(
                "text-generation",
                model="microsoft/DialoGPT-small",
                tokenizer="microsoft/DialoGPT-small",
                device=-1  # Use CPU (change to 0 for GPU)
            )
            print("✅ Loaded DialoGPT for fix suggestions")
        except:
            print("⚠️ Could not load DialoGPT, using rule-based suggestions")
            self.fix_generator = None
        
        # Rule-based fixes database
        self.rule_based_fixes = {
            "sql_injection": [
                "Use parameterized queries: `cursor.execute('SELECT * FROM users WHERE id = %s', (user_id,))`",
                "Use ORM: `User.objects.get(id=user_id)`",
                "Use query builder: `query = users.select().where(users.c.id == user_id)`"
            ],
            "hardcoded_secret": [
                "Store in environment variables: `import os; api_key = os.getenv('API_KEY')`",
                "Use secret management service (AWS Secrets Manager, HashiCorp Vault)",
                "Load from config file (not in version control)"
            ],
            "xss_vulnerability": [
                "Escape HTML: `import html; safe_output = html.escape(user_input)`",
                "Use template auto-escaping: `{{ user_input|escape }}` in Django/Jinja2",
                "Use Content Security Policy (CSP) headers"
            ],
            "command_injection": [
                "Use subprocess with argument list: `subprocess.run(['echo', static_text], check=True)`",
                "Validate and sanitize user input before using in commands",
                "Use shlex.quote() for shell arguments"
            ],
            "insecure_deserialization": [
                "Use safe deserialization: `import json; data = json.loads(user_input)`",
                "Use ast.literal_eval() for Python literals",
                "Implement digital signatures for serialized data"
            ],
            "eval_usage": [
                "Replace eval() with ast.literal_eval(): `import ast; data = ast.literal_eval(user_input)`",
                "Use json.loads() for JSON data",
                "Implement a safe expression evaluator"
            ],
            "syntax_error": [
                "Fix the syntax error before security analysis",
                "Check for missing colons, parentheses, or quotes",
                "Use a linter or IDE to identify syntax errors"
            ]
        }
    
    def generate_ai_fix(self, vulnerable_code: str, issue_type: str) -> Optional[str]:
        """Generate fix using AI model"""
        if self.fix_generator is None:
            return None
        
        prompt = f"""
        Vulnerable code: {vulnerable_code}
        Issue type: {issue_type}
        
        Provide a secure alternative code snippet:
        
        Secure code:"""
        
        try:
            result = self.fix_generator(
                prompt,
                max_length=200,
                num_return_sequences=1,
                temperature=0.7,
                truncation=True
            )
            
            generated_text = result[0]['generated_text']
            # Extract just the fix part
            if "Secure code:" in generated_text:
                fix = generated_text.split("Secure code:")[-1].strip()
                return fix
            return generated_text.strip()
            
        except Exception as e:
            print(f"AI fix generation error: {e}")
            return None
    
    def get_fixes(self, vulnerable_code: str, issue_type: str, num_suggestions: int = 3) -> List[str]:
        """Get multiple fix suggestions for an issue"""
        fixes = []
        
        # Try AI-based fix first
        ai_fix = self.generate_ai_fix(vulnerable_code, issue_type)
        if ai_fix:
            fixes.append(f"🤖 AI Suggestion: {ai_fix}")
        
        # Add rule-based fixes
        if issue_type in self.rule_based_fixes:
            rule_fixes = self.rule_based_fixes[issue_type][:num_suggestions]
            for i, fix in enumerate(rule_fixes, start=1):
                fixes.append(f"🔧 Suggestion {i}: {fix}")
        
        # If no fixes found, provide generic advice
        if not fixes:
            fixes = [
                "Review the code for security best practices",
                "Consult OWASP guidelines for this vulnerability type",
                "Use a security linter or static analysis tool"
            ]
        
        return fixes
    
    def generate_fix_patch(self, original_code: str, issue_line: int, issue_type: str) -> Dict:
        # Generate a complete fix patch with context
        lines = original_code.split('\n')
        
        if issue_line < 1 or issue_line > len(lines):
            return {"error": "Invalid line number"}
        
        vulnerable_line = lines[issue_line - 1]
        fixes = self.get_fixes(vulnerable_line, issue_type)
        
        # Create patch suggestion
        patch_suggestions = []
        for i, fix in enumerate(fixes, 1):
            # Create before/after example
            patch = {
                "title": f"Fix {i} for {issue_type}",
                "vulnerable_line": vulnerable_line,
                "suggestion": fix,
                "context_before": lines[max(0, issue_line-3):issue_line-1],
                "context_after": lines[issue_line:min(len(lines), issue_line+2)]
            }
            patch_suggestions.append(patch)
        
        return {
            "vulnerable_line": issue_line,
            "vulnerable_code": vulnerable_line,
            "issue_type": issue_type,
            "fix_suggestions": fixes,
            "patch_suggestions": patch_suggestions,
            "total_suggestions": len(fixes)
        }

# Test the fix generator
if __name__ == "__main__":
    generator = FixSuggestionGenerator()
    
    test_cases = [
        {
            "code": """query = f"SELECT * FROM users WHERE id = {user_id}" """,
            "line": 1,
            "type": "sql_injection"
        },
        {
            "code": """api_key = "sk_live_1234567890" """,
            "line": 1,
            "type": "hardcoded_secret"
        },
        {
            "code": """return f"<div>{user_input}</div>" """,
            "line": 1,
            "type": "xss_vulnerability"
        }
    ]
    
    print("🔧 Testing Fix Suggestion Generator")
    print("="*50)
    
    for test in test_cases:
        print(f"\n📝 Issue: {test['type']}")
        print(f"Code: {test['code']}")
        
        fixes = generator.get_fixes(test['code'], test['type'])
        
        print("💡 Suggested fixes:")
        for i, fix in enumerate(fixes, 1):
            print(f"  {i}. {fix}")
        
        print("-"*30)