backend/ai/learning/memory_store.py

"""
MEMORY STORE: Persistent storage for adversarial learning data.
Stores and retrieves historical attack results for pattern analysis.
"""

import os
import json
import sqlite3
import logging
from typing import Dict, Any, List, Optional, Tuple, Union
from datetime import datetime, timedelta
from dataclasses import dataclass, asdict
from pathlib import Path

logger = logging.getLogger(__name__)


@dataclass
class AttackRecord:
    """Single attack execution record."""
    attack_id: str
    attack_type: str
    attack_category: str
    target_model: str
    dataset: str
    prompt: str
    success: bool
    safety_score: float
    risk_level: str
    response_text: str
    response_length: int
    inference_time_ms: float
    timestamp: datetime
    metadata: Dict[str, Any]
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for storage."""
        data = asdict(self)
        data['timestamp'] = self.timestamp.isoformat()
        return data
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'AttackRecord':
        """Create from dictionary."""
        data['timestamp'] = datetime.fromisoformat(data['timestamp'])
        return cls(**data)


@dataclass
class PatternMetrics:
    """Metrics for attack patterns."""
    attack_type: str
    total_attempts: int
    successful_attempts: int
    success_rate: float
    avg_safety_score: float
    avg_response_length: int
    avg_inference_time: float
    last_updated: datetime
    weak_indicators: List[str]
    strong_indicators: List[str]


class MemoryStore:
    """
    Persistent memory store for adversarial learning data.
    Stores historical attack results and provides analysis capabilities.
    """
    
    def __init__(self, db_path: Optional[str] = None, enable_persistence: bool = True):
        """
        Initialize memory store.
        
        Args:
            db_path: Path to SQLite database
            enable_persistence: Whether to persist data to disk
        """
        self.enable_persistence = enable_persistence
        self.cache_dir = Path("cache/learning") if enable_persistence else None
        
        if db_path:
            self.db_path = Path(db_path)
        else:
            self.db_path = Path("data/learning_memory.db") if enable_persistence else ":memory:"
        
        # Create directories
        if self.enable_persistence and self.cache_dir:
            self.cache_dir.mkdir(parents=True, exist_ok=True)
            self.db_path.parent.mkdir(parents=True, exist_ok=True)
        
        # Initialize database
        self._init_database()
        
        # In-memory cache for frequent access
        self._attack_cache: List[AttackRecord] = []
        self._pattern_cache: Dict[str, PatternMetrics] = {}
        
        logger.info(f"Memory store initialized: {self.db_path}")
    
    def _init_database(self):
        """Initialize SQLite database with required tables."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                # Attacks table
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS attacks (
                        attack_id TEXT PRIMARY KEY,
                        attack_type TEXT NOT NULL,
                        attack_category TEXT NOT NULL,
                        target_model TEXT NOT NULL,
                        dataset TEXT NOT NULL,
                        prompt TEXT NOT NULL,
                        success BOOLEAN NOT NULL,
                        safety_score REAL NOT NULL,
                        risk_level TEXT NOT NULL,
                        response_text TEXT,
                        response_length INTEGER,
                        inference_time_ms REAL,
                        timestamp TEXT NOT NULL,
                        metadata TEXT
                    )
                """)
                
                # Create indexes separately
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_attacks_timestamp ON attacks(timestamp)")
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_attacks_attack_type ON attacks(attack_type)")
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_attacks_target_model ON attacks(target_model)")
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_attacks_success ON attacks(success)")
                
                # Pattern metrics table
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS pattern_metrics (
                        attack_type TEXT PRIMARY KEY,
                        total_attempts INTEGER NOT NULL,
                        successful_attempts INTEGER NOT NULL,
                        success_rate REAL NOT NULL,
                        avg_safety_score REAL NOT NULL,
                        avg_response_length INTEGER NOT NULL,
                        avg_inference_time REAL NOT NULL,
                        last_updated TEXT NOT NULL,
                        weak_indicators TEXT,
                        strong_indicators TEXT
                    )
                """)
                
                # Learning insights table
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS learning_insights (
                        insight_id TEXT PRIMARY KEY,
                        insight_type TEXT NOT NULL,
                        content TEXT NOT NULL,
                        confidence REAL NOT NULL,
                        created_at TEXT NOT NULL,
                        applies_to_attack_types TEXT
                    )
                """)
                
                # Create indexes for learning insights
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_learning_insights_type ON learning_insights(insight_type)")
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_learning_insights_created_at ON learning_insights(created_at)")
                
                conn.commit()
                logger.info("Database initialized successfully")
                
        except Exception as e:
            logger.error(f"Failed to initialize database: {e}")
            raise
    
    def store_attack(self, attack_record: AttackRecord) -> bool:
        """
        Store a single attack record.
        
        Args:
            attack_record: Attack record to store
            
        Returns:
            True if stored successfully
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                data = attack_record.to_dict()
                
                cursor.execute("""
                    INSERT OR REPLACE INTO attacks 
                    (attack_id, attack_type, attack_category, target_model, dataset,
                     prompt, success, safety_score, risk_level, response_text,
                     response_length, inference_time_ms, timestamp, metadata)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """, (
                    data['attack_id'], data['attack_type'], data['attack_category'],
                    data['target_model'], data['dataset'], data['prompt'],
                    data['success'], data['safety_score'], data['risk_level'],
                    data['response_text'], data['response_length'],
                    data['inference_time_ms'], data['timestamp'], 
                    json.dumps(data['metadata'])
                ))
                
                conn.commit()
                
                # Update cache
                self._attack_cache.append(attack_record)
                
                logger.debug(f"Stored attack record: {attack_record.attack_id}")
                return True
                
        except Exception as e:
            logger.error(f"Failed to store attack record: {e}")
            return False
    
    def store_batch_attacks(self, attack_records: List[AttackRecord]) -> int:
        """
        Store multiple attack records efficiently.
        
        Args:
            attack_records: List of attack records to store
            
        Returns:
            Number of records stored successfully
        """
        if not attack_records:
            return 0
        
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                stored_count = 0
                for record in attack_records:
                    data = record.to_dict()
                    
                    cursor.execute("""
                        INSERT OR REPLACE INTO attacks 
                        (attack_id, attack_type, attack_category, target_model, dataset,
                         prompt, success, safety_score, risk_level, response_text,
                         response_length, inference_time_ms, timestamp, metadata)
                        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """, (
                        data['attack_id'], data['attack_type'], data['attack_category'],
                        data['target_model'], data['dataset'], data['prompt'],
                        data['success'], data['safety_score'], data['risk_level'],
                        data['response_text'], data['response_length'],
                        data['inference_time_ms'], data['timestamp'], 
                        json.dumps(data['metadata'])
                    ))
                    stored_count += 1
                
                conn.commit()
                
                # Update cache
                self._attack_cache.extend(attack_records)
                
                logger.info(f"Stored {stored_count} attack records")
                return stored_count
                
        except Exception as e:
            logger.error(f"Failed to store batch attacks: {e}")
            return 0
    
    def get_attacks_by_model(self, model_name: str, limit: Optional[int] = None) -> List[AttackRecord]:
        """
        Retrieve attacks for a specific model.
        
        Args:
            model_name: Target model name
            limit: Maximum number of records to retrieve
            
        Returns:
            List of attack records
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                query = "SELECT * FROM attacks WHERE target_model = ? ORDER BY timestamp DESC"
                params = [model_name]
                
                if limit:
                    query += " LIMIT ?"
                    params.append(limit)
                
                cursor.execute(query, params)
                rows = cursor.fetchall()
                
                # Convert to AttackRecord objects
                attacks = []
                for row in rows:
                    # Convert row to dict
                    columns = [desc[0] for desc in cursor.description]
                    data = dict(zip(columns, row))
                    data['metadata'] = json.loads(data['metadata']) if data['metadata'] else {}
                    attacks.append(AttackRecord.from_dict(data))
                
                return attacks
                
        except Exception as e:
            logger.error(f"Failed to retrieve attacks for model {model_name}: {e}")
            return []
    
    def get_attacks_by_type(self, attack_type: str, limit: Optional[int] = None) -> List[AttackRecord]:
        """
        Retrieve attacks of a specific type.
        
        Args:
            attack_type: Attack type to retrieve
            limit: Maximum number of records to retrieve
            
        Returns:
            List of attack records
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                query = "SELECT * FROM attacks WHERE attack_type = ? ORDER BY timestamp DESC"
                params = [attack_type]
                
                if limit:
                    query += " LIMIT ?"
                    params.append(limit)
                
                cursor.execute(query, params)
                rows = cursor.fetchall()
                
                # Convert to AttackRecord objects
                attacks = []
                for row in rows:
                    columns = [desc[0] for desc in cursor.description]
                    data = dict(zip(columns, row))
                    data['metadata'] = json.loads(data['metadata']) if data['metadata'] else {}
                    attacks.append(AttackRecord.from_dict(data))
                
                return attacks
                
        except Exception as e:
            logger.error(f"Failed to retrieve attacks for type {attack_type}: {e}")
            return []
    
    def get_recent_attacks(self, hours: int = 24, limit: Optional[int] = None) -> List[AttackRecord]:
        """
        Retrieve recent attacks within specified time window.
        
        Args:
            hours: Number of hours to look back
            limit: Maximum number of records to retrieve
            
        Returns:
            List of recent attack records
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                cutoff_time = (datetime.now() - timedelta(hours=hours)).isoformat()
                
                query = "SELECT * FROM attacks WHERE timestamp >= ? ORDER BY timestamp DESC"
                params = [cutoff_time]
                
                if limit:
                    query += " LIMIT ?"
                    params.append(limit)
                
                cursor.execute(query, params)
                rows = cursor.fetchall()
                
                # Convert to AttackRecord objects
                attacks = []
                for row in rows:
                    columns = [desc[0] for desc in cursor.description]
                    data = dict(zip(columns, row))
                    data['metadata'] = json.loads(data['metadata']) if data['metadata'] else {}
                    attacks.append(AttackRecord.from_dict(data))
                
                return attacks
                
        except Exception as e:
            logger.error(f"Failed to retrieve recent attacks: {e}")
            return []
    
    def get_attack_statistics(self, model_name: Optional[str] = None) -> Dict[str, Any]:
        """
        Get comprehensive attack statistics.
        
        Args:
            model_name: Optional model filter
            
        Returns:
            Dictionary with attack statistics
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                # Base query
                where_clause = "WHERE target_model = ?" if model_name else ""
                params = [model_name] if model_name else []
                
                # Total attacks
                cursor.execute(f"SELECT COUNT(*) FROM attacks {where_clause}", params)
                total_attacks = cursor.fetchone()[0]
                
                # Success rate
                cursor.execute(f"""
                    SELECT 
                        SUM(CASE WHEN success = 1 THEN 1 ELSE 0 END) as successes,
                        AVG(safety_score) as avg_safety,
                        AVG(response_length) as avg_length,
                        AVG(inference_time_ms) as avg_time
                    FROM attacks {where_clause}
                """, params)
                success_data = cursor.fetchone()
                
                # Attack type breakdown
                cursor.execute(f"""
                    SELECT attack_type, COUNT(*) as count, 
                           SUM(CASE WHEN success = 1 THEN 1 ELSE 0 END) as successes
                    FROM attacks {where_clause}
                    GROUP BY attack_type
                    ORDER BY count DESC
                """, params)
                type_breakdown = cursor.fetchall()
                
                # Recent activity
                recent_cutoff = (datetime.now() - timedelta(hours=24)).isoformat()
                cursor.execute(f"""
                    SELECT COUNT(*) FROM attacks 
                    WHERE timestamp >= ? {('AND target_model = ?' if model_name else '')}
                """, [recent_cutoff] + ([model_name] if model_name else []))
                recent_count = cursor.fetchone()[0]
                
                return {
                    "total_attacks": total_attacks,
                    "successful_attacks": success_data[0] if success_data else 0,
                    "success_rate": (success_data[0] / total_attacks * 100) if total_attacks > 0 else 0,
                    "avg_safety_score": success_data[1] or 0,
                    "avg_response_length": success_data[2] or 0,
                    "avg_inference_time_ms": success_data[3] or 0,
                    "attack_type_breakdown": [
                        {
                            "type": row[0],
                            "count": row[1],
                            "successes": row[2],
                            "success_rate": (row[2] / row[1] * 100) if row[1] > 0 else 0
                        }
                        for row in type_breakdown
                    ],
                    "recent_attacks_24h": recent_count,
                    "model_filter": model_name
                }
                
        except Exception as e:
            logger.error(f"Failed to get attack statistics: {e}")
            return {}
    
    def update_pattern_metrics(self, metrics: PatternMetrics) -> bool:
        """
        Update pattern metrics for an attack type.
        
        Args:
            metrics: Pattern metrics to store
            
        Returns:
            True if updated successfully
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                cursor.execute("""
                    INSERT OR REPLACE INTO pattern_metrics 
                    (attack_type, total_attempts, successful_attempts, success_rate,
                     avg_safety_score, avg_response_length, avg_inference_time,
                     last_updated, weak_indicators, strong_indicators)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """, (
                    metrics.attack_type,
                    metrics.total_attempts,
                    metrics.successful_attempts,
                    metrics.success_rate,
                    metrics.avg_safety_score,
                    metrics.avg_response_length,
                    metrics.avg_inference_time,
                    metrics.last_updated.isoformat(),
                    json.dumps(metrics.weak_indicators),
                    json.dumps(metrics.strong_indicators)
                ))
                
                conn.commit()
                
                # Update cache
                self._pattern_cache[metrics.attack_type] = metrics
                
                logger.debug(f"Updated pattern metrics for {metrics.attack_type}")
                return True
                
        except Exception as e:
            logger.error(f"Failed to update pattern metrics: {e}")
            return False
    
    def get_pattern_metrics(self, attack_type: Optional[str] = None) -> Union[PatternMetrics, Dict[str, PatternMetrics]]:
        """
        Retrieve pattern metrics.
        
        Args:
            attack_type: Specific attack type, or None for all
            
        Returns:
            PatternMetrics object or dictionary of all metrics
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                if attack_type:
                    cursor.execute("SELECT * FROM pattern_metrics WHERE attack_type = ?", [attack_type])
                    row = cursor.fetchone()
                    
                    if row:
                        columns = [desc[0] for desc in cursor.description]
                        data = dict(zip(columns, row))
                        return PatternMetrics(
                            attack_type=data['attack_type'],
                            total_attempts=data['total_attempts'],
                            successful_attempts=data['successful_attempts'],
                            success_rate=data['success_rate'],
                            avg_safety_score=data['avg_safety_score'],
                            avg_response_length=data['avg_response_length'],
                            avg_inference_time=data['avg_inference_time'],
                            last_updated=datetime.fromisoformat(data['last_updated']),
                            weak_indicators=json.loads(data['weak_indicators']) if data['weak_indicators'] else [],
                            strong_indicators=json.loads(data['strong_indicators']) if data['strong_indicators'] else []
                        )
                    else:
                        return None
                else:
                    cursor.execute("SELECT * FROM pattern_metrics ORDER BY success_rate DESC")
                    rows = cursor.fetchall()
                    
                    metrics_dict = {}
                    for row in rows:
                        columns = [desc[0] for desc in cursor.description]
                        data = dict(zip(columns, row))
                        metrics_dict[data['attack_type']] = PatternMetrics(
                            attack_type=data['attack_type'],
                            total_attempts=data['total_attempts'],
                            successful_attempts=data['successful_attempts'],
                            success_rate=data['success_rate'],
                            avg_safety_score=data['avg_safety_score'],
                            avg_response_length=data['avg_response_length'],
                            avg_inference_time=data['avg_inference_time'],
                            last_updated=datetime.fromisoformat(data['last_updated']),
                            weak_indicators=json.loads(data['weak_indicators']) if data['weak_indicators'] else [],
                            strong_indicators=json.loads(data['strong_indicators']) if data['strong_indicators'] else []
                        )
                    
                    return metrics_dict
                    
        except Exception as e:
            logger.error(f"Failed to retrieve pattern metrics: {e}")
            return {} if not attack_type else None
    
    def cleanup_old_data(self, days_to_keep: int = 30) -> int:
        """
        Clean up old attack data to manage storage.
        
        Args:
            days_to_keep: Number of days to keep data (0 = no cleanup)
            
        Returns:
            Number of records removed
        """
        try:
            # Check if cleanup is disabled
            if days_to_keep <= 0:
                logger.info("Data cleanup is disabled (LEARNING_CLEANUP_DAYS=0)")
                return 0
            
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                cutoff_time = (datetime.now() - timedelta(days=days_to_keep)).isoformat()
                
                cursor.execute("DELETE FROM attacks WHERE timestamp < ?", [cutoff_time])
                removed_count = cursor.rowcount
                
                conn.commit()
                
                # Clear cache
                self._attack_cache = [r for r in self._attack_cache if r.timestamp >= datetime.fromisoformat(cutoff_time)]
                
                logger.info(f"Cleaned up {removed_count} old attack records")
                return removed_count
                
        except Exception as e:
            logger.error(f"Failed to cleanup old data: {e}")
            return 0
    
    def export_data(self, output_path: str, model_name: Optional[str] = None) -> bool:
        """
        Export attack data to JSON file.
        
        Args:
            output_path: Output file path
            model_name: Optional model filter
            
        Returns:
            True if exported successfully
        """
        try:
            attacks = self.get_attacks_by_model(model_name) if model_name else self.get_recent_attacks(hours=24*30)  # 30 days
            
            export_data = {
                "export_timestamp": datetime.now().isoformat(),
                "model_filter": model_name,
                "total_records": len(attacks),
                "attacks": [attack.to_dict() for attack in attacks]
            }
            
            with open(output_path, 'w') as f:
                json.dump(export_data, f, indent=2)
            
            logger.info(f"Exported {len(attacks)} attack records to {output_path}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to export data: {e}")
            return False


# Singleton instance for global access
_memory_store_instance: Optional[MemoryStore] = None


def get_memory_store(db_path: Optional[str] = None, enable_persistence: bool = True, db_session=None):
    """
    Get memory store instance.
    
    Args:
        db_path: Path to SQLite database (deprecated, use PostgreSQL instead)
        enable_persistence: Whether to enable persistence
        db_session: PostgreSQL database session (preferred)
    
    Returns:
        MemoryStore instance
    """
    global _memory_store_instance
    
    # If database session is provided, reset global instance to ensure PostgreSQL usage
    if db_session is not None:
        _memory_store_instance = None
    
    if _memory_store_instance is None:
        # Use PostgreSQL if session is provided, otherwise fallback to SQLite
        if db_session is not None:
            from .postgresql_memory_store import PostgreSQLMemoryStore
            _memory_store_instance = PostgreSQLMemoryStore(db_session, enable_persistence)
        else:
            logger.warning("⚠️ No database session provided, using SQLite (not recommended for production)")
            _memory_store_instance = MemoryStore(db_path, enable_persistence)
    
    return _memory_store_instance