"""
Exploit Engine Orchestrator for LogicPwn
- Executes exploit chains step-by-step
- Manages session state, retries, validation, and logging
- Extensible, maintainable, and interoperable
"""
import time
from typing import Any, Dict, Optional
from logicpwn.core.exploit_engine.models import ExploitChain, ExploitStep, ExploitResult, ExploitStatus
from logicpwn.core.exploit_engine.payload_generator import inject_payload
from logicpwn.core.exploit_engine.validation_engine import validate_step_success
from logicpwn.core.exploit_engine.exploit_logger import ExploitLogger
from logicpwn.models.request_config import RequestConfig
import requests
from logicpwn.core.performance import monitor_performance, PerformanceMonitor
from logicpwn.core.cache.session_cache import SessionCache
from logicpwn.core.cache.response_cache import ResponseCache
import os

try:
    import yaml
except ImportError:
    yaml = None
import json
import asyncio
import functools

session_cache = SessionCache()
response_cache = ResponseCache()

logger = ExploitLogger()

def extract_data_from_response(response: requests.Response, extractors: Dict[str, Any]) -> Dict[str, Any]:
    """
    Extracts data from a response using a dictionary of extractors.
    """
    extracted_data = {}
    for key, extractor in extractors.items():
        if isinstance(extractor, str):
            # Simple string extraction (e.g., "json.key")
            try:
                extracted_data[key] = response.json()[extractor]
            except (json.JSONDecodeError, KeyError) as e:
                logger.log_warning(f"Failed to extract data for key '{key}' from response: {e}")
        elif callable(extractor):
            # Callable extractor (e.g., a function that takes response and returns data)
            try:
                extracted_data[key] = extractor(response)
            except Exception as e:
                logger.log_warning(f"Failed to extract data for key '{key}' using callable extractor: {e}")
        else:
            logger.log_warning(f"Unsupported extractor type for key '{key}': {type(extractor)}")
    return extracted_data

@monitor_performance("exploit_chain_execution")
def run_exploit_chain(
    session: requests.Session,
    chain: ExploitChain,
    config: Optional[Dict[str, Any]] = None
) -> Any:
    """
    Execute a complete exploit chain with validation and logging.
    - Executes each step sequentially
    - Validates success/failure after each step
    - Updates session state between steps
    - Handles retries and error recovery
    - Generates comprehensive execution report
    """
    chain.session_state = chain.session_state or {}
    chain.execution_log = []
    for step in chain.steps:
        result = None
        for attempt in range(step.retry_count):
            start = time.time()
            try:
                # Inject payloads if any
                req_cfg = step.request_config
                if step.payload_injection_points:
                    req_cfg = inject_payload(req_cfg, step.payload_injection_points, chain.session_state)
                # Prepare request
                req_kwargs = req_cfg.model_dump(exclude_none=True)
                method = req_kwargs.pop("method")
                url = req_kwargs.pop("url")
                # Check response cache first
                cached_response = response_cache.get_response(url, method, req_kwargs.get('params'), req_kwargs.get('headers'))
                if cached_response:
                    response = cached_response
                else:
                    response = session.request(method, url, **req_kwargs)
                    response_cache.set_response(url, method, response, req_kwargs.get('params'), req_kwargs.get('headers'))
                # Validate
                validation = validate_step_success(response, step, chain.session_state)
                # Extract data
                extracted_data = {}
                if hasattr(step, 'data_extractors') and step.data_extractors:
                    extracted_data = extract_data_from_response(response, step.data_extractors)
                    chain.session_state.update(extracted_data)
                # Log
                exec_time = time.time() - start
                result = ExploitResult(
                    step_name=step.name,
                    status=ExploitStatus.SUCCESS if validation.is_valid else ExploitStatus.FAILED,
                    response=response,
                    validation_result=validation,
                    execution_time=exec_time,
                    extracted_data=extracted_data
                )
                logger.log_step(step.name, result.status.value, {
                    "execution_time": exec_time,
                    "matched": validation.matched_indicators,
                    "failed": validation.failed_indicators
                })
                logger.log_request_response(step.name, req_kwargs, response)
                if validation.is_valid:
                    chain.session_state.update(validation.extracted_data)
                    break
                else:
                    logger.log_error(step.name, f"Validation failed: {validation.reasons}")
            except Exception as e:
                exec_time = time.time() - start
                result = ExploitResult(
                    step_name=step.name,
                    status=ExploitStatus.FAILED,
                    error_message=str(e),
                    execution_time=exec_time
                )
                logger.log_error(step.name, str(e))
            if step.critical and (not result or result.status != ExploitStatus.SUCCESS):
                logger.log_error(step.name, "Critical step failed, aborting chain.")
                break
            if step.delay_after_step:
                time.sleep(step.delay_after_step)
        chain.execution_log.append(result)
        if step.critical and result and result.status != ExploitStatus.SUCCESS:
            break
    logger.log_chain(chain.name, [r.model_dump() for r in chain.execution_log])
    return chain.execution_log

async def async_run_exploit_chain(
    session: requests.Session,
    chain: ExploitChain,
    config: Optional[Dict[str, Any]] = None
) -> Any:
    """
    Async version of run_exploit_chain. Supports parallel execution of steps if chain.parallel_execution is True.
    """
    logger = ExploitLogger()
    chain.session_state = chain.session_state or {}
    chain.execution_log = []
    async def run_step(step: ExploitStep):
        result = None
        for attempt in range(step.retry_count):
            start = time.time()
            try:
                req_cfg = step.request_config
                if step.payload_injection_points:
                    req_cfg = inject_payload(req_cfg, step.payload_injection_points, chain.session_state)
                req_kwargs = req_cfg.model_dump(exclude_none=True)
                method = req_kwargs.pop("method")
                url = req_kwargs.pop("url")
                loop = asyncio.get_event_loop()
                req_func = functools.partial(session.request, method, url, **req_kwargs)
                response = await loop.run_in_executor(None, req_func)
                validation = validate_step_success(response, step, chain.session_state)
                extracted_data = {}
                if hasattr(step, 'data_extractors') and step.data_extractors:
                    extracted_data = extract_data_from_response(response, step.data_extractors)
                    chain.session_state.update(extracted_data)
                exec_time = time.time() - start
                result = ExploitResult(
                    step_name=step.name,
                    status=ExploitStatus.SUCCESS if validation.is_valid else ExploitStatus.FAILED,
                    response=response,
                    validation_result=validation,
                    execution_time=exec_time,
                    extracted_data=extracted_data
                )
                logger.log_step(step.name, result.status.value, {
                    "execution_time": exec_time,
                    "matched": validation.matched_indicators,
                    "failed": validation.failed_indicators
                })
                logger.log_request_response(step.name, req_kwargs, response)
                if validation.is_valid:
                    chain.session_state.update(validation.extracted_data)
                    break
                else:
                    logger.log_error(step.name, f"Validation failed: {validation.reasons}")
            except Exception as e:
                exec_time = time.time() - start
                result = ExploitResult(
                    step_name=step.name,
                    status=ExploitStatus.FAILED,
                    error_message=str(e),
                    execution_time=exec_time
                )
                logger.log_error(step.name, str(e))
            if step.critical and (not result or result.status != ExploitStatus.SUCCESS):
                logger.log_error(step.name, "Critical step failed, aborting chain.")
                break
            if step.delay_after_step:
                await asyncio.sleep(step.delay_after_step)
        return result
    if getattr(chain, 'parallel_execution', False):
        results = await asyncio.gather(*(run_step(step) for step in chain.steps))
        chain.execution_log.extend(results)
    else:
        for step in chain.steps:
            result = await run_step(step)
            chain.execution_log.append(result)
            if step.critical and result and result.status != ExploitStatus.SUCCESS:
                break
    logger.log_chain(chain.name, [r.model_dump() for r in chain.execution_log])
    return chain.execution_log

def load_exploit_chain_from_file(file_path: str) -> ExploitChain:
    """
    Load an exploit chain from a YAML or JSON file.
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"Exploit chain file not found: {file_path}")
    with open(file_path, 'r') as f:
        if file_path.endswith('.yaml') or file_path.endswith('.yml'):
            if not yaml:
                raise ImportError("PyYAML is required to load YAML files.")
            data = yaml.safe_load(f)
        elif file_path.endswith('.json'):
            data = json.load(f)
        else:
            raise ValueError("Unsupported file format. Use .yaml, .yml, or .json")
    return ExploitChain.parse_obj(data) 