AI Log Analysis in 2026: A Beginner Guide to Smart Threat...

#!/usr/bin/env python3
"""
AI-Powered Log Analysis System
Production-ready ML-based log analysis with comprehensive feature engineering and multiple models
"""

from typing import List, Dict, Optional, Tuple
import pandas as pd
import numpy as np
from sklearn.ensemble import IsolationForest, RandomForestClassifier
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix
import logging
import json
from datetime import datetime, timedelta
from collections import Counter
import pickle
from pathlib import Path

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class LogFeatureExtractor:
    """Extracts features from log data for ML models."""
    
    def __init__(self):
        self.label_encoders: Dict[str, LabelEncoder] = {}
    
    def extract_statistical_features(self, logs: pd.DataFrame) -> np.ndarray:
        """Extract statistical features from logs.
        
        Args:
            logs: DataFrame with log entries
            
        Returns:
            Array of statistical features
        """
        features = []
        
        # Time-based features
        if 'timestamp' in logs.columns:
            logs['timestamp'] = pd.to_datetime(logs['timestamp'])
            logs['hour'] = logs['timestamp'].dt.hour
            logs['day_of_week'] = logs['timestamp'].dt.dayofweek
            
            # Log frequency per hour
            hourly_counts = logs.groupby('hour').size()
            features.extend([
                hourly_counts.mean() if len(hourly_counts) > 0 else 0,
                hourly_counts.std() if len(hourly_counts) > 0 else 0,
                hourly_counts.max() if len(hourly_counts) > 0 else 0
            ])
        
        # IP-based features
        if 'source_ip' in logs.columns:
            unique_ips = logs['source_ip'].nunique()
            ip_counts = logs['source_ip'].value_counts()
            features.extend([
                unique_ips,
                ip_counts.max() if len(ip_counts) > 0 else 0,
                ip_counts.std() if len(ip_counts) > 0 else 0,
                (ip_counts > 100).sum()  # IPs with > 100 requests
            ])
        
        # Error rate features
        if 'level' in logs.columns:
            error_count = (logs['level'] == 'ERROR').sum()
            warning_count = (logs['level'] == 'WARNING').sum()
            total_count = len(logs)
            
            features.extend([
                error_count / total_count if total_count > 0 else 0,
                warning_count / total_count if total_count > 0 else 0,
                error_count,
                warning_count
            ])
        
        # Path/endpoint features
        if 'path' in logs.columns:
            path_counts = logs['path'].value_counts()
            unique_paths = logs['path'].nunique()
            
            features.extend([
                unique_paths,
                path_counts.max() if len(path_counts) > 0 else 0,
                path_counts.std() if len(path_counts) > 0 else 0,
                (path_counts > 50).sum()  # Popular paths
            ])
            
            # Top path frequencies
            top_paths = path_counts.head(5).values
            features.extend(list(top_paths) + [0] * (5 - len(top_paths)))
        
        # Status code features (if available)
        if 'status_code' in logs.columns:
            status_counts = logs['status_code'].value_counts()
            features.extend([
                (status_counts.get(200, 0) / len(logs)) if len(logs) > 0 else 0,
                (status_counts.get(404, 0) / len(logs)) if len(logs) > 0 else 0,
                (status_counts.get(500, 0) / len(logs)) if len(logs) > 0 else 0,
                (status_counts >= 400).sum()  # Error status codes
            ])
        
        # User agent features (if available)
        if 'user_agent' in logs.columns:
            unique_agents = logs['user_agent'].nunique()
            features.append(unique_agents)
        
        # Request size features (if available)
        if 'request_size' in logs.columns:
            features.extend([
                logs['request_size'].mean(),
                logs['request_size'].std(),
                logs['request_size'].max()
            ])
        
        return np.array(features, dtype=np.float32)
    
    def extract_behavioral_features(self, logs: pd.DataFrame, window_minutes: int = 60) -> np.ndarray:
        """Extract behavioral features from logs.
        
        Args:
            logs: DataFrame with log entries
            window_minutes: Time window for behavioral analysis
            
        Returns:
            Array of behavioral features
        """
        features = []
        
        if 'timestamp' not in logs.columns or 'source_ip' not in logs.columns:
            return np.array([])
        
        logs['timestamp'] = pd.to_datetime(logs['timestamp'])
        logs = logs.sort_values('timestamp')
        
        # Calculate request rate per IP
        ip_request_rates = []
        for ip in logs['source_ip'].unique()[:100]:  # Limit to top 100 IPs
            ip_logs = logs[logs['source_ip'] == ip]
            if len(ip_logs) > 0:
                time_span = (ip_logs['timestamp'].max() - ip_logs['timestamp'].min()).total_seconds() / 60
                rate = len(ip_logs) / max(time_span, 1)
                ip_request_rates.append(rate)
        
        if ip_request_rates:
            features.extend([
                np.mean(ip_request_rates),
                np.std(ip_request_rates),
                np.max(ip_request_rates),
                np.percentile(ip_request_rates, 95)
            ])
        else:
            features.extend([0, 0, 0, 0])
        
        # Failed login attempts (if applicable)
        if 'message' in logs.columns:
            failed_logins = logs['message'].str.contains('failed|denied|unauthorized', case=False, na=False).sum()
            features.append(failed_logins)
        
        return np.array(features, dtype=np.float32)
    
    def extract_features(self, logs: pd.DataFrame) -> np.ndarray:
        """Extract all features from logs.
        
        Args:
            logs: DataFrame with log entries
            
        Returns:
            Combined feature array
        """
        statistical = self.extract_statistical_features(logs)
        behavioral = self.extract_behavioral_features(logs)
        
        # Combine features
        all_features = np.concatenate([statistical, behavioral])
        
        # Handle NaN values
        all_features = np.nan_to_num(all_features, nan=0.0, posinf=0.0, neginf=0.0)
        
        return all_features

class AILogAnalyzer:
    """Production-ready AI-powered log analysis system."""
    
    def __init__(self, model_type: str = "isolation_forest"):
        """Initialize AI log analyzer.
        
        Args:
            model_type: Type of model to use ("isolation_forest" or "random_forest")
        """
        self.model_type = model_type
        self.feature_extractor = LogFeatureExtractor()
        self.scaler = StandardScaler()
        self.is_trained = False
        self.model = None
        self.feature_names: List[str] = []
        self._initialize_model()
    
    def _initialize_model(self):
        """Initialize ML model based on type."""
        if self.model_type == "isolation_forest":
            self.model = IsolationForest(
                contamination=0.1,
                random_state=42,
                n_estimators=100,
                max_samples='auto'
            )
        elif self.model_type == "random_forest":
            self.model = RandomForestClassifier(
                n_estimators=100,
                random_state=42,
                max_depth=10
            )
        else:
            raise ValueError(f"Unknown model type: {self.model_type}")
    
    def prepare_training_data(self, log_files: List[str]) -> Tuple[np.ndarray, Optional[np.ndarray]]:
        """Prepare training data from log files.
        
        Args:
            log_files: List of paths to log files
            
        Returns:
            Tuple of (features, labels) or (features, None) for unsupervised
        """
        all_features = []
        all_labels = []
        
        for log_file in log_files:
            try:
                logs = pd.read_csv(log_file)
                features = self.feature_extractor.extract_features(logs)
                all_features.append(features)
                
                # For supervised learning, extract labels if available
                if 'label' in logs.columns:
                    label = 1 if (logs['label'] == 'anomaly').any() else 0
                    all_labels.append(label)
                    
            except Exception as e:
                logger.error(f"Failed to process {log_file}: {e}")
                continue
        
        features_array = np.array(all_features)
        
        if all_labels:
            labels_array = np.array(all_labels)
            return features_array, labels_array
        else:
            return features_array, None
    
    def train(self, normal_logs: pd.DataFrame, anomaly_logs: Optional[pd.DataFrame] = None):
        """Train model on log data.
        
        Args:
            normal_logs: DataFrame with normal log entries
            anomaly_logs: Optional DataFrame with known anomaly entries
        """
        try:
            logger.info("Extracting features from training data...")
            
            # Extract features
            normal_features = self.feature_extractor.extract_features(normal_logs)
            
            if anomaly_logs is not None and len(anomaly_logs) > 0:
                anomaly_features = self.feature_extractor.extract_features(anomaly_logs)
                
                # Combine features
                X = np.vstack([normal_features.reshape(1, -1), anomaly_features.reshape(1, -1)])
                y = np.array([0] * len(normal_features.reshape(1, -1)) + [1] * len(anomaly_features.reshape(1, -1)))
                
                # For supervised learning
                if self.model_type == "random_forest":
                    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
                    X_scaled = self.scaler.fit_transform(X_train)
                    self.model.fit(X_scaled, y_train)
                    
                    # Evaluate
                    X_test_scaled = self.scaler.transform(X_test)
                    predictions = self.model.predict(X_test_scaled)
                    logger.info(f"Model accuracy: {np.mean(predictions == y_test):.2%}")
                else:
                    # Unsupervised learning
                    X_scaled = self.scaler.fit_transform(X)
                    self.model.fit(X_scaled)
            else:
                # Unsupervised learning on normal data only
                X = normal_features.reshape(1, -1)
                X_scaled = self.scaler.fit_transform(X)
                self.model.fit(X_scaled)
            
            self.is_trained = True
            logger.info("Model trained successfully")
            
        except Exception as e:
            logger.error(f"Training failed: {e}", exc_info=True)
            raise
    
    def analyze(self, logs: pd.DataFrame) -> Dict:
        """Analyze logs for anomalies and threats.
        
        Args:
            logs: DataFrame with log entries to analyze
            
        Returns:
            Analysis result dictionary
        """
        if not self.is_trained:
            raise ValueError("Model not trained. Call train() first.")
        
        try:
            # Extract features
            features = self.feature_extractor.extract_features(logs)
            features_scaled = self.scaler.transform(features.reshape(1, -1))
            
            if self.model_type == "isolation_forest":
                # Isolation Forest prediction
        prediction = self.model.predict(features_scaled)[0]
        score = self.model.score_samples(features_scaled)[0]
        
        is_anomaly = prediction == -1
                anomaly_score = abs(score)
                confidence = min(1.0, anomaly_score / 10.0)  # Normalize confidence
                
            else:
                # Random Forest prediction
                prediction = self.model.predict(features_scaled)[0]
                probabilities = self.model.predict_proba(features_scaled)[0]
                
                is_anomaly = prediction == 1
                confidence = probabilities[1] if len(probabilities) > 1 else 0.0
                anomaly_score = confidence
            
            # Determine threat level
            if is_anomaly:
                if confidence > 0.8:
                    threat_level = "high"
                    recommendation = "immediate_investigation"
                elif confidence > 0.6:
                    threat_level = "medium"
                    recommendation = "priority_investigation"
                else:
                    threat_level = "low"
                    recommendation = "review"
            else:
                threat_level = "normal"
                recommendation = "no_action"
            
            result = {
                'is_anomaly': is_anomaly,
                'threat_level': threat_level,
                'confidence': float(confidence),
                'anomaly_score': float(anomaly_score),
                'recommendation': recommendation,
                'features': features.tolist(),
                'analysis_timestamp': datetime.now().isoformat(),
                'log_count': len(logs)
            }
            
            logger.info(f"Analysis complete: anomaly={is_anomaly}, confidence={confidence:.2f}, threat={threat_level}")
            return result
            
        except Exception as e:
            logger.error(f"Analysis failed: {e}", exc_info=True)
            return {
                'is_anomaly': False,
                'error': str(e),
                'threat_level': 'unknown'
            }
    
    def batch_analyze(self, log_files: List[str]) -> List[Dict]:
        """Analyze multiple log files in batch.
        
        Args:
            log_files: List of log file paths
            
        Returns:
            List of analysis results
        """
        results = []
        
        for log_file in log_files:
            try:
                logs = pd.read_csv(log_file)
                result = self.analyze(logs)
                result['log_file'] = log_file
                results.append(result)
            except Exception as e:
                logger.error(f"Failed to analyze {log_file}: {e}")
                results.append({
                    'log_file': log_file,
                    'error': str(e),
                    'is_anomaly': False
                })
        
        return results
    
    def save_model(self, filepath: str):
        """Save trained model to file.
        
        Args:
            filepath: Path to save model
        """
        try:
            model_data = {
                'model': self.model,
                'scaler': self.scaler,
                'model_type': self.model_type,
                'is_trained': self.is_trained,
                'feature_extractor': self.feature_extractor
            }
            
            with open(filepath, 'wb') as f:
                pickle.dump(model_data, f)
            
            logger.info(f"Model saved to {filepath}")
        except Exception as e:
            logger.error(f"Failed to save model: {e}")
            raise
    
    def load_model(self, filepath: str):
        """Load trained model from file.
        
        Args:
            filepath: Path to model file
        """
        try:
            with open(filepath, 'rb') as f:
                model_data = pickle.load(f)
            
            self.model = model_data['model']
            self.scaler = model_data['scaler']
            self.model_type = model_data['model_type']
            self.is_trained = model_data['is_trained']
            self.feature_extractor = model_data.get('feature_extractor', LogFeatureExtractor())
            
            logger.info(f"Model loaded from {filepath}")
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            raise
    
    def get_model_info(self) -> Dict:
        """Get information about the trained model.
        
        Returns:
            Model information dictionary
        """
        return {
            'model_type': self.model_type,
            'is_trained': self.is_trained,
            'feature_count': self.scaler.n_features_in_ if hasattr(self.scaler, 'n_features_in_') else 0,
            'model_params': self.model.get_params() if hasattr(self.model, 'get_params') else {}
        }


# Example usage
if __name__ == "__main__":
    # Initialize analyzer
    analyzer = AILogAnalyzer(model_type="isolation_forest")
    
    # Create sample training data
    normal_logs = pd.DataFrame({
        'timestamp': pd.date_range('2024-01-01', periods=1000, freq='1min'),
        'source_ip': np.random.choice(['192.168.1.1', '192.168.1.2', '10.0.0.1'], 1000),
        'level': np.random.choice(['INFO', 'DEBUG'], 1000),
        'path': np.random.choice(['/api/users', '/api/data', '/health'], 1000),
        'status_code': np.random.choice([200, 201], 1000)
    })
    
    # Train model
analyzer.train(normal_logs)

# Analyze new logs
    new_logs = pd.DataFrame({
        'timestamp': pd.date_range('2024-01-02', periods=100, freq='1min'),
        'source_ip': ['192.168.1.100'] * 100,  # Suspicious: single IP
        'level': ['ERROR'] * 100,  # Suspicious: all errors
        'path': ['/api/admin'] * 100,  # Suspicious: admin endpoint
        'status_code': [403] * 100  # Suspicious: forbidden
    })
    
result = analyzer.analyze(new_logs)
    print(f"\nAnalysis Result:")
    print(json.dumps(result, indent=2))
    
    # Save model
    analyzer.save_model('ai_log_analyzer_model.pkl')
    
    # Load and use saved model
    new_analyzer = AILogAnalyzer()
    new_analyzer.load_model('ai_log_analyzer_model.pkl')
    result2 = new_analyzer.analyze(new_logs)
    print(f"\nLoaded Model Analysis:")
    print(json.dumps(result2, indent=2))

#!/usr/bin/env python3
"""
Pattern Recognition for Log Analysis
Detects attack patterns and known threat signatures
"""

from typing import List, Dict, Set
import re
import pandas as pd
from collections import Counter
from datetime import datetime
import logging

logger = logging.getLogger(__name__)

class AttackPattern:
    """Represents an attack pattern."""
    
    def __init__(self, name: str, pattern: str, severity: str, description: str):
        self.name = name
        self.pattern = re.compile(pattern, re.IGNORECASE)
        self.severity = severity
        self.description = description
    
    def match(self, text: str) -> bool:
        """Check if pattern matches text."""
        return bool(self.pattern.search(text))

class PatternRecognizer:
    """Recognizes attack patterns in logs."""
    
    def __init__(self):
        self.patterns: List[AttackPattern] = []
        self._load_default_patterns()
    
    def _load_default_patterns(self):
        """Load default attack patterns."""
        default_patterns = [
            AttackPattern(
                name="SQL Injection",
                pattern=r"(union\s+select|';?\s*(drop|delete|insert|update)|exec\s*\(|xp_cmdshell)",
                severity="high",
                description="Potential SQL injection attempt"
            ),
            AttackPattern(
                name="XSS Attack",
                pattern=r"(<script|javascript:|onerror=|onload=)",
                severity="high",
                description="Potential cross-site scripting attack"
            ),
            AttackPattern(
                name="Path Traversal",
                pattern=r"(\.\./|\.\.\\|%2e%2e%2f)",
                severity="medium",
                description="Potential path traversal attempt"
            ),
            AttackPattern(
                name="Command Injection",
                pattern=r"(;|\||&|`|\$\(|\${).*(ls|cat|rm|wget|curl|nc|bash|sh)",
                severity="high",
                description="Potential command injection attempt"
            ),
            AttackPattern(
                name="Brute Force",
                pattern=r"(failed|denied|unauthorized|invalid).*(login|password|authentication)",
                severity="medium",
                description="Potential brute force attack"
            ),
            AttackPattern(
                name="Suspicious User Agent",
                pattern=r"(sqlmap|nikto|nmap|masscan|zap|burp)",
                severity="medium",
                description="Suspicious security tool user agent"
            )
        ]
        
        self.patterns.extend(default_patterns)
        logger.info(f"Loaded {len(self.patterns)} attack patterns")
    
    def add_pattern(self, pattern: AttackPattern):
        """Add custom attack pattern."""
        self.patterns.append(pattern)
        logger.info(f"Added pattern: {pattern.name}")
    
    def detect_patterns(self, logs: pd.DataFrame) -> List[Dict]:
        """Detect attack patterns in logs.
        
        Args:
            logs: DataFrame with log entries
            
        Returns:
            List of detected patterns with details
        """
        detections = []
        
        # Check message/content fields
        text_fields = ['message', 'content', 'request', 'url', 'user_agent']
        available_fields = [f for f in text_fields if f in logs.columns]
        
        for idx, row in logs.iterrows():
            for field in available_fields:
                text = str(row[field]) if pd.notna(row[field]) else ""
                
                for pattern in self.patterns:
                    if pattern.match(text):
                        detection = {
                            'pattern_name': pattern.name,
                            'severity': pattern.severity,
                            'description': pattern.description,
                            'matched_text': text[:100],  # First 100 chars
                            'log_index': idx,
                            'field': field,
                            'timestamp': row.get('timestamp', datetime.now())
                        }
                        detections.append(detection)
                        logger.warning(f"Pattern detected: {pattern.name} in log {idx}")
        
        return detections
    
    def get_pattern_statistics(self, detections: List[Dict]) -> Dict:
        """Get statistics about detected patterns.
        
        Args:
            detections: List of pattern detections
            
        Returns:
            Statistics dictionary
        """
        if not detections:
            return {'total_detections': 0}
        
        pattern_counts = Counter([d['pattern_name'] for d in detections])
        severity_counts = Counter([d['severity'] for d in detections])
        
        return {
            'total_detections': len(detections),
            'unique_patterns': len(pattern_counts),
            'pattern_counts': dict(pattern_counts),
            'severity_counts': dict(severity_counts),
            'high_severity_count': severity_counts.get('high', 0),
            'medium_severity_count': severity_counts.get('medium', 0)
        }


# Example usage
if __name__ == "__main__":
    recognizer = PatternRecognizer()
    
    # Sample logs
    logs = pd.DataFrame({
        'timestamp': [datetime.now()] * 3,
        'message': [
            "SELECT * FROM users WHERE id=1 UNION SELECT password",
            "<script>alert('XSS')</script>",
            "Normal log entry"
        ]
    })
    
    detections = recognizer.detect_patterns(logs)
    print(f"Detected {len(detections)} attack patterns")
    
    for detection in detections:
        print(f"- {detection['pattern_name']}: {detection['description']}")
    
    stats = recognizer.get_pattern_statistics(detections)
    print(f"\nStatistics: {stats}")

#!/usr/bin/env python3
"""
Unit tests for AI Log Analyzer
Comprehensive test coverage with pytest
"""

import pytest
import pandas as pd
import numpy as np
from datetime import datetime
from ai_log_analyzer import AILogAnalyzer, LogFeatureExtractor
from pattern_recognition import PatternRecognizer, AttackPattern

class TestLogFeatureExtractor:
    """Tests for LogFeatureExtractor."""
    
    @pytest.fixture
    def extractor(self):
        return LogFeatureExtractor()
    
    @pytest.fixture
    def sample_logs(self):
        return pd.DataFrame({
            'timestamp': pd.date_range('2024-01-01', periods=100, freq='1min'),
            'source_ip': np.random.choice(['192.168.1.1', '192.168.1.2'], 100),
            'level': np.random.choice(['INFO', 'ERROR'], 100),
            'path': ['/api/users'] * 100,
            'status_code': [200] * 100
        })
    
    def test_extract_statistical_features(self, extractor, sample_logs):
        """Test statistical feature extraction."""
        features = extractor.extract_statistical_features(sample_logs)
        assert len(features) > 0
        assert not np.isnan(features).any()
    
    def test_extract_behavioral_features(self, extractor, sample_logs):
        """Test behavioral feature extraction."""
        features = extractor.extract_behavioral_features(sample_logs)
        assert len(features) >= 0
    
    def test_extract_features(self, extractor, sample_logs):
        """Test complete feature extraction."""
        features = extractor.extract_features(sample_logs)
        assert len(features) > 0
        assert not np.isnan(features).any()

class TestAILogAnalyzer:
    """Tests for AILogAnalyzer."""
    
    @pytest.fixture
    def analyzer(self):
        return AILogAnalyzer(model_type="isolation_forest")
    
    @pytest.fixture
    def normal_logs(self):
        return pd.DataFrame({
            'timestamp': pd.date_range('2024-01-01', periods=100, freq='1min'),
            'source_ip': np.random.choice(['192.168.1.1', '192.168.1.2'], 100),
            'level': ['INFO'] * 100,
            'path': ['/api/users'] * 100,
            'status_code': [200] * 100
        })
    
    @pytest.fixture
    def anomaly_logs(self):
        return pd.DataFrame({
            'timestamp': pd.date_range('2024-01-02', periods=50, freq='1min'),
            'source_ip': ['192.168.1.100'] * 50,
            'level': ['ERROR'] * 50,
            'path': ['/api/admin'] * 50,
            'status_code': [403] * 50
        })
    
    def test_analyzer_initialization(self, analyzer):
        """Test analyzer initialization."""
        assert analyzer.model is not None
        assert analyzer.is_trained is False
    
    def test_train_model(self, analyzer, normal_logs):
        """Test model training."""
        analyzer.train(normal_logs)
        assert analyzer.is_trained is True
    
    def test_analyze_logs(self, analyzer, normal_logs, anomaly_logs):
        """Test log analysis."""
        analyzer.train(normal_logs)
        result = analyzer.analyze(anomaly_logs)
        
        assert 'is_anomaly' in result
        assert 'confidence' in result
        assert 'threat_level' in result
        assert result['is_anomaly'] is True  # Anomaly logs should be detected
    
    def test_save_load_model(self, analyzer, normal_logs, tmp_path):
        """Test model save/load."""
        analyzer.train(normal_logs)
        
        model_path = tmp_path / "test_model.pkl"
        analyzer.save_model(str(model_path))
        
        new_analyzer = AILogAnalyzer()
        new_analyzer.load_model(str(model_path))
        
        assert new_analyzer.is_trained is True
        assert new_analyzer.model is not None

class TestPatternRecognizer:
    """Tests for PatternRecognizer."""
    
    @pytest.fixture
    def recognizer(self):
        return PatternRecognizer()
    
    def test_pattern_detection(self, recognizer):
        """Test pattern detection."""
        logs = pd.DataFrame({
            'message': ["SELECT * FROM users UNION SELECT password"],
            'timestamp': [datetime.now()]
        })
        
        detections = recognizer.detect_patterns(logs)
        assert len(detections) > 0
        assert any(d['pattern_name'] == 'SQL Injection' for d in detections)
    
    def test_add_custom_pattern(self, recognizer):
        """Test adding custom pattern."""
        pattern = AttackPattern(
            name="Test Pattern",
            pattern=r"test",
            severity="low",
            description="Test pattern"
        )
        
        recognizer.add_pattern(pattern)
        assert len(recognizer.patterns) > 0


if __name__ == "__main__":
    pytest.main([__file__, "-v"])

#!/usr/bin/env python3
"""
AI Log Analyzer Cleanup
Production-ready cleanup and resource management
"""

import logging
import os
from pathlib import Path
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)

class AILogAnalyzerCleanup:
    """Handles cleanup operations for AI log analyzer."""
    
    def __init__(self, analyzer, model_dir: str = "./models"):
        """Initialize cleanup handler.
        
        Args:
            analyzer: AILogAnalyzer instance
            model_dir: Directory containing model files
        """
        self.analyzer = analyzer
        self.model_dir = Path(model_dir)
    
    def cleanup_old_models(self, days: int = 30):
        """Remove model files older than specified days.
        
        Args:
            days: Number of days to keep models
        """
        cutoff_date = datetime.now() - timedelta(days=days)
        removed_count = 0
        
        if not self.model_dir.exists():
            logger.warning(f"Model directory does not exist: {self.model_dir}")
            return 0
        
        for model_file in self.model_dir.glob("*.pkl"):
            try:
                file_time = datetime.fromtimestamp(model_file.stat().st_mtime)
                if file_time < cutoff_date:
                    model_file.unlink()
                    removed_count += 1
                    logger.info(f"Removed old model: {model_file}")
            except Exception as e:
                logger.error(f"Failed to remove {model_file}: {e}")
        
        logger.info(f"Cleaned up {removed_count} old model files")
        return removed_count
    
    def cleanup_temp_files(self, temp_dir: str = "./temp"):
        """Clean up temporary files.
        
        Args:
            temp_dir: Temporary directory path
        """
        temp_path = Path(temp_dir)
        if not temp_path.exists():
            return 0
        
        removed_count = 0
        for temp_file in temp_path.glob("*"):
            try:
                if temp_file.is_file():
                    temp_file.unlink()
                    removed_count += 1
            except Exception as e:
                logger.error(f"Failed to remove {temp_file}: {e}")
        
        logger.info(f"Cleaned up {removed_count} temporary files")
        return removed_count
    
    def cleanup(self):
        """Perform complete cleanup."""
        logger.info("Starting AI log analyzer cleanup")
        
        # Clean up old models
        self.cleanup_old_models()
        
        # Clean up temporary files
        self.cleanup_temp_files()
        
        # Clean up analyzer resources
        if hasattr(self.analyzer, 'cleanup'):
            self.analyzer.cleanup()
        
        logger.info("AI log analyzer cleanup complete")