"""
MundiX Agent Orchestrator — The autonomous decision loop for pentesting.

This module is the "brain" of the MundiX autonomous pentesting agent.
It uses the SkillLoader to plan which skills to execute, the AIProvider
to analyze results and make decisions, and the Project to track state.

Architecture:
    ┌─────────────────────────────────────────────┐
    │              Agent Orchestrator              │
    │  ┌─────────┐  ┌──────────┐  ┌───────────┐  │
    │  │  Skill   │  │    AI    │  │  Project  │  │
    │  │  Loader  │  │ Provider │  │   State   │  │
    │  └─────────┘  └──────────┘  └───────────┘  │
    │  ┌─────────┐  ┌──────────┐  ┌───────────┐  │
    │  │ AutoPilot│  │ Executor │  │  Reporter │  │
    │  └─────────┘  └──────────┘  └───────────┘  │
    └─────────────────────────────────────────────┘

Flow (Autopilot 5 — Full Autonomy):
    1. User provides target → Orchestrator creates engagement
    2. Load skills → Generate pentest plan
    3. For each phase:
       a. Select skills for current phase
       b. For each skill:
          - Generate agent prompt with current context
          - AI decides which commands to run
          - Execute commands (respecting autopilot level)
          - Parse output → update findings
          - AI analyzes results → decides next action
       c. Phase complete → advance to next phase
    4. Generate final report

Usage:
    from agent_orchestrator import AgentOrchestrator

    orchestrator = AgentOrchestrator(
        skills_dir="/opt/mundix-cli/skills",
        ai_provider=ai_provider_instance,
        project=project_instance,
        autopilot=autopilot_instance,
    )

    # Start a full autonomous pentest
    await orchestrator.start_engagement("example.com", engagement_type="web")

    # Or run a single skill
    await orchestrator.run_skill("network_port_scan", target="10.0.0.1")
"""

import os
import re
import json
import time
import asyncio
import subprocess
import shlex
from enum import Enum
from typing import List, Dict, Optional, Callable, Tuple, Any
from dataclasses import dataclass, field
from pathlib import Path
from datetime import datetime

from skill_loader import SkillRegistry, Skill, SkillCommand


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

SKILLS_DIR = os.environ.get("MUNDIX_SKILLS_DIR", "/opt/mundix-cli/skills")

# Maximum time (seconds) a single command can run before being killed
COMMAND_TIMEOUT = 300  # 5 minutes

# Maximum retries for a failed command
MAX_COMMAND_RETRIES = 3

# Maximum time (seconds) for an entire skill execution
SKILL_TIMEOUT = 1800  # 30 minutes

# How many lines of command output to keep for AI context
MAX_OUTPUT_LINES = 200

# Engagement types map to which phases to run
ENGAGEMENT_PROFILES = {
    "web": {
        "description": "Web Application Penetration Test",
        "phases": [
            "1_reconnaissance",
            "2_scanning_enumeration",
            "3_vulnerability_analysis",
            "4_exploitation",
            "5_post_exploitation",
            "6_reporting",
        ],
        "skip_skills": [],
        "required_vars": ["target_domain"],
    },
    "internal": {
        "description": "Internal Network Penetration Test",
        "phases": [
            "1_reconnaissance",
            "2_scanning_enumeration",
            "3_vulnerability_analysis",
            "4_exploitation",
            "5_post_exploitation",
            "6_reporting",
        ],
        "skip_skills": ["osint_domain_enumeration", "osint_email_harvesting"],
        "required_vars": ["target_ip", "target_range"],
    },
    "app": {
        "description": "Application Security Assessment",
        "phases": [
            "1_reconnaissance",
            "2_scanning_enumeration",
            "3_vulnerability_analysis",
            "4_exploitation",
            "6_reporting",
        ],
        "skip_skills": [
            "smb_rpc_enumeration", "dns_zone_transfer",
            "linux_privilege_escalation", "windows_privilege_escalation",
            "lateral_movement_psexec", "hash_dumping_and_cracking",
        ],
        "required_vars": ["target_domain"],
    },
}


# ---------------------------------------------------------------------------
# Enums and Data Classes
# ---------------------------------------------------------------------------

class EngagementStatus(str, Enum):
    """Status of the current engagement."""
    IDLE = "idle"
    PLANNING = "planning"
    RUNNING = "running"
    PAUSED = "paused"
    WAITING_INPUT = "waiting_input"
    COMPLETED = "completed"
    FAILED = "failed"
    ABORTED = "aborted"


class SkillStatus(str, Enum):
    """Status of a skill execution."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    SKIPPED = "skipped"


class CommandResult:
    """Result of executing a single command."""
    def __init__(self, command: str, stdout: str, stderr: str,
                 returncode: int, duration: float):
        self.command = command
        self.stdout = stdout
        self.stderr = stderr
        self.returncode = returncode
        self.duration = duration
        self.timestamp = datetime.now().isoformat()

    @property
    def success(self) -> bool:
        return self.returncode == 0

    @property
    def output(self) -> str:
        """Combined output (stdout + stderr), truncated for AI context."""
        combined = self.stdout
        if self.stderr:
            combined += f"\n[STDERR]\n{self.stderr}"
        lines = combined.split('\n')
        if len(lines) > MAX_OUTPUT_LINES:
            half = MAX_OUTPUT_LINES // 2
            lines = lines[:half] + [f"\n... ({len(lines) - MAX_OUTPUT_LINES} lines truncated) ...\n"] + lines[-half:]
        return '\n'.join(lines)

    def to_dict(self) -> Dict:
        return {
            "command": self.command,
            "returncode": self.returncode,
            "success": self.success,
            "duration": round(self.duration, 2),
            "stdout_lines": len(self.stdout.split('\n')),
            "stderr_lines": len(self.stderr.split('\n')) if self.stderr else 0,
            "timestamp": self.timestamp,
        }


@dataclass
class SkillExecution:
    """Tracks the execution state of a single skill."""
    skill_name: str
    status: SkillStatus = SkillStatus.PENDING
    started_at: Optional[str] = None
    completed_at: Optional[str] = None
    commands_executed: List[Dict] = field(default_factory=list)
    findings: List[str] = field(default_factory=list)
    errors: List[str] = field(default_factory=list)
    ai_decisions: List[str] = field(default_factory=list)
    output_files: List[str] = field(default_factory=list)
    skip_reason: Optional[str] = None

    def start(self):
        self.status = SkillStatus.RUNNING
        self.started_at = datetime.now().isoformat()

    def complete(self, findings: List[str] = None):
        self.status = SkillStatus.COMPLETED
        self.completed_at = datetime.now().isoformat()
        if findings:
            self.findings.extend(findings)

    def fail(self, error: str):
        self.status = SkillStatus.FAILED
        self.completed_at = datetime.now().isoformat()
        self.errors.append(error)

    def skip(self, reason: str):
        self.status = SkillStatus.SKIPPED
        self.skip_reason = reason

    def to_dict(self) -> Dict:
        return {
            "skill": self.skill_name,
            "status": self.status.value,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
            "commands_count": len(self.commands_executed),
            "findings_count": len(self.findings),
            "findings": self.findings,
            "errors": self.errors,
            "ai_decisions": self.ai_decisions,
            "output_files": self.output_files,
        }


@dataclass
class EngagementState:
    """Full state of the current pentest engagement."""
    target: str
    engagement_type: str
    status: EngagementStatus = EngagementStatus.IDLE
    workspace: str = ""
    variables: Dict[str, str] = field(default_factory=dict)
    current_phase: str = ""
    current_skill: str = ""
    plan: List[Dict] = field(default_factory=list)
    skill_executions: Dict[str, SkillExecution] = field(default_factory=dict)
    findings: Dict[str, bool] = field(default_factory=dict)
    discovered_hosts: List[str] = field(default_factory=list)
    discovered_ports: List[Dict] = field(default_factory=list)
    discovered_vulns: List[Dict] = field(default_factory=list)
    discovered_creds: List[Dict] = field(default_factory=list)
    timeline: List[Dict] = field(default_factory=list)
    started_at: Optional[str] = None
    completed_at: Optional[str] = None

    def start(self):
        self.status = EngagementStatus.RUNNING
        self.started_at = datetime.now().isoformat()

    def complete(self):
        self.status = EngagementStatus.COMPLETED
        self.completed_at = datetime.now().isoformat()

    def add_timeline_event(self, event_type: str, description: str, data: Dict = None):
        self.timeline.append({
            "timestamp": datetime.now().isoformat(),
            "type": event_type,
            "description": description,
            "data": data or {},
        })

    def get_completed_skills(self) -> List[str]:
        return [
            name for name, ex in self.skill_executions.items()
            if ex.status == SkillStatus.COMPLETED
        ]

    def get_progress(self) -> Dict:
        total = len(self.plan)
        completed = len(self.get_completed_skills())
        failed = len([e for e in self.skill_executions.values() if e.status == SkillStatus.FAILED])
        skipped = len([e for e in self.skill_executions.values() if e.status == SkillStatus.SKIPPED])
        return {
            "total": total,
            "completed": completed,
            "failed": failed,
            "skipped": skipped,
            "remaining": total - completed - failed - skipped,
            "percent": round((completed / total) * 100, 1) if total > 0 else 0,
            "current_phase": self.current_phase,
            "current_skill": self.current_skill,
        }

    def get_context_summary(self) -> str:
        """Generate a context summary for the AI to understand current state."""
        lines = []
        lines.append(f"## Engagement State")
        lines.append(f"**Target:** {self.target}")
        lines.append(f"**Type:** {self.engagement_type}")
        lines.append(f"**Phase:** {self.current_phase}")
        lines.append(f"**Progress:** {self.get_progress()['percent']}%")
        lines.append("")

        if self.discovered_hosts:
            lines.append(f"**Discovered Hosts:** {', '.join(self.discovered_hosts[:20])}")
        if self.discovered_ports:
            lines.append(f"**Open Ports:** {len(self.discovered_ports)} found")
            for p in self.discovered_ports[:15]:
                lines.append(f"  - {p.get('host', '?')}:{p.get('port', '?')} "
                             f"({p.get('service', 'unknown')})")
        if self.discovered_vulns:
            lines.append(f"**Vulnerabilities:** {len(self.discovered_vulns)} found")
            for v in self.discovered_vulns[:10]:
                lines.append(f"  - [{v.get('severity', '?')}] {v.get('name', 'unknown')}")
        if self.discovered_creds:
            lines.append(f"**Credentials:** {len(self.discovered_creds)} found")

        # Recent findings
        completed_skills = self.get_completed_skills()
        if completed_skills:
            lines.append(f"\n**Completed Skills:** {', '.join(completed_skills[-5:])}")
            # Last skill's findings
            last_skill = completed_skills[-1]
            last_exec = self.skill_executions.get(last_skill)
            if last_exec and last_exec.findings:
                lines.append(f"**Last Findings:**")
                for f in last_exec.findings[-5:]:
                    lines.append(f"  - {f}")

        return '\n'.join(lines)

    def save(self, filepath: str):
        """Save engagement state to JSON file."""
        data = {
            "target": self.target,
            "engagement_type": self.engagement_type,
            "status": self.status.value,
            "workspace": self.workspace,
            "variables": self.variables,
            "current_phase": self.current_phase,
            "current_skill": self.current_skill,
            "findings": self.findings,
            "discovered_hosts": self.discovered_hosts,
            "discovered_ports": self.discovered_ports,
            "discovered_vulns": self.discovered_vulns,
            "discovered_creds": self.discovered_creds,
            "timeline": self.timeline,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
            "skill_executions": {
                k: v.to_dict() for k, v in self.skill_executions.items()
            },
            "progress": self.get_progress(),
        }
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    @classmethod
    def load(cls, filepath: str) -> 'EngagementState':
        """Load engagement state from JSON file."""
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)

        state = cls(
            target=data["target"],
            engagement_type=data["engagement_type"],
        )
        state.status = EngagementStatus(data.get("status", "idle"))
        state.workspace = data.get("workspace", "")
        state.variables = data.get("variables", {})
        state.current_phase = data.get("current_phase", "")
        state.current_skill = data.get("current_skill", "")
        state.findings = data.get("findings", {})
        state.discovered_hosts = data.get("discovered_hosts", [])
        state.discovered_ports = data.get("discovered_ports", [])
        state.discovered_vulns = data.get("discovered_vulns", [])
        state.discovered_creds = data.get("discovered_creds", [])
        state.timeline = data.get("timeline", [])
        state.started_at = data.get("started_at")
        state.completed_at = data.get("completed_at")
        return state


# ---------------------------------------------------------------------------
# Command Executor
# ---------------------------------------------------------------------------

class CommandExecutor:
    """
    Executes shell commands with timeout, output capture, and safety checks.
    Integrates with AutoPilot for execution permission.
    """

    # Commands that should NEVER be auto-executed
    DESTRUCTIVE_PATTERNS = [
        r'rm\s+-rf\s+/',
        r'mkfs\.',
        r'dd\s+if=.*of=/dev/',
        r':\(\)\{.*\}',  # fork bomb
        r'chmod\s+-R\s+777\s+/',
        r'shutdown',
        r'reboot',
        r'init\s+0',
        r'halt',
    ]

    # Commands that require elevated caution
    DANGEROUS_PATTERNS = [
        r'rm\s+-rf',
        r'rm\s+-r',
        r'exploit',
        r'msfconsole',
        r'msfvenom',
        r'reverse.?shell',
        r'bind.?shell',
        r'nc\s+-[le]',  # netcat listener
        r'python.*-c.*socket',  # python reverse shell
        r'bash\s+-i\s+>&',  # bash reverse shell
    ]

    def __init__(self, autopilot=None, on_output: Callable = None,
                 on_ask_permission: Callable = None):
        """
        Args:
            autopilot: AutoPilot instance for permission checks
            on_output: Callback for real-time output streaming (line: str) -> None
            on_ask_permission: Callback to ask user permission (cmd: str, reason: str) -> bool
        """
        self.autopilot = autopilot
        self.on_output = on_output
        self.on_ask_permission = on_ask_permission

    def is_destructive(self, command: str) -> bool:
        """Check if a command matches destructive patterns."""
        for pattern in self.DESTRUCTIVE_PATTERNS:
            if re.search(pattern, command, re.IGNORECASE):
                return True
        return False

    def is_dangerous(self, command: str) -> bool:
        """Check if a command matches dangerous patterns."""
        for pattern in self.DANGEROUS_PATTERNS:
            if re.search(pattern, command, re.IGNORECASE):
                return True
        return False

    def should_execute(self, command: str) -> Tuple[bool, str]:
        """
        Determine if a command should be executed based on AutoPilot level.
        Returns (should_execute: bool, reason: str).
        """
        if self.is_destructive(command):
            if self.autopilot and self.autopilot.level < 5:
                return False, f"DESTRUCTIVE command blocked (AP level {self.autopilot.level})"
            # Level 5: even destructive commands run
            return True, "AUTOPILOT: destructive command allowed"

        if self.is_dangerous(command):
            if self.autopilot:
                if self.autopilot.should_skip_dangerous_confirm():
                    return True, "AGGRESSIVE/AUTOPILOT: dangerous command auto-approved"
                else:
                    # Need to ask user
                    return False, "DANGEROUS command requires confirmation"
            return False, "DANGEROUS command (no autopilot configured)"

        # Safe command
        if self.autopilot and self.autopilot.should_auto_execute():
            return True, "Auto-execute (safe command)"

        return False, "Manual mode: requires confirmation"

    def execute(self, command: str, timeout: int = COMMAND_TIMEOUT,
                cwd: str = None, env: Dict = None) -> CommandResult:
        """
        Execute a shell command and return the result.

        Args:
            command: The shell command to execute
            timeout: Maximum seconds to wait
            cwd: Working directory
            env: Additional environment variables
        """
        start_time = time.time()

        # Merge environment
        cmd_env = os.environ.copy()
        if env:
            cmd_env.update(env)

        try:
            process = subprocess.Popen(
                command,
                shell=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                cwd=cwd,
                env=cmd_env,
                text=True,
                bufsize=1,
            )

            stdout_lines = []
            stderr_lines = []

            # Read output with timeout
            try:
                stdout, stderr = process.communicate(timeout=timeout)
                stdout_lines = stdout.split('\n') if stdout else []
                stderr_lines = stderr.split('\n') if stderr else []

                # Stream output if callback provided
                if self.on_output:
                    for line in stdout_lines:
                        if line.strip():
                            self.on_output(line)

            except subprocess.TimeoutExpired:
                process.kill()
                stdout, stderr = process.communicate()
                return CommandResult(
                    command=command,
                    stdout=stdout or "",
                    stderr=f"TIMEOUT: Command exceeded {timeout}s limit\n{stderr or ''}",
                    returncode=-1,
                    duration=timeout,
                )

            duration = time.time() - start_time

            return CommandResult(
                command=command,
                stdout='\n'.join(stdout_lines),
                stderr='\n'.join(stderr_lines),
                returncode=process.returncode,
                duration=duration,
            )

        except Exception as e:
            duration = time.time() - start_time
            return CommandResult(
                command=command,
                stdout="",
                stderr=f"EXECUTION ERROR: {str(e)}",
                returncode=-2,
                duration=duration,
            )

    def execute_with_permission(self, command: str, timeout: int = COMMAND_TIMEOUT,
                                 cwd: str = None) -> Tuple[Optional[CommandResult], str]:
        """
        Execute a command with autopilot permission checks.
        Returns (result or None, status_message).
        """
        should_run, reason = self.should_execute(command)

        if should_run:
            result = self.execute(command, timeout=timeout, cwd=cwd)
            return result, reason

        # Need permission
        if self.on_ask_permission:
            is_dangerous = self.is_dangerous(command)
            granted = self.on_ask_permission(
                command,
                "DANGEROUS" if is_dangerous else "DESTRUCTIVE"
            )
            if granted:
                result = self.execute(command, timeout=timeout, cwd=cwd)
                return result, "User approved execution"
            else:
                return None, "User denied execution"

        return None, reason


# ---------------------------------------------------------------------------
# Output Parser — Extracts structured data from command output
# ---------------------------------------------------------------------------

class OutputParser:
    """Parses command output to extract hosts, ports, vulnerabilities, etc."""

    @staticmethod
    def parse_nmap_output(output: str) -> Dict:
        """Extract hosts, ports, and services from nmap output."""
        results = {"hosts": [], "ports": [], "os_matches": []}

        # Parse open ports: "80/tcp   open  http    Apache httpd 2.4.41"
        port_re = re.compile(
            r'(\d+)/(tcp|udp)\s+(open|filtered)\s+(\S+)\s*(.*)',
            re.IGNORECASE
        )
        current_host = None

        for line in output.split('\n'):
            # Host detection
            host_match = re.search(r'Nmap scan report for\s+(\S+)', line)
            if host_match:
                current_host = host_match.group(1)
                # Extract IP if hostname(IP) format
                ip_match = re.search(r'\((\d+\.\d+\.\d+\.\d+)\)', current_host)
                if ip_match:
                    results["hosts"].append({
                        "hostname": current_host.split('(')[0],
                        "ip": ip_match.group(1),
                    })
                else:
                    results["hosts"].append({"hostname": current_host, "ip": current_host})

            # Port detection
            port_match = port_re.search(line)
            if port_match:
                results["ports"].append({
                    "host": current_host or "unknown",
                    "port": int(port_match.group(1)),
                    "protocol": port_match.group(2),
                    "state": port_match.group(3),
                    "service": port_match.group(4),
                    "version": port_match.group(5).strip(),
                })

            # OS detection
            os_match = re.search(r'OS details?:\s+(.+)', line)
            if os_match:
                results["os_matches"].append(os_match.group(1).strip())

        return results

    @staticmethod
    def parse_gobuster_output(output: str) -> Dict:
        """Extract discovered directories and files from gobuster output."""
        results = {"directories": [], "files": []}

        for line in output.split('\n'):
            # Match: /admin (Status: 200) [Size: 1234]
            match = re.search(r'(/\S+)\s+\(Status:\s*(\d+)\)\s*\[Size:\s*(\d+)\]', line)
            if match:
                path = match.group(1)
                status = int(match.group(2))
                size = int(match.group(3))
                entry = {"path": path, "status": status, "size": size}
                if path.endswith('/') or '.' not in path.split('/')[-1]:
                    results["directories"].append(entry)
                else:
                    results["files"].append(entry)

        return results

    @staticmethod
    def parse_nuclei_output(output: str) -> List[Dict]:
        """Extract vulnerabilities from nuclei output."""
        vulns = []

        for line in output.split('\n'):
            # Nuclei format: [severity] [template-id] [protocol] url
            match = re.search(
                r'\[(\w+)\]\s+\[([^\]]+)\]\s+\[([^\]]+)\]\s+(.+)',
                line
            )
            if match:
                vulns.append({
                    "severity": match.group(1).lower(),
                    "template": match.group(2),
                    "protocol": match.group(3),
                    "url": match.group(4).strip(),
                    "raw": line.strip(),
                })

        return vulns

    @staticmethod
    def parse_sqlmap_output(output: str) -> Dict:
        """Extract SQLi findings from sqlmap output."""
        results = {
            "injectable": False,
            "parameters": [],
            "databases": [],
            "dbms": None,
        }

        for line in output.split('\n'):
            if "is vulnerable" in line.lower() or "injectable" in line.lower():
                results["injectable"] = True
            param_match = re.search(r"Parameter:\s+(\S+)", line)
            if param_match:
                results["parameters"].append(param_match.group(1))
            db_match = re.search(r"available databases.*:\s*\[.*\]\s*(.*)", line)
            if db_match:
                results["databases"].extend(
                    [d.strip().strip("'\"") for d in db_match.group(1).split(',')]
                )
            dbms_match = re.search(r"back-end DBMS:\s+(.+)", line)
            if dbms_match:
                results["dbms"] = dbms_match.group(1).strip()

        return results

    @staticmethod
    def parse_hydra_output(output: str) -> List[Dict]:
        """Extract cracked credentials from hydra output."""
        creds = []

        for line in output.split('\n'):
            # [22][ssh] host: 10.0.0.1   login: admin   password: password123
            match = re.search(
                r'\[\d+\]\[(\w+)\]\s+host:\s+(\S+)\s+login:\s+(\S+)\s+password:\s+(\S+)',
                line
            )
            if match:
                creds.append({
                    "service": match.group(1),
                    "host": match.group(2),
                    "username": match.group(3),
                    "password": match.group(4),
                })

        return creds

    def auto_parse(self, command: str, output: str) -> Dict:
        """Automatically detect the tool and parse its output."""
        cmd_lower = command.lower().split()[0] if command else ""

        if "nmap" in cmd_lower:
            return {"type": "nmap", "data": self.parse_nmap_output(output)}
        elif "gobuster" in cmd_lower or "feroxbuster" in cmd_lower or "dirb" in cmd_lower:
            return {"type": "directory_scan", "data": self.parse_gobuster_output(output)}
        elif "nuclei" in cmd_lower or "nikto" in cmd_lower:
            return {"type": "vuln_scan", "data": self.parse_nuclei_output(output)}
        elif "sqlmap" in cmd_lower:
            return {"type": "sqli", "data": self.parse_sqlmap_output(output)}
        elif "hydra" in cmd_lower or "medusa" in cmd_lower:
            return {"type": "bruteforce", "data": self.parse_hydra_output(output)}
        else:
            return {"type": "raw", "data": {"output": output[:5000]}}


# ---------------------------------------------------------------------------
# Agent Orchestrator — The Brain
# ---------------------------------------------------------------------------

class AgentOrchestrator:
    """
    The autonomous pentesting agent that orchestrates skill execution.

    This is the main class that:
    1. Plans the pentest based on target and engagement type
    2. Loads and sequences skills from the SkillLoader
    3. Generates AI prompts with full context
    4. Executes commands via CommandExecutor
    5. Parses output and updates engagement state
    6. Makes decisions about next steps via AI
    7. Generates the final report
    """

    def __init__(
        self,
        skills_dir: str = SKILLS_DIR,
        ai_provider=None,
        project=None,
        autopilot=None,
        on_status: Callable = None,
        on_output: Callable = None,
        on_finding: Callable = None,
        on_ask_permission: Callable = None,
        on_phase_change: Callable = None,
    ):
        """
        Args:
            skills_dir: Path to the skills/ directory
            ai_provider: AIProvider instance for LLM calls
            project: Project instance for engagement state
            autopilot: AutoPilot instance for execution control
            on_status: Callback for status updates (msg: str) -> None
            on_output: Callback for command output streaming (line: str) -> None
            on_finding: Callback when a finding is discovered (finding: dict) -> None
            on_ask_permission: Callback to ask user (cmd: str, reason: str) -> bool
            on_phase_change: Callback when phase changes (phase: str) -> None
        """
        # Core components
        self.registry = SkillRegistry(skills_dir)
        self.registry.load_all()
        self.parser = OutputParser()
        self.executor = CommandExecutor(
            autopilot=autopilot,
            on_output=on_output,
            on_ask_permission=on_ask_permission,
        )

        # External integrations
        self.ai_provider = ai_provider
        self.project = project
        self.autopilot = autopilot

        # Callbacks
        self.on_status = on_status or (lambda msg: None)
        self.on_output = on_output or (lambda line: None)
        self.on_finding = on_finding or (lambda finding: None)
        self.on_ask_permission = on_ask_permission
        self.on_phase_change = on_phase_change or (lambda phase: None)

        # State
        self.state: Optional[EngagementState] = None
        self._abort_requested = False

    # --- Public API ---

    def start_engagement(
        self,
        target: str,
        engagement_type: str = "web",
        variables: Dict[str, str] = None,
        workspace: str = None,
    ) -> EngagementState:
        """
        Start a new pentest engagement.

        Args:
            target: The target (domain, IP, or URL)
            engagement_type: "web", "internal", or "app"
            variables: Additional variables for skill templates
            workspace: Directory for output files (auto-created if None)
        """
        self._abort_requested = False

        # Validate engagement type
        if engagement_type not in ENGAGEMENT_PROFILES:
            raise ValueError(f"Unknown engagement type: {engagement_type}. "
                             f"Valid: {list(ENGAGEMENT_PROFILES.keys())}")

        profile = ENGAGEMENT_PROFILES[engagement_type]

        # Setup workspace
        if not workspace:
            safe_target = re.sub(r'[^\w.-]', '_', target)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            workspace = f"/tmp/mundix_pentest/{safe_target}_{timestamp}"
        os.makedirs(workspace, exist_ok=True)

        # Setup variables
        vars_dict = variables or {}
        # Auto-detect variable from target
        if "target_domain" not in vars_dict:
            vars_dict["target_domain"] = target
        if "target_ip" not in vars_dict:
            vars_dict["target_ip"] = target
        if "target_range" not in vars_dict:
            vars_dict["target_range"] = target
        if "target_url" not in vars_dict:
            if target.startswith("http"):
                vars_dict["target_url"] = target
            else:
                vars_dict["target_url"] = f"https://{target}"
        vars_dict["workspace"] = workspace

        # Create engagement state
        self.state = EngagementState(
            target=target,
            engagement_type=engagement_type,
            workspace=workspace,
            variables=vars_dict,
        )

        # Build execution plan
        plan = []
        for phase_id in profile["phases"]:
            skills = self.registry.get_by_phase(phase_id)
            for skill in skills:
                if skill.name not in profile.get("skip_skills", []):
                    plan.append({
                        "phase": phase_id,
                        "skill": skill.name,
                        "tools": skill.tools,
                        "commands": len(skill.commands),
                    })
                    # Initialize skill execution tracker
                    self.state.skill_executions[skill.name] = SkillExecution(
                        skill_name=skill.name
                    )

        self.state.plan = plan
        self.state.start()
        self.state.add_timeline_event("engagement_started", f"Target: {target}, Type: {engagement_type}")

        self.on_status(f"[ORCHESTRATOR] Engagement started: {target} ({engagement_type})")
        self.on_status(f"[ORCHESTRATOR] Plan: {len(plan)} skills across {len(profile['phases'])} phases")
        self.on_status(f"[ORCHESTRATOR] Workspace: {workspace}")

        return self.state

    def run_plan(self) -> EngagementState:
        """
        Execute the full pentest plan sequentially.
        This is the main autonomous loop.
        """
        if not self.state:
            raise RuntimeError("No engagement started. Call start_engagement() first.")

        current_phase = None

        for step in self.state.plan:
            if self._abort_requested:
                self.state.status = EngagementStatus.ABORTED
                self.state.add_timeline_event("aborted", "User requested abort")
                break

            phase_id = step["phase"]
            skill_name = step["skill"]

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

            # Execute skill
            self.state.current_skill = skill_name
            self.on_status(f"\n[SKILL] Starting: {skill_name}")

            try:
                self._execute_skill(skill_name)
            except Exception as e:
                self.on_status(f"[ERROR] Skill {skill_name} failed: {e}")
                self.state.skill_executions[skill_name].fail(str(e))
                self.state.add_timeline_event("skill_failed", f"{skill_name}: {e}")

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

        # Engagement complete
        if not self._abort_requested:
            self.state.complete()
            self.state.add_timeline_event("engagement_completed", "All phases finished")
            self.on_status(f"\n{'='*60}")
            self.on_status(f"[COMPLETE] Engagement finished!")
            self.on_status(f"[RESULTS] {json.dumps(self.state.get_progress(), indent=2)}")

        return self.state

    def run_single_skill(self, skill_name: str, variables: Dict[str, str] = None) -> SkillExecution:
        """Run a single skill outside of a full engagement."""
        skill = self.registry.get(skill_name)
        if not skill:
            raise ValueError(f"Skill not found: {skill_name}")

        # Create minimal state if needed
        if not self.state:
            target = (variables or {}).get("target_domain", "unknown")
            self.start_engagement(target, variables=variables)

        if variables:
            self.state.variables.update(variables)

        self._execute_skill(skill_name)
        return self.state.skill_executions.get(skill_name)

    def abort(self):
        """Request abort of the current engagement."""
        self._abort_requested = True
        self.on_status("[ORCHESTRATOR] Abort requested. Finishing current command...")

    def pause(self):
        """Pause the engagement."""
        if self.state:
            self.state.status = EngagementStatus.PAUSED
            self.state.add_timeline_event("paused", "Engagement paused by user")
            self.on_status("[ORCHESTRATOR] Engagement paused.")

    def resume(self):
        """Resume a paused engagement."""
        if self.state and self.state.status == EngagementStatus.PAUSED:
            self.state.status = EngagementStatus.RUNNING
            self.state.add_timeline_event("resumed", "Engagement resumed")
            self.on_status("[ORCHESTRATOR] Engagement resumed.")

    def get_progress(self) -> Dict:
        """Get current progress."""
        if self.state:
            return self.state.get_progress()
        return {"status": "idle"}

    # --- Internal Methods ---

    def _execute_skill(self, skill_name: str):
        """Execute a single skill with all its commands."""
        skill = self.registry.get(skill_name)
        if not skill:
            raise ValueError(f"Skill not found: {skill_name}")

        execution = self.state.skill_executions.get(skill_name)
        if not execution:
            execution = SkillExecution(skill_name=skill_name)
            self.state.skill_executions[skill_name] = execution

        execution.start()

        # Check dependencies
        if not self._check_dependencies(skill):
            execution.skip("Dependencies not met")
            self.on_status(f"[SKIP] {skill_name}: dependencies not met")
            return

        # Generate the AI prompt for this skill
        agent_prompt = skill.to_agent_prompt(self.state.variables)

        # Add engagement context
        context = self.state.get_context_summary()
        full_prompt = f"{agent_prompt}\n\n---\n\n{context}"

        # If we have an AI provider, let the AI decide which commands to run
        if self.ai_provider:
            commands_to_run = self._ai_plan_commands(skill, full_prompt)
        else:
            # Fallback: run all commands from the skill sequentially
            commands_to_run = [
                cmd.render(self.state.variables) if self.state.variables else cmd.raw
                for cmd in skill.commands
            ]

        # Execute each command
        for i, command in enumerate(commands_to_run):
            if self._abort_requested:
                break

            self.on_status(f"  [{i+1}/{len(commands_to_run)}] $ {command}")

            # Execute with permission check
            result, reason = self.executor.execute_with_permission(
                command,
                cwd=self.state.workspace,
            )

            if result is None:
                self.on_status(f"  [BLOCKED] {reason}")
                execution.ai_decisions.append(f"Command blocked: {reason}")
                continue

            # Log the execution
            execution.commands_executed.append(result.to_dict())

            if result.success:
                self.on_status(f"  [OK] Completed in {result.duration:.1f}s")
            else:
                self.on_status(f"  [FAIL] Exit code {result.returncode} ({result.duration:.1f}s)")

                # Try auto-fix if autopilot allows
                if self.autopilot and self.autopilot.should_auto_fix():
                    fixed_result = self._try_auto_fix(command, result)
                    if fixed_result and fixed_result.success:
                        result = fixed_result
                        self.on_status(f"  [AUTO-FIX] Succeeded!")

            # Parse output and update state
            self._process_output(skill, command, result)

        # Skill complete — let AI analyze results if available
        if self.ai_provider and not self._abort_requested:
            findings = self._ai_analyze_results(skill, execution)
            execution.complete(findings)
        else:
            execution.complete()

        self.state.add_timeline_event(
            "skill_completed",
            f"{skill_name}: {len(execution.commands_executed)} commands, "
            f"{len(execution.findings)} findings",
        )

    def _check_dependencies(self, skill: Skill) -> bool:
        """Check if a skill's dependencies are met."""
        for dep in skill.metadata.dependencies:
            # Special handling for common dependencies
            if dep == "target_domain_identified":
                if self.state.variables.get("target_domain"):
                    continue
                return False
            elif dep == "target_ip_identified":
                if self.state.variables.get("target_ip"):
                    continue
                return False
            elif dep == "open_ports_identified":
                if self.state.discovered_ports:
                    continue
                # Still allow if we haven't scanned yet
                return True
            elif dep in self.state.findings:
                if not self.state.findings[dep]:
                    return False
            # Unknown dependency — allow by default
        return True

    def _ai_plan_commands(self, skill: Skill, prompt: str) -> List[str]:
        """
        Ask the AI to plan which commands to execute for this skill.
        Falls back to skill's default commands if AI is unavailable.
        """
        try:
            planning_prompt = (
                f"You are the MundiX Pentesting Agent. Based on the skill below and the "
                f"current engagement state, output ONLY the bash commands to execute, "
                f"one per line. No explanations, no markdown, just raw commands.\n\n"
                f"Replace all placeholders with actual values from the context.\n"
                f"Skip commands that are not relevant to the current state.\n"
                f"Add any additional commands you think are necessary.\n\n"
                f"{prompt}"
            )

            # Use AI to get commands
            response = self.ai_provider.chat_sync(planning_prompt)

            # Extract commands from response
            commands = []
            for line in response.split('\n'):
                line = line.strip()
                # Skip empty lines, comments, and markdown
                if not line or line.startswith('#') or line.startswith('```'):
                    continue
                # Skip lines that look like explanations
                if any(line.startswith(w) for w in ['Note:', 'This ', 'The ', 'We ', 'I ']):
                    continue
                commands.append(line)

            if commands:
                return commands

        except Exception as e:
            self.on_status(f"  [AI] Planning failed: {e}. Using default commands.")

        # Fallback: use skill's default commands
        commands = []
        for cmd in skill.commands:
            try:
                rendered = cmd.render(self.state.variables)
                commands.append(rendered)
            except ValueError:
                commands.append(cmd.raw)
        return commands

    def _ai_analyze_results(self, skill: Skill, execution: SkillExecution) -> List[str]:
        """Ask the AI to analyze the results of a skill execution."""
        try:
            # Build context from execution results
            results_summary = []
            for cmd_result in execution.commands_executed[-5:]:  # Last 5 commands
                results_summary.append(
                    f"Command: {cmd_result['command']}\n"
                    f"Exit code: {cmd_result['returncode']}\n"
                    f"Output lines: {cmd_result['stdout_lines']}"
                )

            analysis_prompt = (
                f"Analyze the results of the '{skill.name}' skill execution.\n\n"
                f"## Skill Analysis Guide\n{skill.analysis_guide}\n\n"
                f"## Execution Results\n" + '\n---\n'.join(results_summary) + "\n\n"
                f"## Current Engagement State\n{self.state.get_context_summary()}\n\n"
                f"Output a JSON array of findings. Each finding should be a short string "
                f"describing what was discovered. Example:\n"
                f'["Port 80 open running Apache 2.4.41", "Port 22 open running OpenSSH 8.2"]\n'
                f"Output ONLY the JSON array, nothing else."
            )

            response = self.ai_provider.chat_sync(analysis_prompt)

            # Parse JSON findings
            try:
                # Extract JSON from response
                json_match = re.search(r'\[.*\]', response, re.DOTALL)
                if json_match:
                    findings = json.loads(json_match.group())
                    if isinstance(findings, list):
                        # Notify findings
                        for f in findings:
                            self.on_finding({"skill": skill.name, "finding": f})
                        return [str(f) for f in findings]
            except json.JSONDecodeError:
                pass

        except Exception as e:
            self.on_status(f"  [AI] Analysis failed: {e}")

        return []

    def _process_output(self, skill: Skill, command: str, result: CommandResult):
        """Parse command output and update engagement state."""
        if not result.success and not result.stdout:
            return

        parsed = self.parser.auto_parse(command, result.stdout)

        if parsed["type"] == "nmap":
            data = parsed["data"]
            for host in data.get("hosts", []):
                if host not in self.state.discovered_hosts:
                    self.state.discovered_hosts.append(host.get("ip", host.get("hostname", "")))
            for port in data.get("ports", []):
                self.state.discovered_ports.append(port)
                # Update variables for next skills
                if port["service"] == "http" or port["port"] in [80, 443, 8080, 8443]:
                    self.state.findings["web_service_found"] = True
                if port["service"] == "ssh" or port["port"] == 22:
                    self.state.findings["ssh_found"] = True
                if port["service"] in ["smb", "microsoft-ds"] or port["port"] in [139, 445]:
                    self.state.findings["smb_found"] = True

        elif parsed["type"] == "vuln_scan":
            for vuln in parsed["data"]:
                self.state.discovered_vulns.append(vuln)
                if vuln.get("severity") in ["critical", "high"]:
                    self.state.findings["critical_vuln_found"] = True

        elif parsed["type"] == "sqli":
            data = parsed["data"]
            if data.get("injectable"):
                self.state.findings["sqli_vulnerability_identified"] = True
                self.state.discovered_vulns.append({
                    "name": "SQL Injection",
                    "severity": "critical",
                    "parameters": data.get("parameters", []),
                    "dbms": data.get("dbms"),
                })

        elif parsed["type"] == "bruteforce":
            for cred in parsed["data"]:
                self.state.discovered_creds.append(cred)
                self.state.findings["credentials_found"] = True

    def _try_auto_fix(self, command: str, result: CommandResult) -> Optional[CommandResult]:
        """Try to automatically fix a failed command."""
        stderr = result.stderr.lower()

        # Command not found → try to install
        if result.returncode == 127 or "command not found" in stderr:
            cmd_name = command.split()[0]
            if cmd_name == "sudo" and len(command.split()) > 1:
                cmd_name = command.split()[1]

            # Try apt install
            install_result = self.executor.execute(
                f"apt-get install -y {cmd_name}",
                timeout=120,
            )
            if install_result.success:
                # Retry original command
                return self.executor.execute(command, cwd=self.state.workspace)

        # Permission denied → try with sudo
        if "permission denied" in stderr and not command.startswith("sudo"):
            return self.executor.execute(f"sudo {command}", cwd=self.state.workspace)

        return None

    # --- Report Generation ---

    def generate_report(self, output_format: str = "markdown") -> str:
        """Generate the final pentest report."""
        if not self.state:
            return "No engagement data available."

        lines = []
        lines.append(f"# MundiX Penetration Test Report")
        lines.append(f"**Generated:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        lines.append(f"**Target:** {self.state.target}")
        lines.append(f"**Type:** {self.state.engagement_type}")
        lines.append(f"**Duration:** {self.state.started_at} → {self.state.completed_at or 'ongoing'}")
        lines.append("")

        # Executive Summary
        progress = self.state.get_progress()
        lines.append(f"## Executive Summary")
        lines.append(f"This penetration test was conducted against **{self.state.target}** "
                      f"using the MundiX autonomous pentesting agent. "
                      f"A total of **{progress['completed']}** skills were executed across "
                      f"**{len(set(s['phase'] for s in self.state.plan))}** phases.")
        lines.append("")

        # Key Findings
        lines.append(f"## Key Findings")
        lines.append("")
        lines.append(f"| Category | Count |")
        lines.append(f"|---|---|")
        lines.append(f"| Hosts Discovered | {len(self.state.discovered_hosts)} |")
        lines.append(f"| Open Ports | {len(self.state.discovered_ports)} |")
        lines.append(f"| Vulnerabilities | {len(self.state.discovered_vulns)} |")
        lines.append(f"| Credentials | {len(self.state.discovered_creds)} |")
        lines.append("")

        # Vulnerabilities
        if self.state.discovered_vulns:
            lines.append(f"## Vulnerabilities")
            lines.append("")
            lines.append(f"| Severity | Name | Details |")
            lines.append(f"|---|---|---|")
            for vuln in self.state.discovered_vulns:
                severity = vuln.get("severity", "info").upper()
                name = vuln.get("name", vuln.get("template", "Unknown"))
                details = vuln.get("url", vuln.get("raw", ""))[:80]
                lines.append(f"| {severity} | {name} | {details} |")
            lines.append("")

        # Open Ports
        if self.state.discovered_ports:
            lines.append(f"## Open Ports")
            lines.append("")
            lines.append(f"| Host | Port | Protocol | Service | Version |")
            lines.append(f"|---|---|---|---|---|")
            for port in self.state.discovered_ports[:50]:
                lines.append(
                    f"| {port.get('host', '?')} | {port.get('port', '?')} | "
                    f"{port.get('protocol', '?')} | {port.get('service', '?')} | "
                    f"{port.get('version', '')} |"
                )
            lines.append("")

        # Credentials
        if self.state.discovered_creds:
            lines.append(f"## Discovered Credentials")
            lines.append("")
            lines.append(f"| Service | Host | Username | Password |")
            lines.append(f"|---|---|---|---|")
            for cred in self.state.discovered_creds:
                lines.append(
                    f"| {cred.get('service', '?')} | {cred.get('host', '?')} | "
                    f"{cred.get('username', '?')} | {cred.get('password', '***')} |"
                )
            lines.append("")

        # Skill Execution Details
        lines.append(f"## Skill Execution Details")
        lines.append("")
        lines.append(f"| Skill | Status | Commands | Findings | Duration |")
        lines.append(f"|---|---|---|---|---|")
        for name, ex in self.state.skill_executions.items():
            duration = ""
            if ex.started_at and ex.completed_at:
                try:
                    start = datetime.fromisoformat(ex.started_at)
                    end = datetime.fromisoformat(ex.completed_at)
                    duration = f"{(end - start).total_seconds():.0f}s"
                except (ValueError, TypeError):
                    pass
            lines.append(
                f"| {name} | {ex.status.value} | {len(ex.commands_executed)} | "
                f"{len(ex.findings)} | {duration} |"
            )
        lines.append("")

        # Timeline
        lines.append(f"## Timeline")
        lines.append("")
        for event in self.state.timeline:
            lines.append(f"- **{event['timestamp']}** [{event['type']}] {event['description']}")
        lines.append("")

        report = '\n'.join(lines)

        # Save report
        if self.state.workspace:
            report_path = os.path.join(self.state.workspace, "pentest_report.md")
            with open(report_path, 'w', encoding='utf-8') as f:
                f.write(report)
            self.on_status(f"[REPORT] Saved to: {report_path}")

        return report


# ---------------------------------------------------------------------------
# CLI Entry Point (for testing)
# ---------------------------------------------------------------------------

def main():
    """CLI entry point for testing the orchestrator."""
    import sys

    skills_dir = sys.argv[1] if len(sys.argv) > 1 else SKILLS_DIR
    target = sys.argv[2] if len(sys.argv) > 2 else "example.com"

    def status_callback(msg):
        print(msg)

    def output_callback(line):
        print(f"    {line}")

    def finding_callback(finding):
        print(f"  [FINDING] {finding}")

    orchestrator = AgentOrchestrator(
        skills_dir=skills_dir,
        on_status=status_callback,
        on_output=output_callback,
        on_finding=finding_callback,
    )

    # Start engagement
    state = orchestrator.start_engagement(target, engagement_type="web")

    # Print plan
    print(f"\n{'='*60}")
    print(f"PENTEST PLAN ({len(state.plan)} skills)")
    print(f"{'='*60}")
    for step in state.plan:
        print(f"  [{step['phase']}] {step['skill']} ({step['commands']} cmds)")

    # Print progress
    print(f"\nProgress: {json.dumps(state.get_progress(), indent=2)}")

    # Print engagement profiles
    print(f"\nAvailable Engagement Types:")
    for name, profile in ENGAGEMENT_PROFILES.items():
        print(f"  {name}: {profile['description']} ({len(profile['phases'])} phases)")

    print(f"\n[OK] Orchestrator initialized successfully.")
    print(f"     Skills: {len(orchestrator.registry.get_all())}")
    print(f"     Plan steps: {len(state.plan)}")
    print(f"     Workspace: {state.workspace}")


if __name__ == "__main__":
    main()
