"""
Data models for penetration testing plans.

This module defines the structure of plans and plan steps that are
generated by AI and executed by the wish system.
"""

import logging
from dataclasses import dataclass, field
from enum import Enum

logger = logging.getLogger(__name__)


class RiskLevel(str, Enum):
    """Risk levels for plan steps."""

    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"


class StepStatus(str, Enum):
    """Execution status of plan steps."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    SKIPPED = "skipped"


@dataclass
class PlanStep:
    """A single step in a penetration testing plan.

    Represents one actionable step that can be executed as part of
    a larger penetration testing workflow.
    """

    tool_name: str
    """Name of the tool to execute (e.g., 'nmap', 'gobuster', 'nikto')"""

    command: str
    """Exact command to execute"""

    purpose: str
    """Human-readable explanation of what this step accomplishes"""

    expected_result: str
    """What results we expect from this command"""

    risk_level: RiskLevel = RiskLevel.LOW
    """Risk level of executing this step"""

    status: StepStatus = StepStatus.PENDING
    """Current execution status"""

    timeout: int | None = None
    """Timeout in seconds (None for no timeout)"""

    requires_confirmation: bool = False
    """Whether this step requires user confirmation before execution"""

    prerequisites: list[str] = field(default_factory=list)
    """List of step IDs that must complete before this step"""

    metadata: dict[str, str] = field(default_factory=dict)
    """Additional metadata for the step"""

    def __post_init__(self) -> None:
        """Validate the plan step after initialization."""
        if not self.tool_name.strip():
            raise ValueError("tool_name cannot be empty")

        if not self.command.strip():
            raise ValueError("command cannot be empty")

        if not self.purpose.strip():
            raise ValueError("purpose cannot be empty")

        # Set high-risk commands to require confirmation
        high_risk_tools = {"metasploit", "msfconsole", "sqlmap", "hydra", "medusa"}
        if self.tool_name.lower() in high_risk_tools and self.risk_level in (RiskLevel.HIGH, RiskLevel.CRITICAL):
            self.requires_confirmation = True

    def is_ready_to_execute(self, completed_steps: list[str]) -> bool:
        """Check if this step is ready to execute based on prerequisites.

        Args:
            completed_steps: List of step IDs that have been completed

        Returns:
            True if all prerequisites are met
        """
        return all(prereq in completed_steps for prereq in self.prerequisites)

    def to_dict(self) -> dict[str, str]:
        """Convert to dictionary for serialization."""
        return {
            "tool_name": self.tool_name,
            "command": self.command,
            "purpose": self.purpose,
            "expected_result": self.expected_result,
            "risk_level": self.risk_level.value,
            "status": self.status.value,
            "timeout": str(self.timeout) if self.timeout else "",
            "requires_confirmation": str(self.requires_confirmation),
            "prerequisites": ",".join(self.prerequisites),
            **self.metadata,
        }


@dataclass
class Plan:
    """A complete penetration testing execution plan.

    Contains multiple steps that together accomplish a penetration
    testing objective based on the current engagement state.
    """

    description: str
    """High-level description of what this plan accomplishes"""

    steps: list[PlanStep]
    """Ordered list of steps to execute"""

    rationale: str
    """Explanation of why this plan was chosen"""

    estimated_duration: int | None = None
    """Estimated execution time in minutes"""

    target_scope: list[str] = field(default_factory=list)
    """List of targets this plan will interact with"""

    mode: str = "recon"
    """Engagement mode this plan is designed for"""

    priority: int = 1
    """Priority level (1=highest, 5=lowest)"""

    created_by: str = "AI"
    """Who/what created this plan"""

    metadata: dict[str, str] = field(default_factory=dict)
    """Additional metadata for the plan"""

    def __post_init__(self) -> None:
        """Validate the plan after initialization."""
        if not self.description.strip():
            raise ValueError("description cannot be empty")

        # Allow empty steps for knowledge-based responses
        # if not self.steps:
        #     raise ValueError("plan must have at least one step")

        if not self.rationale.strip():
            raise ValueError("rationale cannot be empty")

        # Validate all steps
        for i, step in enumerate(self.steps):
            if not isinstance(step, PlanStep):
                raise ValueError(f"step {i} must be a PlanStep instance")

    @property
    def total_steps(self) -> int:
        """Get the total number of steps in the plan."""
        return len(self.steps)

    @property
    def high_risk_steps(self) -> list[PlanStep]:
        """Get all high-risk steps that require special attention."""
        return [step for step in self.steps if step.risk_level in (RiskLevel.HIGH, RiskLevel.CRITICAL)]

    @property
    def completion_percentage(self) -> float:
        """Calculate completion percentage based on step statuses."""
        if not self.steps:
            return 0.0

        completed = sum(1 for step in self.steps if step.status == StepStatus.COMPLETED)
        return (completed / len(self.steps)) * 100.0

    def get_next_step(self) -> PlanStep | None:
        """Get the next step that is ready to execute.

        Returns:
            Next executable step or None if no steps are ready
        """
        completed_step_ids = [f"step_{i}" for i, step in enumerate(self.steps) if step.status == StepStatus.COMPLETED]

        for step in self.steps:
            if step.status == StepStatus.PENDING and step.is_ready_to_execute(completed_step_ids):
                return step

        return None

    def mark_step_completed(self, step_index: int) -> None:
        """Mark a specific step as completed.

        Args:
            step_index: Index of the step to mark as completed

        Raises:
            IndexError: If step_index is out of range
        """
        if 0 <= step_index < len(self.steps):
            self.steps[step_index].status = StepStatus.COMPLETED
            logger.info(f"Marked step {step_index} as completed")
        else:
            raise IndexError(f"Step index {step_index} out of range")

    def mark_step_failed(self, step_index: int) -> None:
        """Mark a specific step as failed.

        Args:
            step_index: Index of the step to mark as failed

        Raises:
            IndexError: If step_index is out of range
        """
        if 0 <= step_index < len(self.steps):
            self.steps[step_index].status = StepStatus.FAILED
            logger.warning(f"Marked step {step_index} as failed")
        else:
            raise IndexError(f"Step index {step_index} out of range")

    def to_dict(self) -> dict:
        """Convert to dictionary for serialization."""
        return {
            "description": self.description,
            "rationale": self.rationale,
            "estimated_duration": str(self.estimated_duration) if self.estimated_duration else "",
            "target_scope": ",".join(self.target_scope),
            "mode": self.mode,
            "priority": str(self.priority),
            "created_by": self.created_by,
            "total_steps": str(self.total_steps),
            "completion_percentage": f"{self.completion_percentage:.1f}%",
            "steps": [step.to_dict() for step in self.steps],
            **self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "Plan":
        """Create Plan from dictionary."""
        # Handle steps - convert dict to PlanStep objects if needed
        steps = []
        if "steps" in data:
            for step_data in data["steps"]:
                if isinstance(step_data, dict):
                    step = PlanStep(
                        tool_name=step_data.get("tool_name", ""),
                        command=step_data.get("command", ""),
                        purpose=step_data.get("purpose", ""),
                        expected_result=step_data.get("expected_result", ""),
                        risk_level=step_data.get("risk_level", RiskLevel.LOW),
                        timeout=step_data.get("timeout"),
                        requires_confirmation=step_data.get("requires_confirmation", False),
                        prerequisites=step_data.get("prerequisites", []),
                        metadata=step_data.get("metadata", {}),
                    )
                else:
                    step = step_data
                steps.append(step)

        # Create plan with basic required fields
        return cls(
            description=data.get("description", ""),
            steps=steps,
            rationale=data.get("rationale", ""),
            estimated_duration=data.get("estimated_duration"),
            target_scope=data.get("target_scope", [])
            if isinstance(data.get("target_scope"), list)
            else data.get("target_scope", "").split(",")
            if data.get("target_scope")
            else [],
            mode=data.get("mode", "recon"),
            priority=int(data.get("priority", 1)) if str(data.get("priority", 1)).isdigit() else 1,
            created_by=data.get("created_by", "AI"),
            metadata=data.get("metadata", {}),
        )
