backend/ai/learning/attack_optimizer.py

"""
ATTACK OPTIMIZER: Optimizes attack generation based on learning insights.
Modifies attack strategies to maximize success rates based on historical data.
"""

import logging
from typing import Dict, Any, List, Optional, Tuple
from datetime import datetime
import random
from dataclasses import dataclass

from .learning_engine import LearningEngine, AttackStrategy, get_learning_engine
from .memory_store import AttackRecord, get_memory_store

logger = logging.getLogger(__name__)


@dataclass
class OptimizedAttackConfig:
    """Configuration for optimized attack generation."""
    attack_type: str
    priority: float  # 0.0 to 1.0
    weight_multiplier: float  # How much to weight this attack type
    prompt_modifications: List[str]
    target_models: List[str]
    success_prediction: float
    confidence: float
    reasoning: str
    last_updated: datetime


class AttackOptimizer:
    """
    Optimizes attack generation based on learning insights and historical patterns.
    Modifies attack strategies to maximize success rates.
    """
    
    def __init__(self, learning_engine: Optional[LearningEngine] = None, db_session=None):
        """
        Initialize attack optimizer.
        
        Args:
            learning_engine: Learning engine instance for accessing insights
            db_session: Database session for PostgreSQL memory store
        """
        self.learning_engine = learning_engine or get_learning_engine(db_session=db_session)
        self.memory_store = get_memory_store(db_session=db_session)
        
        # Optimization parameters
        self.min_confidence_threshold = 0.6  # Minimum confidence to use insights
        self.success_rate_weight = 0.7  # Weight for historical success rate
        self.insight_weight = 0.3  # Weight for learning insights
        
        # Cache for optimized configurations
        self._config_cache: Dict[str, List[OptimizedAttackConfig]] = {}
        self._cache_expiry_minutes = 60  # Cache expires after 1 hour
        
        logger.info("Attack optimizer initialized")
    
    def optimize_attack_strategy(self, target_model: str, available_attack_types: List[str]) -> List[OptimizedAttackConfig]:
        """
        Optimize attack strategy for a specific model.
        
        Args:
            target_model: Target model name
            available_attack_types: List of available attack types
            
        Returns:
            List of optimized attack configurations
        """
        try:
            cache_key = f"{target_model}_{hash(tuple(available_attack_types))}"
            
            # Check cache
            if self._is_cache_valid(cache_key):
                logger.debug(f"Using cached optimization for {target_model}")
                return self._config_cache[cache_key]
            
            # Get adaptive strategy
            strategy = self.learning_engine.build_adaptive_strategy(target_model)
            if not strategy:
                logger.warning(f"No adaptive strategy available for {target_model}")
                return self._create_default_configurations(available_attack_types)
            
            # Get learning insights
            insights = self.learning_engine.generate_learning_insights(target_model)
            
            # Optimize based on strategy and insights
            optimized_configs = self._create_optimized_configurations(strategy, insights, available_attack_types)
            
            # Cache results
            self._config_cache[cache_key] = optimized_configs
            
            logger.info(f"Generated {len(optimized_configs)} optimized configurations for {target_model}")
            return optimized_configs
            
        except Exception as e:
            logger.error(f"Failed to optimize attack strategy for {target_model}: {e}")
            return self._create_default_configurations(available_attack_types)
    
    def get_attack_priorities(self, target_model: str) -> Dict[str, float]:
        """
        Get priority weights for different attack types against a model.
        
        Args:
            target_model: Target model name
            
        Returns:
            Dictionary mapping attack types to priority weights (0.0 to 1.0)
        """
        try:
            # Get adaptive strategy
            strategy = self.learning_engine.build_adaptive_strategy(target_model)
            if not strategy:
                return {}
            
            priorities = {}
            
            # High priority for primary attack types
            for attack_type in strategy.primary_attack_types:
                priorities[attack_type] = 0.9
            
            # Medium priority for secondary attack types
            for attack_type in strategy.secondary_attack_types:
                priorities[attack_type] = 0.6
            
            # Low priority for avoided attack types
            for attack_type in strategy.avoided_attack_types:
                priorities[attack_type] = 0.1
            
            return priorities
            
        except Exception as e:
            logger.error(f"Failed to get attack priorities for {target_model}: {e}")
            return {}
    
    def optimize_prompt_for_attack(self, attack_type: str, base_prompt: str, target_model: str) -> str:
        """
        Optimize a prompt for a specific attack type and target model.
        
        Args:
            attack_type: Type of attack
            base_prompt: Base prompt to optimize
            target_model: Target model name
            
        Returns:
            Optimized prompt
        """
        try:
            # Get strategy for target model
            strategy = self.learning_engine.build_adaptive_strategy(target_model)
            if not strategy:
                return base_prompt
            
            # Get historical data for this attack type and model
            attacks = self.memory_store.get_attacks_by_model(target_model)
            relevant_attacks = [a for a in attacks if a.attack_type == attack_type and a.success]
            
            if not relevant_attacks:
                return base_prompt
            
            # Extract optimization patterns
            optimized_prompt = self._apply_prompt_optimizations(base_prompt, relevant_attacks, strategy)
            
            return optimized_prompt
            
        except Exception as e:
            logger.error(f"Failed to optimize prompt: {e}")
            return base_prompt
    
    def update_optimization_from_results(self, target_model: str, attack_results: List[AttackRecord]) -> bool:
        """
        Update optimization based on new attack results.
        
        Args:
            target_model: Target model name
            attack_results: New attack results
            
        Returns:
            True if optimization updated successfully
        """
        try:
            # Store results and update learning
            strategy_id = None
            if target_model in self.learning_engine._strategy_cache:
                strategy_id = self.learning_engine._strategy_cache[target_model].strategy_id
            
            if strategy_id:
                success = self.learning_engine.update_strategy_from_results(strategy_id, attack_results)
                
                # Clear cache to force re-optimization
                for cache_key in list(self._config_cache.keys()):
                    if cache_key.startswith(target_model):
                        del self._config_cache[cache_key]
                
                logger.info(f"Updated optimization for {target_model} from {len(attack_results)} results")
                return success
            else:
                logger.warning(f"No strategy found for {target_model}")
                return False
                
        except Exception as e:
            logger.error(f"Failed to update optimization: {e}")
            return False
    
    def get_optimization_report(self, target_model: str) -> Dict[str, Any]:
        """
        Get comprehensive optimization report for a model.
        
        Args:
            target_model: Target model name
            
        Returns:
            Optimization report
        """
        try:
            # Get strategy
            strategy = self.learning_engine.build_adaptive_strategy(target_model)
            
            # Get priorities
            priorities = self.get_attack_priorities(target_model)
            
            # Get insights
            insights = self.learning_engine.generate_learning_insights(target_model)
            
            # Get historical performance
            recent_attacks = self.memory_store.get_attacks_by_model(target_model, limit=100)
            
            report = {
                "target_model": target_model,
                "timestamp": datetime.now().isoformat(),
                "has_strategy": strategy is not None,
                "strategy_confidence": strategy.confidence if strategy else 0.0,
                "predicted_success_rate": strategy.success_prediction if strategy else 0.0,
                "attack_priorities": priorities,
                "primary_attack_types": strategy.primary_attack_types if strategy else [],
                "secondary_attack_types": strategy.secondary_attack_types if strategy else [],
                "avoided_attack_types": strategy.avoided_attack_types if strategy else [],
                "supporting_insights": len(insights),
                "recent_performance": {
                    "total_attacks": len(recent_attacks),
                    "success_rate": sum(1 for a in recent_attacks if a.success) / len(recent_attacks) * 100 if recent_attacks else 0,
                    "avg_safety_score": sum(a.safety_score for a in recent_attacks) / len(recent_attacks) if recent_attacks else 0
                },
                "optimization_status": "active" if strategy else "inactive"
            }
            
            return report
            
        except Exception as e:
            logger.error(f"Failed to generate optimization report: {e}")
            return {"error": str(e)}
    
    def _create_optimized_configurations(self, strategy: AttackStrategy, insights: List, available_attack_types: List[str]) -> List[OptimizedAttackConfig]:
        """Create optimized attack configurations based on strategy and insights."""
        configs = []
        
        # Primary attack types (highest priority)
        for attack_type in strategy.primary_attack_types:
            if attack_type in available_attack_types:
                config = self._create_attack_config(attack_type, strategy, insights, priority=0.9)
                configs.append(config)
        
        # Secondary attack types (medium priority)
        for attack_type in strategy.secondary_attack_types:
            if attack_type in available_attack_types and attack_type not in [c.attack_type for c in configs]:
                config = self._create_attack_config(attack_type, strategy, insights, priority=0.6)
                configs.append(config)
        
        # Available attack types not in strategy (low priority)
        for attack_type in available_attack_types:
            if attack_type not in [c.attack_type for c in configs] and attack_type not in strategy.avoided_attack_types:
                config = self._create_attack_config(attack_type, strategy, insights, priority=0.3)
                configs.append(config)
        
        # Sort by priority
        configs.sort(key=lambda x: x.priority, reverse=True)
        
        return configs
    
    def _create_attack_config(self, attack_type: str, strategy: AttackStrategy, insights: List, priority: float) -> OptimizedAttackConfig:
        """Create optimized configuration for a specific attack type."""
        # Get historical performance
        attacks = self.memory_store.get_attacks_by_model(strategy.target_model)
        relevant_attacks = [a for a in attacks if a.attack_type == attack_type]
        
        success_rate = 0.0
        if relevant_attacks:
            success_rate = sum(1 for a in relevant_attacks if a.success) / len(relevant_attacks) * 100
        
        # Calculate confidence based on data volume and strategy confidence
        data_confidence = min(1.0, len(relevant_attacks) / 20.0)  # More data = higher confidence
        final_confidence = (data_confidence * 0.6) + (strategy.confidence * 0.4)
        
        # Generate prompt modifications
        prompt_modifications = self._generate_prompt_modifications(attack_type, relevant_attacks, insights)
        
        # Calculate weight multiplier
        weight_multiplier = self._calculate_weight_multiplier(attack_type, strategy, success_rate)
        
        # Generate reasoning
        reasoning = self._generate_config_reasoning(attack_type, strategy, success_rate, insights)
        
        return OptimizedAttackConfig(
            attack_type=attack_type,
            priority=priority,
            weight_multiplier=weight_multiplier,
            prompt_modifications=prompt_modifications,
            target_models=[strategy.target_model],
            success_prediction=success_rate,
            confidence=final_confidence,
            reasoning=reasoning,
            last_updated=datetime.now()
        )
    
    def _generate_prompt_modifications(self, attack_type: str, relevant_attacks: List[AttackRecord], insights: List) -> List[str]:
        """Generate prompt modifications based on historical success patterns."""
        modifications = []
        
        if not relevant_attacks:
            return modifications
        
        # Analyze successful prompts
        successful_attacks = [a for a in relevant_attacks if a.success]
        
        if successful_attacks:
            # Length patterns
            avg_length = sum(len(a.prompt) for a in successful_attacks) / len(successful_attacks)
            if avg_length > 200:
                modifications.append("Use longer, more detailed prompts")
            elif avg_length < 100:
                modifications.append("Use shorter, more direct prompts")
            
            # Common keywords
            all_text = " ".join([a.prompt.lower() for a in successful_attacks])
            words = all_text.split()
            word_freq = {}
            for word in words:
                word_freq[word] = word_freq.get(word, 0) + 1
            
            # Find frequently used words
            common_words = [word for word, freq in sorted(word_freq.items(), key=lambda x: x[1], reverse=True)[:5]]
            if common_words:
                modifications.append(f"Consider using keywords: {', '.join(common_words[:3])}")
        
        # Add insight-based modifications
        for insight in insights:
            if attack_type in insight.applies_to_attack_types and insight.confidence >= 0.7:
                if "weakness" in insight.insight_type:
                    modifications.append(f"Exploit identified weakness: {insight.content}")
                elif "pattern" in insight.insight_type:
                    modifications.append(f"Apply successful pattern: {insight.content}")
        
        return modifications
    
    def _calculate_weight_multiplier(self, attack_type: str, strategy: AttackStrategy, success_rate: float) -> float:
        """Calculate weight multiplier for attack type."""
        base_multiplier = 1.0
        
        # Adjust based on strategy priority
        if attack_type in strategy.primary_attack_types:
            base_multiplier *= 1.5
        elif attack_type in strategy.secondary_attack_types:
            base_multiplier *= 1.2
        elif attack_type in strategy.avoided_attack_types:
            base_multiplier *= 0.3
        
        # Adjust based on historical success rate
        if success_rate > 70:
            base_multiplier *= 1.3
        elif success_rate < 30:
            base_multiplier *= 0.7
        
        # Adjust based on strategy confidence
        if strategy.confidence > 0.8:
            base_multiplier *= 1.1
        elif strategy.confidence < 0.5:
            base_multiplier *= 0.9
        
        # Ensure reasonable bounds
        return max(0.1, min(3.0, base_multiplier))
    
    def _generate_config_reasoning(self, attack_type: str, strategy: AttackStrategy, success_rate: float, insights: List) -> str:
        """Generate reasoning for attack configuration."""
        reasoning_parts = []
        
        # Base reasoning
        if attack_type in strategy.primary_attack_types:
            reasoning_parts.append(f"Primary attack type with predicted success rate of {success_rate:.1f}%")
        elif attack_type in strategy.secondary_attack_types:
            reasoning_parts.append(f"Secondary attack type with predicted success rate of {success_rate:.1f}%")
        else:
            reasoning_parts.append(f"Alternative attack type with predicted success rate of {success_rate:.1f}%")
        
        # Strategy-based reasoning
        if strategy.confidence > 0.8:
            reasoning_parts.append(f"High confidence strategy ({strategy.confidence:.1f})")
        
        # Insight-based reasoning
        relevant_insights = [i for i in insights if attack_type in i.applies_to_attack_types]
        if relevant_insights:
            reasoning_parts.append(f"Supported by {len(relevant_insights)} learning insights")
        
        # Data-based reasoning
        reasoning_parts.append(f"Based on {len(self.memory_store.get_attacks_by_type(attack_type))} historical attacks")
        
        return "; ".join(reasoning_parts)
    
    def _create_default_configurations(self, available_attack_types: List[str]) -> List[OptimizedAttackConfig]:
        """Create default configurations when no learning data is available."""
        configs = []
        
        for attack_type in available_attack_types:
            config = OptimizedAttackConfig(
                attack_type=attack_type,
                priority=0.5,  # Default priority
                weight_multiplier=1.0,  # Default weight
                prompt_modifications=["Default configuration - no learning data available"],
                target_models=[],  # No specific targets
                success_prediction=50.0,  # Default prediction
                confidence=0.3,  # Low confidence without data
                reasoning="Default configuration - insufficient historical data",
                last_updated=datetime.now()
            )
            configs.append(config)
        
        return configs
    
    def _apply_prompt_optimizations(self, base_prompt: str, successful_attacks: List[AttackRecord], strategy: AttackStrategy) -> str:
        """Apply optimizations to a prompt based on successful historical attacks."""
        if not successful_attacks:
            return base_prompt
        
        optimized_prompt = base_prompt
        
        # Length optimization
        avg_length = sum(len(a.prompt) for a in successful_attacks) / len(successful_attacks)
        current_length = len(base_prompt)
        
        if current_length < avg_length * 0.8:
            # Prompt is too short, add detail
            optimized_prompt += " Please provide a detailed and comprehensive response."
        elif current_length > avg_length * 1.2:
            # Prompt is too long, make it more concise
            optimized_prompt = base_prompt[:int(avg_length)] + "..."
        
        # Keyword optimization
        all_successful_text = " ".join([a.prompt.lower() for a in successful_attacks])
        words = all_successful_text.split()
        word_freq = {}
        for word in words:
            word_freq[word] = word_freq.get(word, 0) + 1
        
        # Add common successful words if not present
        common_words = [word for word, freq in sorted(word_freq.items(), key=lambda x: x[1], reverse=True)[:3]]
        for word in common_words:
            if word not in optimized_prompt.lower():
                optimized_prompt += f" Consider {word}."
        
        return optimized_prompt
    
    def _is_cache_valid(self, cache_key: str) -> bool:
        """Check if cached configuration is still valid."""
        if cache_key not in self._config_cache:
            return False
        
        # Check if any configuration is older than cache expiry
        for config in self._config_cache[cache_key]:
            age_minutes = (datetime.now() - config.last_updated).total_seconds() / 60
            if age_minutes > self._cache_expiry_minutes:
                return False
        
        return True


# Singleton instance for global access
_attack_optimizer_instance: Optional[AttackOptimizer] = None


def get_attack_optimizer(learning_engine: Optional[LearningEngine] = None, db_session=None) -> AttackOptimizer:
    """
    Get global attack optimizer instance.
    
    Args:
        learning_engine: Learning engine instance
        db_session: Database session for PostgreSQL memory store
        
    Returns:
        AttackOptimizer instance
    """
    global _attack_optimizer_instance
    
    # If database session is provided, reset global instance to ensure PostgreSQL usage
    if db_session is not None:
        _attack_optimizer_instance = None
    
    if _attack_optimizer_instance is None:
        _attack_optimizer_instance = AttackOptimizer(learning_engine, db_session)
    
    return _attack_optimizer_instance