"""
AIPTX Beast Mode - Exploit Chain Builder
=========================================

LLM-powered multi-step attack chain planning. Analyzes discovered vulnerabilities
and constructs optimal exploitation paths to achieve maximum impact.

Key Features:
- Automatic chain selection based on vulnerability type
- LLM-powered custom chain planning
- Dependency resolution between steps
- Fallback path generation
- Risk assessment for each chain
"""

from __future__ import annotations

import json
import logging
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Optional

import litellm

from aipt_v2.exploitation.chain_templates import (
    CHAIN_TEMPLATES,
    ChainTemplate,
    ChainStepTemplate,
    StepType,
    get_chains_for_vuln_type,
)

logger = logging.getLogger(__name__)


class ChainStatus(str, Enum):
    """Status of an exploit chain."""
    PLANNED = "planned"
    IN_PROGRESS = "in_progress"
    STEP_FAILED = "step_failed"
    COMPLETED = "completed"
    ABORTED = "aborted"


class StepStatus(str, Enum):
    """Status of a chain step."""
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    SUCCESS = "success"
    FAILED = "failed"
    SKIPPED = "skipped"


@dataclass
class ChainStep:
    """
    A single step in an exploit chain.

    Each step represents one action in the multi-step exploitation,
    with its own techniques, conditions, and outputs.
    """
    step_id: str
    name: str
    step_type: StepType
    description: str
    techniques: list[str]
    required_conditions: list[str] = field(default_factory=list)
    produces: list[str] = field(default_factory=list)
    status: StepStatus = StepStatus.PENDING
    selected_technique: str | None = None
    output: dict[str, Any] = field(default_factory=dict)
    error: str | None = None
    started_at: str | None = None
    completed_at: str | None = None
    fallback_techniques: list[str] = field(default_factory=list)
    retry_count: int = 0
    max_retries: int = 2

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "step_id": self.step_id,
            "name": self.name,
            "step_type": self.step_type.value,
            "description": self.description,
            "techniques": self.techniques,
            "required_conditions": self.required_conditions,
            "produces": self.produces,
            "status": self.status.value,
            "selected_technique": self.selected_technique,
            "output": self.output,
            "error": self.error,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
        }

    @classmethod
    def from_template(cls, template: ChainStepTemplate) -> "ChainStep":
        """Create a ChainStep from a template."""
        return cls(
            step_id=str(uuid.uuid4())[:8],
            name=template.name,
            step_type=template.step_type,
            description=template.description,
            techniques=template.techniques.copy(),
            required_conditions=template.required_conditions.copy(),
            produces=template.produces.copy(),
            fallback_techniques=template.fallback_techniques.copy(),
        )


@dataclass
class AttackChain:
    """
    A complete exploit chain with multiple steps.

    Represents the full exploitation path from initial vulnerability
    to final objective (RCE, data exfil, privesc, etc.).
    """
    chain_id: str
    name: str
    description: str
    target: str
    entry_vuln: dict[str, Any]  # The vulnerability that starts this chain
    steps: list[ChainStep]
    status: ChainStatus = ChainStatus.PLANNED
    current_step_index: int = 0
    context: dict[str, Any] = field(default_factory=dict)  # Accumulated data
    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
    started_at: str | None = None
    completed_at: str | None = None
    estimated_impact: str = ""
    stealth_rating: int = 5
    risk_score: float = 0.0

    def current_step(self) -> ChainStep | None:
        """Get the current step to execute."""
        if 0 <= self.current_step_index < len(self.steps):
            return self.steps[self.current_step_index]
        return None

    def advance_step(self) -> ChainStep | None:
        """Advance to the next step."""
        self.current_step_index += 1
        return self.current_step()

    def get_available_conditions(self) -> list[str]:
        """Get all conditions produced by completed steps."""
        conditions = []
        for step in self.steps:
            if step.status == StepStatus.SUCCESS:
                conditions.extend(step.produces)
        # Add context data as conditions
        conditions.extend(self.context.keys())
        return conditions

    def check_step_conditions(self, step: ChainStep) -> bool:
        """Check if step's required conditions are met."""
        available = self.get_available_conditions()
        return all(cond in available for cond in step.required_conditions)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "chain_id": self.chain_id,
            "name": self.name,
            "description": self.description,
            "target": self.target,
            "entry_vuln": self.entry_vuln,
            "steps": [s.to_dict() for s in self.steps],
            "status": self.status.value,
            "current_step_index": self.current_step_index,
            "context": self.context,
            "created_at": self.created_at,
            "estimated_impact": self.estimated_impact,
            "stealth_rating": self.stealth_rating,
            "risk_score": self.risk_score,
        }


# LLM prompt for custom chain planning
CHAIN_PLANNER_PROMPT = """You are an expert penetration tester planning a multi-step exploitation chain.

Given the discovered vulnerability and target information, plan an optimal attack chain
to achieve the maximum impact while considering stealth and reliability.

VULNERABILITY DETAILS:
{vuln_details}

TARGET INFORMATION:
{target_info}

OBJECTIVE:
{objective}

AVAILABLE CHAIN TEMPLATES:
{templates}

Your task:
1. Analyze the vulnerability and determine the best exploitation path
2. Select or customize a chain template
3. Order the steps for maximum success probability
4. Identify fallback options for each step

Respond with a JSON object:
{{
    "chain_name": "descriptive_chain_name",
    "description": "What this chain accomplishes",
    "steps": [
        {{
            "name": "step_name",
            "step_type": "exploit|recon|escalation|exfil|persist",
            "description": "What this step does",
            "techniques": ["technique1", "technique2"],
            "required_conditions": ["condition1"],
            "produces": ["output1", "output2"],
            "fallback_techniques": ["fallback1"]
        }}
    ],
    "estimated_impact": "Critical/High/Medium/Low",
    "stealth_rating": 1-10,
    "reasoning": "Why this chain was chosen"
}}

IMPORTANT:
- Each step must have clear conditions and outputs
- Include fallback techniques for critical steps
- Consider detection risk at each step
- Output valid JSON only"""


class ExploitChainBuilder:
    """
    Builds multi-step exploit chains based on discovered vulnerabilities.

    Uses templates for common chains and LLM for custom chain planning.
    """

    def __init__(
        self,
        model: str | None = None,
        api_key: str | None = None,
        api_base: str | None = None,
    ):
        """
        Initialize the chain builder.

        Args:
            model: LLM model for custom chain planning
            api_key: Optional API key override
            api_base: Optional API base URL override
        """
        self.model = model or "claude-3-haiku-20240307"
        self.api_key = api_key
        self.api_base = api_base
        self._chain_cache: dict[str, AttackChain] = {}

    def build_from_template(
        self,
        template_name: str,
        target: str,
        entry_vuln: dict[str, Any],
        context: dict[str, Any] | None = None,
    ) -> AttackChain | None:
        """
        Build an exploit chain from a predefined template.

        Args:
            template_name: Name of the template to use
            target: Target URL or identifier
            entry_vuln: The vulnerability that starts this chain
            context: Additional context data

        Returns:
            AttackChain or None if template not found
        """
        template = CHAIN_TEMPLATES.get(template_name)
        if not template:
            logger.warning(f"Chain template not found: {template_name}")
            return None

        # Create steps from template
        steps = [ChainStep.from_template(step_template) for step_template in template.steps]

        chain = AttackChain(
            chain_id=str(uuid.uuid4())[:8],
            name=template.name,
            description=template.description,
            target=target,
            entry_vuln=entry_vuln,
            steps=steps,
            context=context or {},
            estimated_impact=template.estimated_impact,
            stealth_rating=template.stealth_rating,
            risk_score=self._calculate_risk_score(template),
        )

        # Add entry vulnerability data to context
        chain.context["entry_vuln_type"] = entry_vuln.get("type", "unknown")
        chain.context["entry_vuln_id"] = entry_vuln.get("id", "")

        self._chain_cache[chain.chain_id] = chain
        logger.info(f"Built chain '{chain.name}' with {len(steps)} steps")

        return chain

    def build_from_vulnerability(
        self,
        vuln_type: str,
        target: str,
        entry_vuln: dict[str, Any],
        objective: str = "rce",
        context: dict[str, Any] | None = None,
    ) -> AttackChain | None:
        """
        Automatically select and build a chain based on vulnerability type.

        Args:
            vuln_type: Type of vulnerability (sqli, xss, ssrf, etc.)
            target: Target URL or identifier
            entry_vuln: The vulnerability that starts this chain
            objective: Desired outcome (rce, data_exfil, session_hijack)
            context: Additional context data

        Returns:
            Best matching AttackChain or None
        """
        # Get matching templates
        matching_chains = get_chains_for_vuln_type(vuln_type)

        if not matching_chains:
            logger.warning(f"No chain templates found for vuln type: {vuln_type}")
            return None

        # Filter by objective if specified
        objective_keywords = {
            "rce": ["rce", "shell", "command"],
            "data_exfil": ["data", "exfil", "dump"],
            "session": ["session", "hijack", "token"],
            "privesc": ["root", "privilege", "escalation"],
        }

        if objective in objective_keywords:
            keywords = objective_keywords[objective]
            filtered = [
                chain for chain in matching_chains
                if any(kw in chain.name.lower() or kw in chain.description.lower() for kw in keywords)
            ]
            if filtered:
                matching_chains = filtered

        # Select the best template (first one after sorting)
        selected_template = matching_chains[0]

        return self.build_from_template(
            template_name=selected_template.name,
            target=target,
            entry_vuln=entry_vuln,
            context=context,
        )

    async def plan_custom_chain(
        self,
        entry_vuln: dict[str, Any],
        target_info: dict[str, Any],
        objective: str = "maximum_impact",
    ) -> AttackChain | None:
        """
        Use LLM to plan a custom exploit chain.

        Args:
            entry_vuln: The discovered vulnerability
            target_info: Information about the target
            objective: Desired outcome

        Returns:
            Custom-planned AttackChain or None on failure
        """
        try:
            # Prepare template summaries for the prompt
            template_summaries = []
            for name, template in CHAIN_TEMPLATES.items():
                template_summaries.append(
                    f"- {name}: {template.description} "
                    f"(Impact: {template.estimated_impact}, Stealth: {template.stealth_rating}/10)"
                )

            prompt = CHAIN_PLANNER_PROMPT.format(
                vuln_details=json.dumps(entry_vuln, indent=2),
                target_info=json.dumps(target_info, indent=2),
                objective=objective,
                templates="\n".join(template_summaries),
            )

            messages = [
                {"role": "user", "content": prompt},
            ]

            completion_kwargs: dict[str, Any] = {
                "model": self.model,
                "messages": messages,
                "temperature": 0.3,
                "timeout": 60,
            }

            if self.api_key:
                completion_kwargs["api_key"] = self.api_key
            if self.api_base:
                completion_kwargs["api_base"] = self.api_base

            response = await litellm.acompletion(**completion_kwargs)
            content = response.choices[0].message.content

            # Parse the JSON response
            # Find JSON block in response
            json_start = content.find("{")
            json_end = content.rfind("}") + 1
            if json_start == -1 or json_end == 0:
                logger.error("No JSON found in LLM response")
                return None

            chain_plan = json.loads(content[json_start:json_end])

            # Build chain from LLM plan
            steps = []
            for step_data in chain_plan.get("steps", []):
                step = ChainStep(
                    step_id=str(uuid.uuid4())[:8],
                    name=step_data.get("name", "unknown"),
                    step_type=StepType(step_data.get("step_type", "exploit")),
                    description=step_data.get("description", ""),
                    techniques=step_data.get("techniques", []),
                    required_conditions=step_data.get("required_conditions", []),
                    produces=step_data.get("produces", []),
                    fallback_techniques=step_data.get("fallback_techniques", []),
                )
                steps.append(step)

            chain = AttackChain(
                chain_id=str(uuid.uuid4())[:8],
                name=chain_plan.get("chain_name", "custom_chain"),
                description=chain_plan.get("description", ""),
                target=target_info.get("url", target_info.get("host", "")),
                entry_vuln=entry_vuln,
                steps=steps,
                estimated_impact=chain_plan.get("estimated_impact", "Unknown"),
                stealth_rating=chain_plan.get("stealth_rating", 5),
            )

            self._chain_cache[chain.chain_id] = chain
            logger.info(f"LLM planned custom chain '{chain.name}' with {len(steps)} steps")

            return chain

        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse LLM chain plan: {e}")
            return None
        except Exception as e:
            logger.exception(f"Error planning custom chain: {e}")
            return None

    def get_chain(self, chain_id: str) -> AttackChain | None:
        """Get a cached chain by ID."""
        return self._chain_cache.get(chain_id)

    def list_chains(self) -> list[AttackChain]:
        """List all cached chains."""
        return list(self._chain_cache.values())

    def _calculate_risk_score(self, template: ChainTemplate) -> float:
        """
        Calculate a risk score for a chain template.

        Lower score = safer to execute.
        """
        # Base risk from complexity
        risk = template.complexity * 0.1

        # Adjust for stealth (higher stealth = lower risk)
        risk -= template.stealth_rating * 0.05

        # Adjust for impact
        impact_scores = {
            "critical": 0.3,
            "high": 0.2,
            "medium": 0.1,
            "low": 0.0,
        }
        impact_str = template.estimated_impact.lower()
        for level, score in impact_scores.items():
            if level in impact_str:
                risk += score
                break

        return max(0.0, min(1.0, risk))

    def get_recommended_chains(
        self,
        vulnerabilities: list[dict[str, Any]],
        max_chains: int = 5,
    ) -> list[tuple[AttackChain, float]]:
        """
        Get recommended exploit chains based on discovered vulnerabilities.

        Args:
            vulnerabilities: List of discovered vulnerabilities
            max_chains: Maximum number of chains to recommend

        Returns:
            List of (chain, confidence_score) tuples
        """
        recommendations = []

        for vuln in vulnerabilities:
            vuln_type = vuln.get("type", vuln.get("vulnerability_type", ""))
            matching = get_chains_for_vuln_type(vuln_type)

            for template in matching:
                # Calculate confidence based on template match quality
                confidence = 1.0

                # Reduce confidence for high complexity
                confidence -= template.complexity * 0.05

                # Increase for high stealth
                confidence += template.stealth_rating * 0.02

                # Build the chain
                chain = self.build_from_template(
                    template_name=template.name,
                    target=vuln.get("url", vuln.get("target", "")),
                    entry_vuln=vuln,
                )

                if chain:
                    recommendations.append((chain, confidence))

        # Sort by confidence and deduplicate
        recommendations.sort(key=lambda x: x[1], reverse=True)

        # Remove duplicates by chain name
        seen_names = set()
        unique_recommendations = []
        for chain, confidence in recommendations:
            if chain.name not in seen_names:
                seen_names.add(chain.name)
                unique_recommendations.append((chain, confidence))
                if len(unique_recommendations) >= max_chains:
                    break

        return unique_recommendations


# Convenience functions
def build_chain(
    vuln_type: str,
    target: str,
    entry_vuln: dict[str, Any],
    objective: str = "rce",
) -> AttackChain | None:
    """Quick function to build a chain."""
    builder = ExploitChainBuilder()
    return builder.build_from_vulnerability(vuln_type, target, entry_vuln, objective)


def get_chains_for_vuln(vuln: dict[str, Any]) -> list[str]:
    """Get available chain names for a vulnerability."""
    vuln_type = vuln.get("type", vuln.get("vulnerability_type", ""))
    chains = get_chains_for_vuln_type(vuln_type)
    return [c.name for c in chains]


__all__ = [
    "ChainStatus",
    "StepStatus",
    "ChainStep",
    "AttackChain",
    "ExploitChainBuilder",
    "build_chain",
    "get_chains_for_vuln",
]
