#!/usr/bin/env python3
"""
souleyez.core.next_steps - Context-aware next steps suggestion engine

Generates manual and automated next step suggestions based on:
- Current job findings
- Other tools' discoveries in the engagement
- Already queued/running/completed jobs
- Available credentials
- Auto-chain rules that will trigger automatically
"""

import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional, Set

logger = logging.getLogger(__name__)


class StepType(Enum):
    """Type of next step suggestion."""

    MANUAL = "manual"  # User must run manually
    AUTO = "auto"  # Will be auto-chained
    INFO = "info"  # Informational/advisory


@dataclass
class NextStep:
    """A suggested next step based on job findings."""

    title: str  # "Try SQL injection on login form"
    reason: str  # "Login page detected with POST form"
    commands: List[str] = field(default_factory=list)  # ["sqlmap -u ..."]
    step_type: StepType = StepType.MANUAL
    priority: int = 3  # 1-5 (1=highest)
    category: str = ""  # "sqli", "auth", "privesc", etc.
    target: str = ""  # URL or IP this applies to
    requires_creds: bool = False  # Needs credentials to execute
    cred_type: Optional[str] = None  # "mysql", "ssh", "smb", etc.
    tool: str = ""  # Suggested tool name
    auto_job_id: Optional[int] = None  # If AUTO, the queued job ID


@dataclass
class NextStepsContext:
    """Context gathered from the engagement for smart filtering."""

    engagement_id: int
    all_jobs: List[Dict[str, Any]] = field(default_factory=list)
    findings: List[Dict[str, Any]] = field(default_factory=list)
    credentials: List[Dict[str, Any]] = field(default_factory=list)
    queued_tools: Set[str] = field(default_factory=set)  # tool|target combos
    completed_scans: Dict[str, Set[str]] = field(
        default_factory=dict
    )  # tool -> set of targets
    pending_targets: Dict[str, Set[str]] = field(
        default_factory=dict
    )  # tool -> set of targets


class NextStepsEngine:
    """
    Engine for generating context-aware next step suggestions.

    Integrates with:
    - Job queue (to know what's already running/queued)
    - Findings database (to know what's been discovered)
    - Credentials database (to suggest credential-based attacks)
    - Tool chaining rules (to mark auto-triggered steps)
    """

    def __init__(
        self,
        findings_manager=None,
        credentials_manager=None,
    ):
        """
        Initialize the NextStepsEngine.

        Args:
            findings_manager: Optional FindingsManager instance
            credentials_manager: Optional CredentialsManager instance
        """
        self.findings_manager = findings_manager
        self.credentials_manager = credentials_manager

    def generate_steps(
        self,
        job: Dict[str, Any],
        engagement_id: int,
        handler=None,
    ) -> List[NextStep]:
        """
        Generate next step suggestions for a completed job.

        Args:
            job: The completed job dict
            engagement_id: The engagement ID
            handler: Optional handler instance for tool-specific suggestions

        Returns:
            List of NextStep suggestions, filtered and prioritized
        """
        # Gather context
        context = self._gather_context(engagement_id)

        # Get tool-specific suggestions from handler
        steps = []
        if handler and hasattr(handler, "generate_next_steps"):
            try:
                handler_steps = handler.generate_next_steps(job, context)
                if handler_steps:
                    steps.extend(handler_steps)
            except Exception as e:
                logger.warning(f"Handler generate_next_steps failed: {e}")

        # Apply cross-tool intelligence
        steps = self._apply_cross_tool_intelligence(steps, context, job)

        # Filter steps that will be auto-chained
        steps = self._filter_auto_chained(steps, job)

        # Filter already-handled steps
        steps = self._filter_already_handled(steps, context)

        # Sort by priority
        steps.sort(key=lambda s: s.priority)

        return steps

    def get_steps_for_display(
        self,
        job: Dict[str, Any],
        engagement_id: int,
        handler=None,
        max_steps: int = 10,
    ) -> str:
        """
        Generate formatted next steps output for display.

        Args:
            job: The completed job dict
            engagement_id: The engagement ID
            handler: Optional handler instance
            max_steps: Maximum steps to show

        Returns:
            Formatted string for display
        """
        steps = self.generate_steps(job, engagement_id, handler)

        if not steps:
            return ""

        lines = []
        lines.append("")
        lines.append("=" * 70)
        lines.append("SUGGESTED NEXT STEPS")
        lines.append("=" * 70)
        lines.append("")

        for step in steps[:max_steps]:
            # Step type indicator
            if step.step_type == StepType.MANUAL:
                type_str = "[MANUAL]"
            elif step.step_type == StepType.AUTO:
                type_str = "[AUTO]"
            else:
                type_str = "[INFO]"

            # Priority indicator
            priority_map = {1: "CRITICAL", 2: "HIGH", 3: "MEDIUM", 4: "LOW", 5: "INFO"}
            priority_str = priority_map.get(step.priority, "MEDIUM")

            lines.append(f"{type_str} {step.title} (Priority: {priority_str})")
            lines.append(f"  Reason: {step.reason}")

            if step.step_type == StepType.AUTO and step.auto_job_id:
                lines.append(f"  Queued as Job #{step.auto_job_id}")
            elif step.commands:
                lines.append("  Commands:")
                for cmd in step.commands[:3]:  # Show max 3 commands
                    lines.append(f"    {cmd}")

            if step.requires_creds:
                cred_note = f" ({step.cred_type})" if step.cred_type else ""
                lines.append(f"  Note: Requires credentials{cred_note}")

            lines.append("")

        if len(steps) > max_steps:
            lines.append(f"  ... and {len(steps) - max_steps} more suggestions")
            lines.append("")

        lines.append("=" * 70)

        return "\n".join(lines)

    def _gather_context(self, engagement_id: int) -> NextStepsContext:
        """
        Gather engagement context for filtering and intelligence.

        Args:
            engagement_id: The engagement ID

        Returns:
            NextStepsContext with all relevant data
        """
        context = NextStepsContext(engagement_id=engagement_id)

        # Import here to avoid circular imports
        try:
            from souleyez.engine.background import get_all_jobs
        except ImportError:
            logger.warning("Could not import get_all_jobs")
            return context

        # Get all jobs for this engagement
        all_jobs = get_all_jobs()
        context.all_jobs = [
            j for j in all_jobs if j.get("engagement_id") == engagement_id
        ]

        # Build queued/completed maps
        for job in context.all_jobs:
            tool = job.get("tool", "")
            target = job.get("target", "")
            status = job.get("status", "")

            tool_target_key = f"{tool}|{target}"

            if status in ("pending", "running", "queued"):
                context.queued_tools.add(tool_target_key)
                if tool not in context.pending_targets:
                    context.pending_targets[tool] = set()
                context.pending_targets[tool].add(target)
            elif status in ("done", "warning"):
                if tool not in context.completed_scans:
                    context.completed_scans[tool] = set()
                context.completed_scans[tool].add(target)

        # Get findings
        if self.findings_manager:
            try:
                context.findings = self.findings_manager.list_findings(engagement_id)
            except Exception as e:
                logger.warning(f"Could not get findings: {e}")

        # Get credentials
        if self.credentials_manager:
            try:
                context.credentials = self.credentials_manager.list_credentials(
                    engagement_id, decrypt=False
                )
            except Exception as e:
                logger.warning(f"Could not get credentials: {e}")

        return context

    def _filter_auto_chained(
        self,
        steps: List[NextStep],
        job: Dict[str, Any],
    ) -> List[NextStep]:
        """
        Filter out suggestions that will be auto-triggered by chain rules.

        If a chain rule exists that will trigger the same tool, we shouldn't
        show it as a manual suggestion.

        Args:
            steps: List of suggested steps
            job: The current job that generated these suggestions

        Returns:
            Filtered list excluding auto-chained tools
        """
        try:
            from souleyez.core.tool_chaining import ToolChaining
        except ImportError:
            logger.warning("Could not import ChainEngine for filtering")
            return steps

        current_tool = job.get("tool", "")

        # Get all tools that will be auto-triggered from the current tool
        # This includes direct chains and indirect chains (e.g., nmap -> http_fingerprint -> nikto)
        auto_triggered_tools = self._get_auto_triggered_tools(current_tool)

        if not auto_triggered_tools:
            return steps

        filtered = []
        for step in steps:
            # If this step's tool is in the auto-triggered set, skip it
            if step.tool and step.tool.lower() in auto_triggered_tools:
                logger.debug(
                    f"Filtering out '{step.title}' - {step.tool} is auto-chained from {current_tool}"
                )
                continue
            filtered.append(step)

        return filtered

    def _get_auto_triggered_tools(self, from_tool: str, depth: int = 0) -> Set[str]:
        """
        Get all tools that will be auto-triggered from a given tool.

        Follows chain rules recursively to find all downstream tools.

        Args:
            from_tool: The source tool
            depth: Recursion depth (to prevent infinite loops)

        Returns:
            Set of tool names that will be auto-triggered
        """
        if depth > 5:  # Prevent infinite recursion
            return set()

        triggered = set()

        try:
            from souleyez.core.tool_chaining import ToolChaining

            engine = ToolChaining()

            # Find all rules that trigger from this tool
            for rule in engine.rules:
                if rule.trigger_tool == from_tool and rule.enabled:
                    target_tool = rule.target_tool.lower()
                    triggered.add(target_tool)

                    # Recursively get tools triggered by the target
                    downstream = self._get_auto_triggered_tools(target_tool, depth + 1)
                    triggered.update(downstream)

        except Exception as e:
            logger.debug(f"Error getting auto-triggered tools: {e}")

        return triggered

    def _filter_already_handled(
        self,
        steps: List[NextStep],
        context: NextStepsContext,
    ) -> List[NextStep]:
        """
        Filter out steps that are already queued or completed.

        Args:
            steps: List of suggested steps
            context: Engagement context

        Returns:
            Filtered list of steps
        """
        filtered = []

        for step in steps:
            tool_target_key = f"{step.tool}|{step.target}"

            # Check if already queued
            if tool_target_key in context.queued_tools:
                # Mark as AUTO if queued, find the job ID
                for job in context.all_jobs:
                    if (
                        job.get("tool") == step.tool
                        and job.get("target") == step.target
                        and job.get("status") in ("pending", "running", "queued")
                    ):
                        step.step_type = StepType.AUTO
                        step.auto_job_id = job.get("id")
                        filtered.append(step)
                        break
                continue

            # Check if already completed (skip unless different args matter)
            completed_targets = context.completed_scans.get(step.tool, set())
            if step.target in completed_targets:
                # Already scanned this target with this tool
                continue

            filtered.append(step)

        return filtered

    def _apply_cross_tool_intelligence(
        self,
        steps: List[NextStep],
        context: NextStepsContext,
        current_job: Dict[str, Any],
    ) -> List[NextStep]:
        """
        Apply cross-tool intelligence to boost/add suggestions.

        Args:
            steps: List of suggested steps
            context: Engagement context
            current_job: The current job that generated these steps

        Returns:
            Enhanced list of steps
        """
        # Build a quick lookup of what we have
        has_mysql_creds = any(
            c.get("service") == "mysql" and c.get("status") == "valid"
            for c in context.credentials
        )
        has_smb_creds = any(
            c.get("service") in ("smb", "cifs") and c.get("status") == "valid"
            for c in context.credentials
        )
        has_ssh_creds = any(
            c.get("service") == "ssh" and c.get("status") == "valid"
            for c in context.credentials
        )
        has_kerberos_users = any(
            f.get("finding_type") == "kerberos_user" for f in context.findings
        )
        has_admin_panel = any(
            "admin" in (f.get("path") or "").lower() for f in context.findings
        )

        for step in steps:
            # Boost priority if we have related credentials
            if step.cred_type == "mysql" and has_mysql_creds:
                step.priority = max(1, step.priority - 1)
                step.reason += " (MySQL credentials available)"

            if step.cred_type == "smb" and has_smb_creds:
                step.priority = max(1, step.priority - 1)
                step.reason += " (SMB credentials available)"

            if step.cred_type == "ssh" and has_ssh_creds:
                step.priority = max(1, step.priority - 1)
                step.reason += " (SSH credentials available)"

            # Boost SQLi suggestions if admin panel found
            if step.category == "sqli" and has_admin_panel:
                step.priority = max(1, step.priority - 1)

            # Boost Kerberos attacks if users found
            if step.category == "kerberos" and has_kerberos_users:
                step.priority = max(1, step.priority - 1)

        return steps


# Convenience function for use in handlers
def create_next_step(
    title: str,
    reason: str,
    commands: Optional[List[str]] = None,
    step_type: StepType = StepType.MANUAL,
    priority: int = 3,
    category: str = "",
    target: str = "",
    requires_creds: bool = False,
    cred_type: Optional[str] = None,
    tool: str = "",
) -> NextStep:
    """
    Convenience function to create a NextStep.

    Args:
        title: Short description of what to try
        reason: Why this step is suggested
        commands: Example commands to run
        step_type: MANUAL, AUTO, or INFO
        priority: 1-5 (1=highest)
        category: Category like "sqli", "auth", "privesc"
        target: URL or IP this applies to
        requires_creds: Whether credentials are needed
        cred_type: Type of credentials needed
        tool: Suggested tool name

    Returns:
        NextStep instance
    """
    return NextStep(
        title=title,
        reason=reason,
        commands=commands or [],
        step_type=step_type,
        priority=priority,
        category=category,
        target=target,
        requires_creds=requires_creds,
        cred_type=cred_type,
        tool=tool,
    )
