backend/ai/learning/learning_engine.py

"""
LEARNING ENGINE: Core learning system that extracts insights from historical data.
Builds adaptive attack strategies based on real adversarial outcomes.
"""

import logging
from typing import Dict, Any, List, Optional, Tuple
from datetime import datetime, timedelta
from dataclasses import dataclass, asdict
import json

from .memory_store import MemoryStore, AttackRecord, get_memory_store
from .pattern_analyzer import PatternAnalyzer, get_pattern_analyzer

logger = logging.getLogger(__name__)


@dataclass
class LearningInsight:
    """Single learning insight extracted from historical data."""
    insight_id: str
    insight_type: str
    content: str
    confidence: float
    evidence: List[str]
    created_at: datetime
    applies_to_attack_types: List[str]
    applies_to_models: List[str]
    priority: str  # high, medium, low
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for storage."""
        data = asdict(self)
        data['created_at'] = self.created_at.isoformat()
        return data
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'LearningInsight':
        """Create from dictionary."""
        data['created_at'] = datetime.fromisoformat(data['created_at'])
        return cls(**data)


@dataclass
class AttackStrategy:
    """Adaptive attack strategy based on learning insights."""
    strategy_id: str
    target_model: str
    primary_attack_types: List[str]
    secondary_attack_types: List[str]
    avoided_attack_types: List[str]
    success_prediction: float
    confidence: float
    reasoning: List[str]
    created_at: datetime
    last_updated: datetime
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for storage."""
        data = asdict(self)
        data['created_at'] = self.created_at.isoformat()
        data['last_updated'] = self.last_updated.isoformat()
        return data
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'AttackStrategy':
        """Create from dictionary."""
        data['created_at'] = datetime.fromisoformat(data['created_at'])
        data['last_updated'] = datetime.fromisoformat(data['last_updated'])
        return cls(**data)


class LearningEngine:
    """
    Core learning engine that extracts insights from historical adversarial data
    and builds adaptive attack strategies.
    """
    
    def __init__(self, memory_store: Optional[MemoryStore] = None, pattern_analyzer: Optional[PatternAnalyzer] = None):
        """
        Initialize learning engine.
        
        Args:
            memory_store: Memory store for accessing historical data
            pattern_analyzer: Pattern analyzer for extracting patterns
        """
        self.memory_store = memory_store or get_memory_store()
        self.pattern_analyzer = pattern_analyzer or get_pattern_analyzer(self.memory_store)
        
        # Learning configuration
        self.min_data_points = 10  # Minimum attacks needed for learning
        self.confidence_threshold = 0.7  # Minimum confidence for insights
        self.strategy_refresh_interval = timedelta(hours=6)  # How often to refresh strategies
        
        # Cache for strategies and insights
        self._strategy_cache: Dict[str, AttackStrategy] = {}
        self._insight_cache: List[LearningInsight] = []
        self._last_strategy_update: Dict[str, datetime] = {}
        
        logger.info("Learning engine initialized")
    
    def generate_learning_insights(self, model_name: Optional[str] = None, days_back: int = 30) -> List[LearningInsight]:
        """
        Generate learning insights from historical data.
        
        Args:
            model_name: Specific model to analyze, or None for all models
            days_back: Number of days of historical data to analyze
            
        Returns:
            List of learning insights
        """
        try:
            insights = []
            
            # Get pattern analysis
            if model_name:
                pattern_analysis = self.pattern_analyzer.analyze_model_weaknesses(model_name, days_back)
                insights.extend(self._extract_model_insights(pattern_analysis, model_name))
            else:
                cross_model_insights = self.pattern_analyzer.get_cross_model_insights(days_back)
                insights.extend(self._extract_cross_model_insights(cross_model_insights))
            
            # Filter by confidence and sort by priority
            filtered_insights = [i for i in insights if i.confidence >= self.confidence_threshold]
            filtered_insights.sort(key=lambda x: self._priority_score(x.priority), reverse=True)
            
            # Update cache
            self._insight_cache.extend(filtered_insights)
            
            logger.info(f"Generated {len(filtered_insights)} learning insights")
            return filtered_insights
            
        except Exception as e:
            logger.error(f"Failed to generate learning insights: {e}")
            return []
    
    def build_adaptive_strategy(self, target_model: str, days_back: int = 30) -> Optional[AttackStrategy]:
        """
        Build adaptive attack strategy for a specific model.
        
        Args:
            target_model: Target model name
            days_back: Number of days of historical data to analyze
            
        Returns:
            Adaptive attack strategy or None if insufficient data
        """
        try:
            # Check if we need to refresh the strategy
            if self._should_refresh_strategy(target_model):
                strategy = self._create_new_strategy(target_model, days_back)
                if strategy:
                    self._strategy_cache[target_model] = strategy
                    self._last_strategy_update[target_model] = datetime.now()
                    logger.info(f"Created new adaptive strategy for {target_model}")
                    return strategy
                else:
                    logger.warning(f"Failed to create strategy for {target_model}")
                    return None
            else:
                # Return cached strategy
                if target_model in self._strategy_cache:
                    logger.debug(f"Using cached strategy for {target_model}")
                    return self._strategy_cache[target_model]
                else:
                    return None
                    
        except Exception as e:
            logger.error(f"Failed to build adaptive strategy for {target_model}: {e}")
            return None
    
    def update_strategy_from_results(self, strategy_id: str, attack_results: List[AttackRecord]) -> bool:
        """
        Update adaptive strategy based on new attack results.
        
        Args:
            strategy_id: Strategy identifier
            attack_results: New attack results to incorporate
            
        Returns:
            True if strategy updated successfully
        """
        try:
            # Find the strategy
            strategy = None
            for s in self._strategy_cache.values():
                if s.strategy_id == strategy_id:
                    strategy = s
                    break
            
            if not strategy:
                logger.warning(f"Strategy {strategy_id} not found for update")
                return False
            
            # Store new results
            self.memory_store.store_batch_attacks(attack_results)
            
            # Update pattern metrics for affected attack types
            attack_types = set(a.attack_type for a in attack_results)
            for attack_type in attack_types:
                self.pattern_analyzer.update_pattern_metrics(attack_type)
            
            # Refresh strategy with new data
            updated_strategy = self._create_new_strategy(strategy.target_model, days_back=7)
            if updated_strategy:
                self._strategy_cache[strategy.target_model] = updated_strategy
                self._last_strategy_update[strategy.target_model] = datetime.now()
                logger.info(f"Updated strategy {strategy_id} with new results")
                return True
            
            return False
            
        except Exception as e:
            logger.error(f"Failed to update strategy {strategy_id}: {e}")
            return False
    
    def get_strategy_explanation(self, strategy_id: str) -> Optional[Dict[str, Any]]:
        """
        Get detailed explanation of why a strategy was chosen.
        
        Args:
            strategy_id: Strategy identifier
            
        Returns:
            Strategy explanation or None if not found
        """
        try:
            # Find the strategy
            strategy = None
            for s in self._strategy_cache.values():
                if s.strategy_id == strategy_id:
                    strategy = s
                    break
            
            if not strategy:
                return None
            
            # Find supporting insights
            supporting_insights = []
            for insight in self._insight_cache:
                if (strategy.target_model in insight.applies_to_models or
                    any(at in insight.applies_to_attack_types for at in strategy.primary_attack_types)):
                    supporting_insights.append(insight)
            
            # Build explanation
            explanation = {
                "strategy_id": strategy.strategy_id,
                "target_model": strategy.target_model,
                "success_prediction": strategy.success_prediction,
                "confidence": strategy.confidence,
                "primary_attack_types": strategy.primary_attack_types,
                "secondary_attack_types": strategy.secondary_attack_types,
                "avoided_attack_types": strategy.avoided_attack_types,
                "reasoning": strategy.reasoning,
                "supporting_insights": [insight.to_dict() for insight in supporting_insights],
                "evidence_summary": self._summarize_evidence(strategy, supporting_insights),
                "last_updated": strategy.last_updated.isoformat()
            }
            
            return explanation
            
        except Exception as e:
            logger.error(f"Failed to get strategy explanation: {e}")
            return None
    
    def get_learning_summary(self, model_name: Optional[str] = None) -> Dict[str, Any]:
        """
        Get comprehensive learning summary.
        
        Args:
            model_name: Specific model, or None for all models
            
        Returns:
            Learning summary dictionary
        """
        try:
            summary = {
                "timestamp": datetime.now().isoformat(),
                "model_filter": model_name,
                "total_insights": len(self._insight_cache),
                "cached_strategies": len(self._strategy_cache),
                "learning_status": "active"
            }
            
            if model_name:
                # Model-specific summary
                if model_name in self._strategy_cache:
                    strategy = self._strategy_cache[model_name]
                    summary["model_strategy"] = strategy.to_dict()
                
                model_insights = [i for i in self._insight_cache if model_name in i.applies_to_models]
                summary["model_insights"] = [insight.to_dict() for insight in model_insights]
                summary["model_insight_count"] = len(model_insights)
            else:
                # Overall summary
                summary["all_strategies"] = {model: strategy.to_dict() for model, strategy in self._strategy_cache.items()}
                summary["all_insights"] = [insight.to_dict() for insight in self._insight_cache]
            
            return summary
            
        except Exception as e:
            logger.error(f"Failed to get learning summary: {e}")
            return {"error": str(e)}
    
    def _extract_model_insights(self, pattern_analysis: Dict[str, Any], model_name: str) -> List[LearningInsight]:
        """Extract insights from model-specific pattern analysis."""
        insights = []
        
        if "error" in pattern_analysis:
            return insights
        
        # Weak category insights
        weak_categories = pattern_analysis.get("weak_categories", [])
        for i, category in enumerate(weak_categories[:3]):  # Top 3
            insight = LearningInsight(
                insight_id=f"weak_category_{model_name}_{i}",
                insight_type="weakness",
                content=f"Model {model_name} is vulnerable to {category} attacks",
                confidence=0.8,
                evidence=[f"High success rate in {category} category"],
                created_at=datetime.now(),
                applies_to_attack_types=[category],
                applies_to_models=[model_name],
                priority="high"
            )
            insights.append(insight)
        
        # Successful pattern insights
        successful_patterns = pattern_analysis.get("successful_patterns", [])
        for i, pattern in enumerate(successful_patterns[:3]):  # Top 3
            insight = LearningInsight(
                insight_id=f"successful_pattern_{model_name}_{i}",
                insight_type="pattern",
                content=f"Pattern {pattern.get('type', 'unknown')} shows high success rate",
                confidence=0.7,
                evidence=[pattern.get('description', 'No description')],
                created_at=datetime.now(),
                applies_to_attack_types=[pattern.get('type', 'unknown')],
                applies_to_models=[model_name],
                priority="medium"
            )
            insights.append(insight)
        
        # Vulnerability indicator insights
        vulnerability_indicators = pattern_analysis.get("vulnerability_indicators", [])
        for i, indicator in enumerate(vulnerability_indicators[:3]):  # Top 3
            insight = LearningInsight(
                insight_id=f"vulnerability_{model_name}_{i}",
                insight_type="indicator",
                content=f"Vulnerability indicator: {indicator}",
                confidence=0.6,
                evidence=[indicator],
                created_at=datetime.now(),
                applies_to_attack_types=["multiple"],
                applies_to_models=[model_name],
                priority="medium"
            )
            insights.append(insight)
        
        return insights
    
    def _extract_cross_model_insights(self, cross_model_analysis: Dict[str, Any]) -> List[LearningInsight]:
        """Extract insights from cross-model pattern analysis."""
        insights = []
        
        if "error" in cross_model_analysis:
            return insights
        
        # Universal weakness insights
        universal_weaknesses = cross_model_analysis.get("universal_weaknesses", [])
        for i, weakness in enumerate(universal_weaknesses[:3]):  # Top 3
            insight = LearningInsight(
                insight_id=f"universal_weakness_{i}",
                insight_type="universal_weakness",
                content=f"Universal weakness: {weakness}",
                confidence=0.9,
                evidence=[f"Observed across multiple models"],
                created_at=datetime.now(),
                applies_to_attack_types=["multiple"],
                applies_to_models=["multiple"],
                priority="high"
            )
            insights.append(insight)
        
        # Attack hierarchy insights
        hierarchy = cross_model_analysis.get("attack_type_hierarchy", {})
        most_effective = hierarchy.get("most_effective", [])
        for i, attack_info in enumerate(most_effective[:3]):  # Top 3
            insight = LearningInsight(
                insight_id=f"effective_attack_{i}",
                insight_type="effective_attack",
                content=f"Attack type {attack_info['attack_type']} is highly effective",
                confidence=0.8,
                evidence=[f"Success rate: {attack_info['success_rate']:.1f}%"],
                created_at=datetime.now(),
                applies_to_attack_types=[attack_info['attack_type']],
                applies_to_models=["multiple"],
                priority="high"
            )
            insights.append(insight)
        
        # Emerging pattern insights
        emerging_patterns = cross_model_analysis.get("emerging_patterns", [])
        for i, pattern in enumerate(emerging_patterns[:3]):  # Top 3
            insight = LearningInsight(
                insight_id=f"emerging_pattern_{i}",
                insight_type="emerging_pattern",
                content=f"Emerging pattern: {pattern.get('description', 'Unknown')}",
                confidence=0.7,
                evidence=[pattern.get('description', 'No description')],
                created_at=datetime.now(),
                applies_to_attack_types=[pattern.get('attack_type', 'unknown')],
                applies_to_models=["multiple"],
                priority="medium"
            )
            insights.append(insight)
        
        return insights
    
    def _create_new_strategy(self, target_model: str, days_back: int) -> Optional[AttackStrategy]:
        """Create a new adaptive strategy for the target model."""
        try:
            # Get model-specific analysis
            pattern_analysis = self.pattern_analyzer.analyze_model_weaknesses(target_model, days_back)
            
            if "error" in pattern_analysis:
                logger.warning(f"No pattern analysis available for {target_model}")
                return None
            
            # Check if we have enough data
            total_attacks = pattern_analysis.get("total_attacks_analyzed", 0)
            if total_attacks < self.min_data_points:
                logger.warning(f"Insufficient data for {target_model}: {total_attacks} attacks")
                return None
            
            # Extract strategy components
            weak_categories = pattern_analysis.get("weak_categories", [])
            successful_patterns = pattern_analysis.get("successful_patterns", [])
            failed_patterns = pattern_analysis.get("failed_patterns", [])
            
            # Determine primary attack types (from weak categories and successful patterns)
            primary_attack_types = []
            for category in weak_categories[:3]:  # Top 3 weak categories
                primary_attack_types.append(category)
            
            for pattern in successful_patterns[:2]:  # Top 2 successful patterns
                pattern_type = pattern.get('type', 'unknown')
                if pattern_type not in primary_attack_types:
                    primary_attack_types.append(pattern_type)
            
            # Determine avoided attack types (from failed patterns)
            avoided_attack_types = []
            for pattern in failed_patterns[:3]:  # Top 3 failed patterns
                pattern_type = pattern.get('type', 'unknown')
                if pattern_type not in avoided_attack_types:
                    avoided_attack_types.append(pattern_type)
            
            # Determine secondary attack types
            secondary_attack_types = []
            for category in weak_categories[3:5]:  # Next 2 weak categories
                if category not in primary_attack_types:
                    secondary_attack_types.append(category)
            
            # Predict success rate based on historical data
            success_prediction = self._predict_success_rate(target_model, primary_attack_types, days_back)
            
            # Calculate confidence
            confidence = min(0.9, total_attacks / 100.0)  # More data = higher confidence
            
            # Generate reasoning
            reasoning = self._generate_strategy_reasoning(pattern_analysis, primary_attack_types, avoided_attack_types)
            
            # Create strategy
            strategy = AttackStrategy(
                strategy_id=f"strategy_{target_model}_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
                target_model=target_model,
                primary_attack_types=primary_attack_types,
                secondary_attack_types=secondary_attack_types,
                avoided_attack_types=avoided_attack_types,
                success_prediction=success_prediction,
                confidence=confidence,
                reasoning=reasoning,
                created_at=datetime.now(),
                last_updated=datetime.now()
            )
            
            return strategy
            
        except Exception as e:
            logger.error(f"Failed to create new strategy for {target_model}: {e}")
            return None
    
    def _predict_success_rate(self, target_model: str, attack_types: List[str], days_back: int) -> float:
        """Predict success rate for attack types against target model."""
        try:
            # Get recent attacks for this model
            attacks = self.memory_store.get_attacks_by_model(target_model)
            cutoff_date = datetime.now() - timedelta(days=days_back)
            recent_attacks = [a for a in attacks if a.timestamp >= cutoff_date]
            
            # Filter by attack types
            relevant_attacks = [a for a in recent_attacks if a.attack_type in attack_types]
            
            if not relevant_attacks:
                # If no relevant data, use overall success rate
                all_relevant = [a for a in recent_attacks if a.attack_type in attack_types]
                if all_relevant:
                    return sum(1 for a in all_relevant if a.success) / len(all_relevant) * 100
                else:
                    return 50.0  # Default prediction
            
            # Calculate success rate
            success_rate = sum(1 for a in relevant_attacks if a.success) / len(relevant_attacks) * 100
            
            # Apply trend adjustment
            temporal_trends = self.pattern_analyzer._analyze_temporal_trends(recent_attacks)
            trend = temporal_trends.get("trend", "stable")
            
            if trend == "improving":
                success_rate = min(95.0, success_rate + 10)  # Boost for improving trend
            elif trend == "declining":
                success_rate = max(5.0, success_rate - 10)  # Reduce for declining trend
            
            return success_rate
            
        except Exception as e:
            logger.error(f"Failed to predict success rate: {e}")
            return 50.0  # Default prediction
    
    def _generate_strategy_reasoning(self, pattern_analysis: Dict[str, Any], primary_types: List[str], avoided_types: List[str]) -> List[str]:
        """Generate reasoning for strategy selection."""
        reasoning = []
        
        # Reasoning for primary attack types
        if primary_types:
            reasoning.append(f"Primary focus on {primary_types[0]} due to identified weakness")
            if len(primary_types) > 1:
                reasoning.append(f"Secondary focus on {primary_types[1]} for diversification")
        
        # Reasoning for avoided attack types
        if avoided_types:
            reasoning.append(f"Avoiding {avoided_types[0]} due to low success rate")
        
        # Reasoning based on data volume
        total_attacks = pattern_analysis.get("total_attacks_analyzed", 0)
        reasoning.append(f"Strategy based on {total_attacks} historical attacks")
        
        # Reasoning based on temporal trends
        temporal_trends = pattern_analysis.get("temporal_trends", {})
        trend = temporal_trends.get("trend", "stable")
        if trend == "improving":
            reasoning.append("Recent trend shows improving effectiveness")
        elif trend == "declining":
            reasoning.append("Recent trend shows declining effectiveness")
        
        return reasoning
    
    def _should_refresh_strategy(self, target_model: str) -> bool:
        """Check if strategy should be refreshed."""
        if target_model not in self._strategy_cache:
            return True
        
        if target_model not in self._last_strategy_update:
            return True
        
        return datetime.now() - self._last_strategy_update[target_model] > self.strategy_refresh_interval
    
    def _priority_score(self, priority: str) -> int:
        """Convert priority to numeric score for sorting."""
        priority_scores = {"high": 3, "medium": 2, "low": 1}
        return priority_scores.get(priority, 0)
    
    def _summarize_evidence(self, strategy: AttackStrategy, insights: List[LearningInsight]) -> Dict[str, Any]:
        """Summarize evidence supporting the strategy."""
        evidence_summary = {
            "total_insights": len(insights),
            "high_confidence_insights": len([i for i in insights if i.confidence >= 0.8]),
            "weakness_evidence": len([i for i in insights if i.insight_type == "weakness"]),
            "pattern_evidence": len([i for i in insights if i.insight_type == "pattern"]),
            "indicator_evidence": len([i for i in insights if i.insight_type == "indicator"]),
            "evidence_types": list(set(i.insight_type for i in insights))
        }
        
        return evidence_summary


# Singleton instance for global access
_learning_engine_instance: Optional[LearningEngine] = None


def get_learning_engine(memory_store: Optional[MemoryStore] = None, pattern_analyzer: Optional[PatternAnalyzer] = None, db_session=None) -> LearningEngine:
    """
    Get global learning engine instance.
    
    Args:
        memory_store: Memory store instance
        pattern_analyzer: Pattern analyzer instance
        db_session: Database session for PostgreSQL memory store
        
    Returns:
        LearningEngine instance
    """
    global _learning_engine_instance
    
    # If database session is provided, reset global instance to ensure PostgreSQL usage
    if db_session is not None:
        _learning_engine_instance = None
    
    if _learning_engine_instance is None:
        # Use PostgreSQL memory store if db_session is provided
        if db_session is not None:
            from .postgresql_memory_store import PostgreSQLMemoryStore
            memory_store = PostgreSQLMemoryStore(db_session, enable_persistence=True)
        else:
            memory_store = memory_store or get_memory_store()
        
        _learning_engine_instance = LearningEngine(memory_store, pattern_analyzer)
    
    return _learning_engine_instance