"""
MundiX Multi-Agent System
==========================

Specialized agents that execute pentest phases in parallel where possible.
Each agent has a role, a set of skills, and shares state via the Knowledge Graph.

Architecture:
    Coordinator
    ├── ReconAgent      (Phase 1: OSINT, passive recon)
    ├── ScanAgent       (Phase 2: Port scan, dir discovery — parallelized per host)
    ├── VulnAgent       (Phase 3: CVE search, vuln scanning — parallelized per host)
    ├── ExploitAgent    (Phase 4: SQLi, brute force — sequential, requires confirmation)
    ├── PostExAgent     (Phase 5: Privesc, lateral movement — sequential)
    └── ReporterAgent   (Phase 6: Report generation)

Parallelism rules:
    - Between phases: SEQUENTIAL (each phase depends on previous)
    - Within a phase: PARALLEL per host (scan host A and host B simultaneously)
    - Exploitation: SEQUENTIAL (destructive, needs human oversight at lower autopilot)

Usage:
    from multi_agent import AgentCoordinator
    
    coordinator = AgentCoordinator(orchestrator)
    coordinator.run_parallel()  # Enhanced run_plan with parallelism
"""

import json
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed, Future
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Dict, List, Optional, Callable, Any, Set

try:
    from knowledge_graph import KnowledgeGraph, NodeType
except ImportError:
    KnowledgeGraph = None
    NodeType = None


# ---------------------------------------------------------------------------
# Agent Roles
# ---------------------------------------------------------------------------

class AgentRole(Enum):
    RECON = "recon"
    SCANNER = "scanner"
    VULN_ANALYST = "vuln_analyst"
    EXPLOITER = "exploiter"
    POST_EXPLOITER = "post_exploiter"
    REPORTER = "reporter"


# Phase → Agent Role mapping
PHASE_AGENT_MAP = {
    "1_reconnaissance": AgentRole.RECON,
    "2_scanning": AgentRole.SCANNER,
    "3_vulnerability_analysis": AgentRole.VULN_ANALYST,
    "4_exploitation": AgentRole.EXPLOITER,
    "5_post_exploitation": AgentRole.POST_EXPLOITER,
    "6_reporting": AgentRole.REPORTER,
}

# Which phases allow parallel execution per host
PARALLELIZABLE_PHASES = {
    "2_scanning",
    "3_vulnerability_analysis",
}

# Phases that need sequential (human-supervised) execution
SEQUENTIAL_PHASES = {
    "4_exploitation",
    "5_post_exploitation",
}


# ---------------------------------------------------------------------------
# Specialist Agent
# ---------------------------------------------------------------------------

@dataclass
class AgentTask:
    """A unit of work for a specialist agent."""
    skill_name: str
    target_host: Optional[str] = None  # None = all hosts
    phase: str = ""
    priority: int = 5  # 1=highest

    def __repr__(self):
        host = f"@{self.target_host}" if self.target_host else ""
        return f"Task({self.skill_name}{host})"


@dataclass
class AgentResult:
    """Result of an agent task execution."""
    task: AgentTask
    success: bool
    duration: float = 0.0
    findings_count: int = 0
    error: Optional[str] = None


class SpecialistAgent:
    """
    A specialist agent that executes skills within a specific domain.
    Wraps the orchestrator's _execute_skill with host-specific context.
    """

    def __init__(self, role: AgentRole, orchestrator, agent_id: str = None):
        self.role = role
        self.orchestrator = orchestrator
        self.agent_id = agent_id or f"agent-{role.value}-{id(self) % 10000:04d}"
        self.tasks_completed: List[AgentResult] = []
        self.active = False
        self._lock = threading.Lock()

    def execute_task(self, task: AgentTask) -> AgentResult:
        """Execute a single skill task, optionally scoped to a specific host."""
        self.active = True
        start_time = time.time()

        try:
            # If task is host-specific, temporarily scope variables
            original_vars = None
            if task.target_host and self.orchestrator.state:
                original_vars = dict(self.orchestrator.state.variables)
                self.orchestrator.state.variables["target_ip"] = task.target_host
                self.orchestrator.state.variables["target_domain"] = task.target_host

            # Execute the skill
            self.orchestrator._execute_skill(task.skill_name)

            # Check result
            exec_state = self.orchestrator.state.skill_executions.get(task.skill_name)
            success = exec_state and exec_state.status.value == "completed"
            findings = len(exec_state.findings) if exec_state else 0

            # Restore variables
            if original_vars and self.orchestrator.state:
                self.orchestrator.state.variables = original_vars

            result = AgentResult(
                task=task, success=success,
                duration=time.time() - start_time,
                findings_count=findings,
            )

        except Exception as e:
            result = AgentResult(
                task=task, success=False,
                duration=time.time() - start_time,
                error=str(e),
            )

        with self._lock:
            self.tasks_completed.append(result)
        self.active = False
        return result

    def get_stats(self) -> Dict[str, Any]:
        """Get agent execution stats."""
        completed = [r for r in self.tasks_completed if r.success]
        failed = [r for r in self.tasks_completed if not r.success]
        total_time = sum(r.duration for r in self.tasks_completed)
        total_findings = sum(r.findings_count for r in self.tasks_completed)
        return {
            "agent_id": self.agent_id,
            "role": self.role.value,
            "tasks_total": len(self.tasks_completed),
            "tasks_success": len(completed),
            "tasks_failed": len(failed),
            "total_time": round(total_time, 1),
            "total_findings": total_findings,
            "active": self.active,
        }


# ---------------------------------------------------------------------------
# Agent Coordinator
# ---------------------------------------------------------------------------

class AgentCoordinator:
    """
    Coordinates multiple specialist agents for parallel pentest execution.
    
    The coordinator:
    1. Analyzes the engagement plan
    2. Creates specialist agents per role
    3. Distributes work (parallelized within phases, sequential between phases)
    4. Merges results into the shared Knowledge Graph
    5. Makes inter-phase adaptation decisions
    """

    def __init__(self, orchestrator, max_workers: int = 4):
        """
        Args:
            orchestrator: AgentOrchestrator instance (must have started engagement)
            max_workers: Maximum parallel threads for host-parallel execution
        """
        self.orchestrator = orchestrator
        self.max_workers = max_workers
        self.agents: Dict[AgentRole, SpecialistAgent] = {}
        self._executor: Optional[ThreadPoolExecutor] = None

        # Create specialist agents
        for role in AgentRole:
            self.agents[role] = SpecialistAgent(role, orchestrator)

    def run_parallel(self) -> Dict[str, Any]:
        """
        Execute the pentest plan with multi-agent parallelism.
        
        Returns summary dict with agent stats and results.
        """
        state = self.orchestrator.state
        if not state:
            raise RuntimeError("No engagement started.")

        self.orchestrator.on_status(
            f"\n[MULTI-AGENT] 🚀 Coordinator initialized with {len(self.agents)} specialist agents"
        )

        # Group plan by phase
        phase_skills: Dict[str, List[str]] = {}
        for step in state.plan:
            phase_skills.setdefault(step["phase"], []).append(step["skill"])

        results_summary = {"phases": {}, "agents": {}}

        for phase_id in phase_skills:
            if self.orchestrator._abort_requested:
                state.status = self.orchestrator.state.status  # Keep abort status
                break

            # Phase transition
            state.current_phase = phase_id
            self.orchestrator.on_phase_change(phase_id)
            self.orchestrator.on_status(f"\n{'='*60}")
            self.orchestrator.on_status(f"[PHASE] {phase_id.replace('_', ' ').upper()}")
            self.orchestrator.on_status(f"{'='*60}")
            state.add_timeline_event("phase_started", f"Phase: {phase_id}")

            # Get skills (with AI adaptation)
            skill_names = list(phase_skills[phase_id])
            if self.orchestrator.ai_provider and state.get_completed_skills():
                adapted = self.orchestrator._ai_adapt_phase(phase_id, skill_names)
                if adapted:
                    skill_names = adapted

            agent_role = PHASE_AGENT_MAP.get(phase_id, AgentRole.RECON)
            agent = self.agents[agent_role]

            # Decide execution strategy
            if phase_id in PARALLELIZABLE_PHASES and len(state.discovered_hosts) > 1:
                phase_results = self._execute_parallel(
                    agent, skill_names, state.discovered_hosts, phase_id
                )
            else:
                phase_results = self._execute_sequential(
                    agent, skill_names, phase_id
                )

            results_summary["phases"][phase_id] = {
                "skills": len(skill_names),
                "success": sum(1 for r in phase_results if r.success),
                "failed": sum(1 for r in phase_results if not r.success),
                "duration": round(sum(r.duration for r in phase_results), 1),
            }

            # Post-phase hooks (same as original run_plan)
            if state.findings.get("waf_detected") and phase_id.startswith("2_"):
                self.orchestrator.on_status(
                    f"\n[ADAPT] ⚡ WAF detected! Adjusting tactics for exploitation phase."
                )
                state.variables["waf_evasion"] = "true"

            if phase_id.startswith("2_") and self.orchestrator.threat_intel and state.discovered_ports:
                self.orchestrator._enrich_with_threat_intel()

            # Save state
            state_file = os.path.join(state.workspace, "engagement_state.json")
            state.save(state_file)

        # Completion
        if not self.orchestrator._abort_requested:
            state.complete()
            state.add_timeline_event("engagement_completed", "Multi-agent execution finished")

            self.orchestrator._save_engagement_learnings()

            if self.orchestrator.knowledge_graph:
                kg = self.orchestrator.knowledge_graph
                kg_path = os.path.join(state.workspace, "knowledge_graph.json")
                kg.save(kg_path)
                chains = kg.find_attack_chains()
                self.orchestrator.on_status(f"\n[KG] {len(kg.nodes)} nodes, {len(kg.edges)} edges")
                if chains:
                    self.orchestrator.on_status(f"[KG] 🎯 {len(chains)} attack chain(s)")

        # Agent stats
        for role, agent in self.agents.items():
            stats = agent.get_stats()
            if stats["tasks_total"] > 0:
                results_summary["agents"][role.value] = stats

        self.orchestrator.on_status(
            f"\n[MULTI-AGENT] ✅ Complete — "
            f"{sum(s['success'] for s in results_summary['phases'].values())} skills succeeded, "
            f"{sum(s['failed'] for s in results_summary['phases'].values())} failed"
        )

        return results_summary

    def _execute_sequential(self, agent: SpecialistAgent,
                           skill_names: List[str], phase_id: str) -> List[AgentResult]:
        """Execute skills sequentially (standard path)."""
        results = []
        failed_count = 0

        for skill_name in skill_names:
            if self.orchestrator._abort_requested:
                break
            if failed_count >= 5:
                self.orchestrator.on_status(
                    f"[ADAPT] Too many failures in {phase_id}. Moving on."
                )
                break

            self.orchestrator.state.current_skill = skill_name
            self.orchestrator.on_status(f"\n[{agent.role.value.upper()}] Starting: {skill_name}")

            task = AgentTask(skill_name=skill_name, phase=phase_id)
            result = agent.execute_task(task)
            results.append(result)

            if result.success:
                failed_count = 0
            else:
                failed_count += 1
                if result.error:
                    self.orchestrator.on_status(f"  [ERROR] {result.error}")

        return results

    def _execute_parallel(self, agent: SpecialistAgent,
                         skill_names: List[str], hosts: List[str],
                         phase_id: str) -> List[AgentResult]:
        """
        Execute skills in parallel across multiple hosts.
        Each host gets its own thread executing the skill list.
        """
        self.orchestrator.on_status(
            f"[PARALLEL] 🔀 Splitting {len(skill_names)} skills across {len(hosts)} hosts "
            f"(max {self.max_workers} threads)"
        )

        results = []
        
        # For host-parallelizable skills, create per-host tasks
        # Some skills (like OSINT) don't need per-host splitting
        host_specific_skills = []
        global_skills = []
        
        for skill_name in skill_names:
            skill = self.orchestrator.registry.get(skill_name)
            if skill and any(t in skill.tools for t in ["nmap", "masscan", "gobuster", 
                            "feroxbuster", "ffuf", "nuclei", "nikto", "testssl"]):
                host_specific_skills.append(skill_name)
            else:
                global_skills.append(skill_name)

        # Run global skills first (sequentially)
        for skill_name in global_skills:
            if self.orchestrator._abort_requested:
                break
            self.orchestrator.on_status(f"\n[{agent.role.value.upper()}] {skill_name} (global)")
            task = AgentTask(skill_name=skill_name, phase=phase_id)
            result = agent.execute_task(task)
            results.append(result)

        # Run host-specific skills in parallel
        if host_specific_skills and not self.orchestrator._abort_requested:
            tasks = []
            for host in hosts[:10]:  # Cap at 10 hosts to avoid resource exhaustion
                for skill_name in host_specific_skills:
                    tasks.append(AgentTask(
                        skill_name=skill_name,
                        target_host=host,
                        phase=phase_id,
                    ))

            self.orchestrator.on_status(
                f"[PARALLEL] Queuing {len(tasks)} tasks across {min(len(hosts), 10)} hosts"
            )

            # Use ThreadPoolExecutor for I/O-bound tasks (command execution)
            with ThreadPoolExecutor(max_workers=min(self.max_workers, len(tasks))) as executor:
                futures: Dict[Future, AgentTask] = {}
                for task in tasks:
                    future = executor.submit(agent.execute_task, task)
                    futures[future] = task

                for future in as_completed(futures):
                    task = futures[future]
                    try:
                        result = future.result(timeout=300)  # 5 min max per task
                        results.append(result)
                        status = "✅" if result.success else "❌"
                        self.orchestrator.on_status(
                            f"  {status} {task.skill_name}@{task.target_host} "
                            f"({result.duration:.1f}s, {result.findings_count} findings)"
                        )
                    except Exception as e:
                        results.append(AgentResult(
                            task=task, success=False, error=str(e)
                        ))
                        self.orchestrator.on_status(
                            f"  ❌ {task.skill_name}@{task.target_host}: {e}"
                        )

        return results

    def get_dashboard(self) -> str:
        """Get a text dashboard of agent status."""
        lines = ["## Multi-Agent Dashboard\n"]
        for role, agent in self.agents.items():
            stats = agent.get_stats()
            if stats["tasks_total"] > 0:
                status = "🟢 Active" if stats["active"] else "⚪ Idle"
                lines.append(
                    f"**{role.value.title()}** ({status}): "
                    f"{stats['tasks_success']}/{stats['tasks_total']} OK, "
                    f"{stats['total_findings']} findings, {stats['total_time']}s"
                )
        return '\n'.join(lines)
