"""
Correlation Engine Module
"""

from pathlib import Path
from typing import List, Dict, Any
from collections import defaultdict

from ..utils.logger import log_progress
from ..utils.fs import save_json

class CorrelationEngine:
    """Correlate findings across different analysis modules"""
    
    def __init__(self, output_dir: Path):
        self.output_dir = output_dir
    
    def correlate_findings(self, analysis_results: Dict[str, Any]) -> Dict[str, Any]:
        """Correlate findings from different analysis modules"""
        log_progress("Correlating intelligence across modules")
        
        correlated = {
            'high_confidence_endpoints': [],
            'correlated_secrets': [],
            'critical_sinks': [],
            'file_risk_scores': {},
            'summary': {}
        }
        
        # Extract data from analysis results
        endpoints = analysis_results.get('endpoints', [])
        secrets = analysis_results.get('secrets', [])
        sinks = analysis_results.get('sinks', [])
        
        # Correlate endpoints found by multiple methods
        endpoint_sources = defaultdict(list)
        
        for endpoint in endpoints:
            url = endpoint.get('url', endpoint.get('match', ''))
            if url:
                endpoint_sources[url].append(endpoint)
        
        # High confidence endpoints (found by multiple sources)
        for url, sources in endpoint_sources.items():
            if len(sources) > 1:
                confidence_score = len(sources) * 10
                
                correlated['high_confidence_endpoints'].append({
                    'url': url,
                    'sources': [s.get('tool', s.get('pattern_name', 'unknown')) for s in sources],
                    'confidence_score': confidence_score,
                    'source_count': len(sources)
                })
        
        # Correlate secrets with high-risk patterns
        for secret in secrets:
            risk_score = self.calculate_secret_risk(secret)
            
            if risk_score >= 7:  # High risk threshold
                correlated['correlated_secrets'].append({
                    **secret,
                    'risk_score': risk_score,
                    'risk_level': 'high' if risk_score >= 8 else 'medium'
                })
        
        # Identify critical sinks
        for sink in sinks:
            if sink.get('risk') in ['critical', 'high']:
                correlated['critical_sinks'].append(sink)
        
        # Calculate file risk scores
        file_risks = defaultdict(int)
        
        for endpoint in endpoints:
            filename = endpoint.get('filename', '')
            if filename:
                file_risks[filename] += 2
        
        for secret in secrets:
            filename = secret.get('filename', secret.get('source_file', ''))
            if filename:
                file_risks[filename] += 5
        
        for sink in sinks:
            filename = sink.get('filename', '')
            if filename:
                risk_multiplier = {'critical': 8, 'high': 5, 'medium': 3, 'low': 1}
                file_risks[filename] += risk_multiplier.get(sink.get('risk', 'low'), 1)
        
        correlated['file_risk_scores'] = dict(file_risks)
        
        # Generate summary
        correlated['summary'] = {
            'total_endpoints': len(endpoints),
            'high_confidence_endpoints': len(correlated['high_confidence_endpoints']),
            'total_secrets': len(secrets),
            'high_risk_secrets': len(correlated['correlated_secrets']),
            'total_sinks': len(sinks),
            'critical_sinks': len(correlated['critical_sinks']),
            'highest_risk_files': sorted(
                file_risks.items(), 
                key=lambda x: x[1], 
                reverse=True
            )[:10]
        }
        
        # Save correlated results
        save_json(correlated, self.output_dir / "correlation_results.json")
        
        log_progress(f"Correlation complete: {len(correlated['high_confidence_endpoints'])} high-confidence findings")
        
        return correlated
    
    def calculate_secret_risk(self, secret: Dict) -> int:
        """Calculate risk score for a secret (0-10)"""
        risk_score = 5  # Base score
        
        secret_type = secret.get('type', '').lower()
        secret_value = secret.get('value', '').lower()
        
        # High-risk secret types
        if any(keyword in secret_type for keyword in ['api_key', 'secret_key', 'private_key', 'aws']):
            risk_score += 3
        
        # Medium-risk types
        elif any(keyword in secret_type for keyword in ['token', 'password', 'auth']):
            risk_score += 2
        
        # Check value characteristics
        if len(secret_value) > 30:  # Long secrets are often real
            risk_score += 1
        
        if any(char.isdigit() for char in secret_value) and any(char.isalpha() for char in secret_value):
            risk_score += 1  # Mixed alphanumeric
        
        # Reduce score for common false positives
        if any(fp in secret_value for fp in ['example', 'test', 'demo', 'placeholder']):
            risk_score -= 2
        
        return max(0, min(10, risk_score))  # Clamp to 0-10 range