#!/usr/bin/env python3
"""
Exploit suggestion engine.
Orchestrates version extraction, exploit matching, and ranking.
"""

import re
import time
from typing import Dict, List, Optional

from souleyez.intelligence.exploit_knowledge import ExploitKnowledgeBase
from souleyez.intelligence.service_parser import ServiceVersionExtractor
from souleyez.storage import exploit_attempts

# Module-level cache shared across all instances
_SUGGESTION_CACHE = {}
_CACHE_TIMEOUT = 30


def _generate_exploit_identifier(exploit: Dict) -> str:
    """
    Generate a unique identifier for an exploit.

    Priority order:
    1. MSF module path (most specific)
    2. CVE ID
    3. EDB ID
    4. Title hash (fallback)

    Returns:
        Unique identifier string like "msf:exploit/unix/ftp/vsftpd_234_backdoor"
        or "cve:CVE-2011-2523" or "edb:49757" or "title:<hash>"
    """
    if exploit.get("msf_module"):
        return f"msf:{exploit['msf_module']}"
    elif exploit.get("cve"):
        return f"cve:{exploit['cve']}"
    elif exploit.get("edb_id"):
        return f"edb:{exploit['edb_id']}"
    else:
        # Fallback: use title (should be rare)
        title_hash = hash(exploit.get("title", "unknown"))
        return f"title:{title_hash}"


class ExploitSuggestionEngine:
    """Generate exploit suggestions for engagement hosts."""

    def __init__(self, use_searchsploit: bool = True):
        """
        Initialize exploit suggestion engine.

        Args:
            use_searchsploit: Enable searchsploit integration (default: True)
        """
        self.version_extractor = ServiceVersionExtractor()
        self.knowledge_base = ExploitKnowledgeBase(use_searchsploit=use_searchsploit)

    def generate_suggestions(self, engagement_id: int, host_id: int = None) -> Dict:
        """
        Generate exploit suggestions for engagement or specific host.
        Results cached for 30 seconds to improve dashboard performance.

        Returns:
            {
                'hosts': [
                    {
                        'host_id': 123,
                        'ip': '10.0.0.5',
                        'hostname': 'example.com',
                        'services': [
                            {
                                'port': 21,
                                'service': 'ftp',
                                'version': 'vsftpd 2.3.4',
                                'product': 'vsftpd',
                                'version_number': '2.3.4',
                                'exploits': [
                                    {
                                        'title': 'vsftpd 2.3.4 Backdoor',
                                        'severity': 'critical',
                                        'msf_module': 'exploit/unix/ftp/vsftpd_234_backdoor',
                                        'cve': 'CVE-2011-2523',
                                        'description': '...',
                                        'match_type': 'exact'
                                    }
                                ],
                                'techniques': [
                                    'Anonymous FTP Access',
                                    'FTP Bounce Attack'
                                ]
                            }
                        ]
                    }
                ]
            }
        """
        # Check cache first
        cache_key = f"engagement_{engagement_id}_host_{host_id}"
        if cache_key in _SUGGESTION_CACHE:
            cached_result, cached_time = _SUGGESTION_CACHE[cache_key]
            if time.time() - cached_time < _CACHE_TIMEOUT:
                return cached_result

        # Cache miss or expired - do the analysis
        from souleyez.storage.findings import FindingsManager
        from souleyez.storage.hosts import HostManager
        from souleyez.storage.wazuh_vulns import WazuhVulnsManager

        hm = HostManager()
        fm = FindingsManager()
        wazuh_mgr = WazuhVulnsManager()

        # Fetch Wazuh vulnerabilities for this engagement
        wazuh_vulns = wazuh_mgr.list_vulnerabilities(engagement_id, limit=10000)

        # Get hosts to analyze
        if host_id:
            hosts = [hm.get_host(host_id)]
        else:
            hosts = hm.list_hosts(engagement_id)

        # Fetch all findings ONCE (not per-host - N+1 query fix)
        all_findings = fm.list_findings(engagement_id)

        results = []

        for host in hosts:
            if not host:
                continue

            h_id = host.get("id")
            host_ip = host.get("ip") or host.get("ip_address")

            # Filter findings for this host
            host_findings = [f for f in all_findings if f.get("host") == host_ip]

            # Get Wazuh vulnerabilities for this host
            host_wazuh_vulns = [
                v
                for v in wazuh_vulns
                if v.get("host_id") == h_id
                or v.get("agent_ip") == host_ip
                or v.get("host_ip") == host_ip
            ]

            # Build set of CVEs found by scans (from findings)
            scan_cves = set()
            for finding in host_findings:
                # Extract CVE from finding title or refs
                title = finding.get("title", "")
                refs = finding.get("refs", "")
                for text in [title, refs]:
                    cve_match = re.search(r"CVE-\d{4}-\d{4,}", text, re.IGNORECASE)
                    if cve_match:
                        scan_cves.add(cve_match.group(0).upper())

            # Build set of CVEs from Wazuh
            wazuh_cves = {
                v.get("cve_id").upper() for v in host_wazuh_vulns if v.get("cve_id")
            }

            # Extract versions from all sources
            services = self.version_extractor.combine_sources(h_id, host_findings)

            # Generate exploits for each service
            enriched_services = []
            seen_cves_in_services = set()

            for svc in services:
                exploits = self.knowledge_base.find_exploits(
                    svc["service"], svc["version_number"], svc.get("product")
                )

                # Enrich exploits with attempt status and detection source
                for exploit in exploits:
                    exploit_id = _generate_exploit_identifier(exploit)
                    exploit["identifier"] = exploit_id

                    # Look up attempt status
                    status = exploit_attempts.get_attempt_status(
                        engagement_id=engagement_id,
                        host_id=h_id,
                        exploit_identifier=exploit_id,
                        service_id=svc.get("service_id"),
                    )
                    exploit["attempt_status"] = status or "not_tried"

                    # Set detection source
                    cve = exploit.get("cve", "").upper() if exploit.get("cve") else None
                    if cve:
                        seen_cves_in_services.add(cve)
                        in_scan = cve in scan_cves
                        in_wazuh = cve in wazuh_cves
                        if in_scan and in_wazuh:
                            exploit["detection_source"] = "both"
                        elif in_wazuh:
                            exploit["detection_source"] = "wazuh"
                        else:
                            exploit["detection_source"] = "scan"
                    else:
                        exploit["detection_source"] = "scan"

                techniques = self.knowledge_base.get_techniques(svc["service"])

                enriched_services.append(
                    {**svc, "exploits": exploits, "techniques": techniques}
                )

            # Add Wazuh-only CVE exploits (not already covered by service exploits)
            wazuh_only_exploits = []
            for wazuh_vuln in host_wazuh_vulns:
                cve_id = wazuh_vuln.get("cve_id")
                if not cve_id:
                    continue
                cve_upper = cve_id.upper()
                if cve_upper in seen_cves_in_services:
                    continue  # Already covered by service-based exploits

                # Find exploits for this Wazuh CVE
                cve_exploits = self.knowledge_base.find_exploits_by_cve(cve_id)
                for exploit in cve_exploits:
                    exploit_id = _generate_exploit_identifier(exploit)
                    exploit["identifier"] = exploit_id

                    # Look up attempt status
                    status = exploit_attempts.get_attempt_status(
                        engagement_id=engagement_id,
                        host_id=h_id,
                        exploit_identifier=exploit_id,
                    )
                    exploit["attempt_status"] = status or "not_tried"

                    # Set detection source
                    in_scan = cve_upper in scan_cves
                    if in_scan:
                        exploit["detection_source"] = "both"
                    else:
                        exploit["detection_source"] = "wazuh"

                    # Add package context from Wazuh
                    exploit["wazuh_package"] = wazuh_vuln.get("package_name")
                    exploit["wazuh_version"] = wazuh_vuln.get("package_version")

                    wazuh_only_exploits.append(exploit)
                    seen_cves_in_services.add(cve_upper)

            # Add Wazuh-only exploits as a pseudo-service if any exist
            if wazuh_only_exploits:
                enriched_services.append(
                    {
                        "port": None,
                        "service": "wazuh_detected",
                        "version": None,
                        "product": None,
                        "version_number": None,
                        "exploits": wazuh_only_exploits,
                        "techniques": [],
                        "source": "wazuh",
                    }
                )

            # Sort all exploits to prioritize 'both' detection source
            for svc in enriched_services:
                if svc.get("exploits"):
                    svc["exploits"] = self._prioritize_exploits(svc["exploits"])

            # Only include hosts with exploits or techniques
            if any(s.get("exploits") or s.get("techniques") for s in enriched_services):
                results.append(
                    {
                        "host_id": h_id,
                        "ip": host_ip,
                        "hostname": host.get("hostname"),
                        "services": enriched_services,
                        "wazuh_cve_count": len(wazuh_cves),
                    }
                )

        result = {"hosts": results}

        # Cache the result
        _SUGGESTION_CACHE[cache_key] = (result, time.time())

        return result

    def get_host_summary(self, engagement_id: int, host_id: int) -> Dict:
        """
        Get exploit summary for a specific host.

        Returns:
            {
                'total_exploits': 5,
                'critical': 2,
                'high': 2,
                'medium': 1,
                'services_with_exploits': 3,
                'top_exploits': [...]
            }
        """
        suggestions = self.generate_suggestions(engagement_id, host_id)

        if not suggestions["hosts"]:
            return {
                "total_exploits": 0,
                "critical": 0,
                "high": 0,
                "medium": 0,
                "low": 0,
                "services_with_exploits": 0,
                "top_exploits": [],
            }

        host_data = suggestions["hosts"][0]
        all_exploits = []
        services_with_exploits = 0

        for svc in host_data["services"]:
            if svc.get("exploits"):
                services_with_exploits += 1
                all_exploits.extend(svc["exploits"])

        # Count by severity
        severity_counts = {
            "critical": len(
                [e for e in all_exploits if e.get("severity") == "critical"]
            ),
            "high": len([e for e in all_exploits if e.get("severity") == "high"]),
            "medium": len([e for e in all_exploits if e.get("severity") == "medium"]),
            "low": len([e for e in all_exploits if e.get("severity") == "low"]),
        }

        # Get top 5 exploits
        top_exploits = sorted(
            all_exploits,
            key=lambda x: (
                {"critical": 4, "high": 3, "medium": 2, "low": 1}.get(
                    x.get("severity", "low"), 0
                ),
                {"exact": 3, "range": 2, "generic": 1}.get(
                    x.get("match_type", "generic"), 0
                ),
            ),
            reverse=True,
        )[:5]

        return {
            "total_exploits": len(all_exploits),
            **severity_counts,
            "services_with_exploits": services_with_exploits,
            "top_exploits": top_exploits,
        }

    def get_engagement_summary(self, engagement_id: int) -> Dict:
        """
        Get exploit summary for entire engagement.

        Returns:
            {
                'total_hosts': 5,
                'hosts_with_exploits': 3,
                'total_exploits': 25,
                'critical': 5,
                'high': 10,
                'top_vulnerable_hosts': [...]
            }
        """
        suggestions = self.generate_suggestions(engagement_id)

        total_hosts = len(suggestions["hosts"])
        all_exploits = []

        host_exploit_counts = []

        for host in suggestions["hosts"]:
            host_exploits = []
            for svc in host["services"]:
                if svc.get("exploits"):
                    host_exploits.extend(svc["exploits"])

            all_exploits.extend(host_exploits)

            if host_exploits:
                host_exploit_counts.append(
                    {
                        "host_id": host["host_id"],
                        "ip": host["ip"],
                        "hostname": host.get("hostname"),
                        "exploit_count": len(host_exploits),
                        "critical_count": len(
                            [
                                e
                                for e in host_exploits
                                if e.get("severity") == "critical"
                            ]
                        ),
                    }
                )

        # Sort hosts by exploit count
        top_vulnerable = sorted(
            host_exploit_counts,
            key=lambda x: (x["critical_count"], x["exploit_count"]),
            reverse=True,
        )[:10]

        return {
            "total_hosts": total_hosts,
            "hosts_with_exploits": len(host_exploit_counts),
            "total_exploits": len(all_exploits),
            "critical": len(
                [e for e in all_exploits if e.get("severity") == "critical"]
            ),
            "high": len([e for e in all_exploits if e.get("severity") == "high"]),
            "medium": len([e for e in all_exploits if e.get("severity") == "medium"]),
            "low": len([e for e in all_exploits if e.get("severity") == "low"]),
            "top_vulnerable_hosts": top_vulnerable,
        }

    def _prioritize_exploits(self, exploits: List[Dict]) -> List[Dict]:
        """
        Prioritize exploits by detection source and severity.

        Priority order:
        1. detection_source: 'both' > 'wazuh' > 'scan'
        2. severity: critical > high > medium > low > info
        3. match_type: exact > range > generic

        Args:
            exploits: List of exploit dicts

        Returns:
            Sorted list of exploits
        """
        source_order = {"both": 3, "wazuh": 2, "scan": 1}
        severity_order = {"critical": 5, "high": 4, "medium": 3, "low": 2, "info": 1}
        match_order = {"exact": 3, "range": 2, "generic": 1}

        def sort_key(exploit):
            source_score = source_order.get(exploit.get("detection_source", "scan"), 0)
            severity_score = severity_order.get(exploit.get("severity", "low"), 0)
            match_score = match_order.get(exploit.get("match_type", "generic"), 0)
            return (source_score, severity_score, match_score)

        return sorted(exploits, key=sort_key, reverse=True)
