"""
Payload injection system for LogicPwn Exploit Engine.
- Supports static, dynamic, and template payloads
- Injects into URL, headers, body, JSON path, and form data
- Extensible and interoperable with core models
"""
from typing import List, Dict, Any
from logicpwn.models.request_config import RequestConfig
from .models import PayloadInjectionPoint, PayloadType
import copy
import json
import base64
import urllib.parse
import random
import string
import uuid

try:
    from jinja2 import Template
except ImportError:
    Template = None

class PayloadGenerator:
    @staticmethod
    def random_string(length=8):
        return ''.join(random.choices(string.ascii_letters + string.digits, k=length))

    @staticmethod
    def random_int(min_val=0, max_val=10000):
        return random.randint(min_val, max_val)

    @staticmethod
    def random_uuid():
        return str(uuid.uuid4())

    @staticmethod
    def render_template(template_str: str, context: dict) -> str:
        if Template:
            return Template(template_str).render(**context)
        else:
            return template_str.format(**context)


def inject_payload(
    request_config: RequestConfig,
    injection_points: List[PayloadInjectionPoint],
    session_state: Dict[str, Any]
) -> RequestConfig:
    """
    Inject payloads into request at specified locations.
    Supports injection into:
    - URL parameters and path
    - HTTP headers
    - JSON body fields (using dot notation)
    - Form data
    - Raw body content
    
    Enhanced with proper error handling and validation.
    """
    config = copy.deepcopy(request_config)
    
    for point in injection_points:
        try:
            value = point.payload_value
            
            # Handle None payload values
            if value is None:
                if point.payload_type == PayloadType.STATIC:
                    raise ValueError(f"Static payload cannot have None value for {point.target_field}")
                # Generate default value based on payload type
                value = ""
            
            # Advanced payload types
            if point.payload_type.value == "random":
                if point.target_field.lower().endswith("id"):
                    value = str(PayloadGenerator.random_int())
                elif point.target_field.lower().endswith("uuid"):
                    value = PayloadGenerator.random_uuid()
                else:
                    value = PayloadGenerator.random_string()
            elif point.payload_type.value == "fuzz":
                # Simple fuzz: random string with special chars
                fuzz_chars = string.ascii_letters + string.digits + "'\"<>!@#$%^&*()"
                value = ''.join(random.choices(fuzz_chars, k=12))
            elif point.payload_type.value == "template":
                if value is None:
                    raise ValueError(f"Template payload requires payload_value for {point.target_field}")
                value = PayloadGenerator.render_template(value, session_state)
            elif point.payload_type.value == "dynamic":
                # Dynamic payload based on session state
                if point.target_field in session_state:
                    value = str(session_state[point.target_field])
                else:
                    # Generate a reasonable default
                    value = PayloadGenerator.random_string()
            
            # Convert to string if needed
            if not isinstance(value, str):
                value = str(value)
            
            # Encoding
            if point.encoding == "url":
                value = urllib.parse.quote_plus(value)
            elif point.encoding == "base64":
                value = base64.b64encode(value.encode()).decode()
            elif point.encoding == "html":
                import html
                value = html.escape(value)
            
            # Injection with enhanced error handling
            if point.location == "url":
                # Safe URL injection with placeholder validation
                placeholder = f"{{{point.target_field}}}"
                if placeholder not in config.url:
                    raise ValueError(f"URL placeholder {placeholder} not found in {config.url}")
                config.url = config.url.replace(placeholder, value)
                
            elif point.location == "headers":
                config.headers = config.headers or {}
                config.headers[point.target_field] = value
                
            elif point.location == "body":
                config.raw_body = value
                
            elif point.location == "json_path":
                if not config.json_data:
                    config.json_data = {}
                try:
                    _set_json_path(config.json_data, point.target_field, value)
                except Exception as e:
                    raise ValueError(f"Failed to set JSON path {point.target_field}: {e}")
                    
            elif point.location == "form_data":
                config.data = config.data or {}
                config.data[point.target_field] = value
            else:
                raise ValueError(f"Unsupported injection location: {point.location}")
                
        except Exception as e:
            # Enhanced error reporting
            raise ValueError(f"Payload injection failed for {point.target_field} at {point.location}: {e}")
    
    return config

def _set_json_path(obj: Dict[str, Any], path: str, value: Any):
    """
    Set a value in a nested dict using dot notation path (e.g., 'user.id').
    """
    keys = path.split('.')
    for k in keys[:-1]:
        if k not in obj or not isinstance(obj[k], dict):
            obj[k] = {}
        obj = obj[k]
    obj[keys[-1]] = value 