"""
MIESC Vulnerability Clusterer
Agrupa vulnerabilidades similares para análisis más eficiente.
"""

import re
import hashlib
from typing import Dict, Any, List, Optional, Tuple, Set
from dataclasses import dataclass, field
from collections import defaultdict
from difflib import SequenceMatcher


@dataclass
class VulnerabilityCluster:
    """Un cluster de vulnerabilidades relacionadas."""
    id: str
    category: str
    root_cause: str
    severity: str
    findings: List[Dict[str, Any]]
    affected_functions: Set[str]
    affected_files: Set[str]
    common_patterns: List[str]
    remediation: str
    confidence: float

    def to_dict(self) -> Dict[str, Any]:
        return {
            'id': self.id,
            'category': self.category,
            'root_cause': self.root_cause,
            'severity': self.severity,
            'finding_count': len(self.findings),
            'affected_functions': list(self.affected_functions),
            'affected_files': list(self.affected_files),
            'common_patterns': self.common_patterns,
            'remediation': self.remediation,
            'confidence': round(self.confidence, 3),
        }


class VulnerabilityClusterer:
    """
    Agrupa vulnerabilidades por:
    1. Categoría semántica (reentrancy, access control, etc.)
    2. Causa raíz común
    3. Ubicación en el código
    4. Patrones de código similares

    Beneficios:
    - Reduce duplicados semánticos
    - Facilita remediación sistemática
    - Mejora comprensión del riesgo global
    """

    # Categorías de vulnerabilidades
    CATEGORIES = {
        'reentrancy': {
            'patterns': ['reentrancy', 'reentrant', 'call.value', 'external-call'],
            'root_cause': "State changes after external calls",
            'remediation': "Use checks-effects-interactions pattern or ReentrancyGuard",
        },
        'access_control': {
            'patterns': ['access-control', 'unprotected', 'missing-authorization', 'tx.origin'],
            'root_cause': "Missing or incorrect access restrictions",
            'remediation': "Implement proper access control modifiers (onlyOwner, role-based)",
        },
        'arithmetic': {
            'patterns': ['overflow', 'underflow', 'integer', 'arithmetic', 'divide-by-zero'],
            'root_cause': "Unsafe arithmetic operations",
            'remediation': "Use SafeMath or Solidity 0.8+ with built-in overflow checks",
        },
        'input_validation': {
            'patterns': ['validation', 'user-input', 'unchecked', 'require'],
            'root_cause': "Missing or insufficient input validation",
            'remediation': "Add require statements to validate all external inputs",
        },
        'denial_of_service': {
            'patterns': ['dos', 'denial', 'gas', 'loop', 'unbounded'],
            'root_cause': "Operations that can exhaust gas or block execution",
            'remediation': "Limit loop iterations, use pull over push patterns",
        },
        'oracle_manipulation': {
            'patterns': ['oracle', 'price', 'manipulation', 'flash-loan'],
            'root_cause': "Reliance on manipulable price oracles",
            'remediation': "Use time-weighted average prices (TWAP) or multiple oracle sources",
        },
        'front_running': {
            'patterns': ['front-run', 'mev', 'sandwich', 'transaction-order'],
            'root_cause': "Transaction ordering dependencies",
            'remediation': "Implement commit-reveal schemes or use private mempools",
        },
        'logic_errors': {
            'patterns': ['logic', 'incorrect', 'wrong', 'bug', 'flaw'],
            'root_cause': "Flawed business logic implementation",
            'remediation': "Review and test business logic thoroughly",
        },
        'code_quality': {
            'patterns': ['naming', 'style', 'convention', 'unused', 'dead-code'],
            'root_cause': "Poor code quality and maintainability",
            'remediation': "Follow Solidity style guide and best practices",
        },
        'upgradability': {
            'patterns': ['upgrade', 'proxy', 'delegatecall', 'storage-collision'],
            'root_cause': "Unsafe upgrade patterns",
            'remediation': "Use established upgrade patterns (UUPS, Transparent Proxy)",
        },
    }

    def __init__(self, similarity_threshold: float = 0.7):
        self.similarity_threshold = similarity_threshold
        self._clusters: List[VulnerabilityCluster] = []

    def _categorize_finding(self, finding: Dict[str, Any]) -> str:
        """Determina la categoría de un hallazgo."""
        vuln_type = finding.get('type', finding.get('check', '')).lower()
        message = finding.get('message', '').lower()
        combined = f"{vuln_type} {message}"

        best_category = 'uncategorized'
        best_score = 0

        for category, info in self.CATEGORIES.items():
            score = sum(
                1 for pattern in info['patterns']
                if pattern in combined
            )
            if score > best_score:
                best_score = score
                best_category = category

        return best_category

    def _extract_patterns(self, findings: List[Dict[str, Any]]) -> List[str]:
        """Extrae patrones comunes de un grupo de hallazgos."""
        patterns = []

        # Extraer tipos únicos
        types = set()
        for f in findings:
            vtype = f.get('type', f.get('check', ''))
            if vtype:
                types.add(vtype)
        patterns.extend(list(types)[:3])

        # Extraer palabras clave comunes de mensajes
        word_counts = defaultdict(int)
        stop_words = {'the', 'a', 'an', 'is', 'are', 'was', 'were', 'be', 'been',
                     'being', 'have', 'has', 'had', 'do', 'does', 'did', 'will',
                     'would', 'could', 'should', 'may', 'might', 'must', 'shall',
                     'can', 'need', 'to', 'of', 'in', 'for', 'on', 'with', 'at',
                     'by', 'from', 'or', 'and', 'not', 'no', 'this', 'that', 'it'}

        for f in findings:
            message = f.get('message', '')
            words = re.findall(r'\b[a-zA-Z]{4,}\b', message.lower())
            for word in words:
                if word not in stop_words:
                    word_counts[word] += 1

        # Top palabras
        top_words = sorted(word_counts.items(), key=lambda x: -x[1])[:5]
        patterns.extend([w for w, _ in top_words])

        return patterns[:8]

    def _compute_similarity(self, f1: Dict[str, Any], f2: Dict[str, Any]) -> float:
        """Calcula similaridad entre dos hallazgos."""
        # Similaridad de tipo
        type1 = f1.get('type', '').lower()
        type2 = f2.get('type', '').lower()
        type_sim = SequenceMatcher(None, type1, type2).ratio()

        # Similaridad de mensaje
        msg1 = f1.get('message', '').lower()
        msg2 = f2.get('message', '').lower()
        msg_sim = SequenceMatcher(None, msg1[:200], msg2[:200]).ratio()

        # Mismo archivo
        file1 = f1.get('location', {}).get('file', '')
        file2 = f2.get('location', {}).get('file', '')
        same_file = 1.0 if file1 == file2 and file1 else 0.0

        # Misma categoría
        cat1 = self._categorize_finding(f1)
        cat2 = self._categorize_finding(f2)
        same_cat = 1.0 if cat1 == cat2 else 0.0

        # Ponderación
        return (type_sim * 0.3 + msg_sim * 0.3 + same_file * 0.2 + same_cat * 0.2)

    def _get_highest_severity(self, findings: List[Dict[str, Any]]) -> str:
        """Obtiene la severidad más alta de un grupo."""
        severity_order = ['critical', 'high', 'medium', 'low', 'informational', 'info']

        highest = 'informational'
        for f in findings:
            sev = f.get('severity', 'medium').lower()
            if sev in severity_order:
                if severity_order.index(sev) < severity_order.index(highest):
                    highest = sev

        return highest

    def cluster(self, findings: List[Dict[str, Any]]) -> List[VulnerabilityCluster]:
        """
        Agrupa hallazgos en clusters.

        Args:
            findings: Lista de hallazgos de múltiples herramientas

        Returns:
            Lista de VulnerabilityCluster
        """
        if not findings:
            return []

        # Primero agrupar por categoría
        category_groups: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
        for finding in findings:
            category = self._categorize_finding(finding)
            category_groups[category].append(finding)

        self._clusters = []

        for category, group_findings in category_groups.items():
            # Sub-agrupar por similaridad dentro de la categoría
            subclusters = self._subcluster(group_findings)

            for subgroup in subclusters:
                if not subgroup:
                    continue

                # Extraer información del cluster
                affected_functions = set()
                affected_files = set()

                for f in subgroup:
                    loc = f.get('location', {})
                    if loc.get('function'):
                        affected_functions.add(loc['function'])
                    if loc.get('file'):
                        affected_files.add(loc['file'])

                # Obtener info de la categoría
                cat_info = self.CATEGORIES.get(category, {
                    'root_cause': 'Unknown root cause',
                    'remediation': 'Manual review required',
                })

                cluster_id = hashlib.md5(
                    f"{category}_{len(self._clusters)}".encode()
                ).hexdigest()[:8]

                cluster = VulnerabilityCluster(
                    id=f"CLUSTER-{cluster_id}",
                    category=category,
                    root_cause=cat_info['root_cause'],
                    severity=self._get_highest_severity(subgroup),
                    findings=subgroup,
                    affected_functions=affected_functions,
                    affected_files=affected_files,
                    common_patterns=self._extract_patterns(subgroup),
                    remediation=cat_info['remediation'],
                    confidence=self._compute_cluster_confidence(subgroup),
                )
                self._clusters.append(cluster)

        # Ordenar por severidad
        severity_order = {'critical': 0, 'high': 1, 'medium': 2, 'low': 3, 'informational': 4}
        self._clusters.sort(key=lambda c: severity_order.get(c.severity, 5))

        return self._clusters

    def _subcluster(self, findings: List[Dict[str, Any]]) -> List[List[Dict[str, Any]]]:
        """Sub-agrupa hallazgos por similaridad."""
        if len(findings) <= 1:
            return [findings]

        # Simple clustering por similaridad
        clusters: List[List[Dict[str, Any]]] = []
        used: Set[int] = set()

        for i, f1 in enumerate(findings):
            if i in used:
                continue

            cluster = [f1]
            used.add(i)

            for j, f2 in enumerate(findings[i+1:], start=i+1):
                if j not in used:
                    if self._compute_similarity(f1, f2) >= self.similarity_threshold:
                        cluster.append(f2)
                        used.add(j)

            clusters.append(cluster)

        return clusters

    def _compute_cluster_confidence(self, findings: List[Dict[str, Any]]) -> float:
        """Calcula confianza del cluster."""
        if not findings:
            return 0.0

        # Más hallazgos = más confianza
        count_factor = min(len(findings) / 5, 1.0) * 0.3

        # Más herramientas diferentes = más confianza
        tools = set(f.get('tool', 'unknown') for f in findings)
        tool_factor = min(len(tools) / 3, 1.0) * 0.4

        # Severidad alta = más confianza
        severity = self._get_highest_severity(findings)
        severity_factor = {'critical': 0.3, 'high': 0.25, 'medium': 0.2, 'low': 0.1}.get(severity, 0.05)

        return min(count_factor + tool_factor + severity_factor, 0.95)

    def get_summary(self) -> Dict[str, Any]:
        """Obtiene resumen de los clusters."""
        if not self._clusters:
            return {'total_clusters': 0}

        category_counts = defaultdict(int)
        severity_counts = defaultdict(int)

        for cluster in self._clusters:
            category_counts[cluster.category] += 1
            severity_counts[cluster.severity] += 1

        return {
            'total_clusters': len(self._clusters),
            'total_findings': sum(len(c.findings) for c in self._clusters),
            'by_category': dict(category_counts),
            'by_severity': dict(severity_counts),
            'clusters': [c.to_dict() for c in self._clusters],
        }

    def get_remediation_plan(self) -> List[Dict[str, Any]]:
        """Genera plan de remediación priorizado."""
        plan = []

        for i, cluster in enumerate(self._clusters):
            plan.append({
                'priority': i + 1,
                'cluster_id': cluster.id,
                'category': cluster.category,
                'severity': cluster.severity,
                'finding_count': len(cluster.findings),
                'affected_files': list(cluster.affected_files)[:5],
                'root_cause': cluster.root_cause,
                'remediation': cluster.remediation,
                'estimated_effort': self._estimate_effort(cluster),
            })

        return plan

    def _estimate_effort(self, cluster: VulnerabilityCluster) -> str:
        """Estima esfuerzo de remediación."""
        file_count = len(cluster.affected_files)
        finding_count = len(cluster.findings)

        if file_count <= 1 and finding_count <= 2:
            return "Low (< 1 hour)"
        elif file_count <= 3 and finding_count <= 5:
            return "Medium (1-4 hours)"
        elif file_count <= 5 and finding_count <= 10:
            return "High (4-8 hours)"
        else:
            return "Very High (> 8 hours)"
