"""Execution strategies for Tru-Guard."""

import asyncio
from abc import ABC, abstractmethod
from typing import Any

from trusys.exceptions import ScannerExecutionError, ScannerTimeoutError
from trusys.scanners.base import BaseScanner
from trusys.types import ExecutionMode, ScanContext, ScanResult, ScanType, Violation
from trusys.utils.logging import Timer, get_logger


def _sync_latency_metadata(result: ScanResult) -> None:
    """Sync latency field in metadata with execution_time_ms.
    
    Sets latency in milliseconds (same as execution_time_ms).
    
    Args:
        result: ScanResult to update
    """
    # Always update latency if metadata has new format (has "action" field)
    # This ensures latency is synced even if scanner set it to 0.0 initially
    if "action" in result.metadata:
        # Force update latency from execution_time_ms
        result.metadata["latency"] = result.execution_time_ms


class ExecutionStrategy(ABC):
    """Abstract base class for execution strategies."""

    @abstractmethod
    async def execute(
        self,
        scanners: list[BaseScanner],
        content: dict[str, Any],
        context: ScanContext,
        timeout: float | None = None,
    ) -> list[ScanResult]:
        """Execute scanners using this strategy.

        Args:
            scanners: List of scanners to execute
            content: Content dict to scan (contains 'prompt' and/or 'response')
            context: Scan context
            timeout: Optional per-scanner timeout in seconds

        Returns:
            List of scan results
        """
        pass


class ParallelStrategy(ExecutionStrategy):
    """Execute all scanners concurrently."""

    def __init__(self, max_concurrent: int | None = None) -> None:
        """Initialize the parallel strategy.

        Args:
            max_concurrent: Maximum concurrent scanners (None = unlimited)
        """
        self.max_concurrent = max_concurrent

    async def execute(
        self,
        scanners: list[BaseScanner],
        content: dict[str, Any],
        context: ScanContext,
        timeout: float | None = None,
    ) -> list[ScanResult]:
        """Execute all scanners in parallel.

        Args:
            scanners: List of scanners to execute
            content: Content dict to scan (contains 'prompt' and/or 'response')
            context: Scan context
            timeout: Optional per-scanner timeout in seconds

        Returns:
            List of scan results
        """
        if not scanners:
            return []

        logger = get_logger()

        async def run_scanner(scanner: BaseScanner) -> ScanResult:
            """Run a single scanner with timeout handling."""
            logger.scanner_started(scanner.name, context.scan_type.value)

            with Timer() as timer:
                try:
                    if timeout:
                        result = await asyncio.wait_for(
                            scanner.scan_async(content, context),
                            timeout=timeout,
                        )
                    else:
                        result = await scanner.scan_async(content, context)

                    result.execution_time_ms = timer.elapsed_ms
                    _sync_latency_metadata(result)
                    logger.scanner_completed(
                        scanner.name,
                        result.passed,
                        timer.elapsed_ms,
                        len(result.violations),
                    )
                    return result

                except asyncio.TimeoutError:
                    logger.scanner_timeout(scanner.name, timeout or 0)
                    raise ScannerTimeoutError(scanner.name, timeout or 0)

                except Exception as e:
                    logger.scanner_error(scanner.name, e)
                    if isinstance(e, (ScannerTimeoutError, ScannerExecutionError)):
                        raise
                    raise ScannerExecutionError(scanner.name, e) from e

        # Use semaphore for concurrency limiting if specified
        if self.max_concurrent:
            semaphore = asyncio.Semaphore(self.max_concurrent)

            async def limited_run(scanner: BaseScanner) -> ScanResult:
                async with semaphore:
                    return await run_scanner(scanner)

            tasks = [limited_run(s) for s in scanners]
        else:
            tasks = [run_scanner(s) for s in scanners]

        # Gather results, capturing exceptions
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Process results
        processed: list[ScanResult] = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                # Create a failed result for the exception
                scanner = scanners[i]
                error_result = ScanResult(
                    passed=False,
                    scanner_name=scanner.name,
                    scan_type=context.scan_type,
                    violations=[
                        Violation(
                            rule="scanner_error",
                            severity="critical",
                            message=str(result),
                        )
                    ],
                    execution_time_ms=0.0,
                    metadata={
                        "error": str(result),
                        "error_type": type(result).__name__,
                        "assertion": {"type": scanner.name, "output": ""},
                        "score": 0.0,
                        "reason": f"Scanner error: {str(result)}",
                        "latency": 0.0,
                        "action": {"type": "ignore"},
                    },
                )
                processed.append(error_result)
            else:
                # Ensure latency is synced even if it wasn't updated
                _sync_latency_metadata(result)
                processed.append(result)

        return processed


class SequentialStrategy(ExecutionStrategy):
    """Execute scanners one at a time."""

    async def execute(
        self,
        scanners: list[BaseScanner],
        content: dict[str, Any],
        context: ScanContext,
        timeout: float | None = None,
    ) -> list[ScanResult]:
        """Execute scanners sequentially.

        Args:
            scanners: List of scanners to execute
            content: Content dict to scan (contains 'prompt' and/or 'response')
            context: Scan context
            timeout: Optional per-scanner timeout in seconds

        Returns:
            List of scan results
        """
        logger = get_logger()
        results: list[ScanResult] = []

        for scanner in scanners:
            logger.scanner_started(scanner.name, context.scan_type.value)

            with Timer() as timer:
                try:
                    if timeout:
                        result = await asyncio.wait_for(
                            scanner.scan_async(content, context),
                            timeout=timeout,
                        )
                    else:
                        result = await scanner.scan_async(content, context)

                    result.execution_time_ms = timer.elapsed_ms
                    _sync_latency_metadata(result)
                    logger.scanner_completed(
                        scanner.name,
                        result.passed,
                        timer.elapsed_ms,
                        len(result.violations),
                    )
                    results.append(result)

                except asyncio.TimeoutError:
                    logger.scanner_timeout(scanner.name, timeout or 0)
                    timeout_msg = f"Scanner timed out after {timeout}s"
                    timeout_result = ScanResult(
                        passed=False,
                        scanner_name=scanner.name,
                        scan_type=context.scan_type,
                        violations=[
                            Violation(
                                rule="timeout",
                                severity="critical",
                                message=timeout_msg,
                            )
                        ],
                        execution_time_ms=timer.elapsed_ms,
                        metadata={
                            "error": timeout_msg,
                            "assertion": {"type": scanner.name, "output": ""},
                            "score": 0.0,
                            "reason": timeout_msg,
                            "latency": timer.elapsed_ms,
                            "action": {"type": "ignore"},
                        },
                    )
                    results.append(timeout_result)

                except Exception as e:
                    logger.scanner_error(scanner.name, e)
                    error_result = ScanResult(
                        passed=False,
                        scanner_name=scanner.name,
                        scan_type=context.scan_type,
                        violations=[
                            Violation(
                                rule="scanner_error",
                                severity="critical",
                                message=str(e),
                            )
                        ],
                        execution_time_ms=timer.elapsed_ms,
                        metadata={
                            "error": str(e),
                            "error_type": type(e).__name__,
                            "assertion": {"type": scanner.name, "output": ""},
                            "score": 0.0,
                            "reason": f"Scanner error: {str(e)}",
                            "latency": timer.elapsed_ms,
                            "action": {"type": "ignore"},
                        },
                    )
                    results.append(error_result)

        return results


class PipelineStrategy(ExecutionStrategy):
    """Execute scanners sequentially, passing modified content to next."""

    def __init__(self, early_exit: bool = False) -> None:
        """Initialize the pipeline strategy.

        Args:
            early_exit: Stop on first failure if True
        """
        self.early_exit = early_exit

    async def execute(
        self,
        scanners: list[BaseScanner],
        content: dict[str, Any],
        context: ScanContext,
        timeout: float | None = None,
    ) -> list[ScanResult]:
        """Execute scanners in pipeline mode.

        Args:
            scanners: List of scanners to execute
            content: Content dict to scan (contains 'prompt' and/or 'response')
            context: Scan context
            timeout: Optional per-scanner timeout in seconds

        Returns:
            List of scan results
        """
        logger = get_logger()
        results: list[ScanResult] = []
        current_content = content.copy()
        
        # Determine which field to modify based on scan type
        content_field = "prompt" if context.scan_type == ScanType.INPUT else "response"

        for scanner in scanners:
            logger.scanner_started(scanner.name, context.scan_type.value)

            with Timer() as timer:
                try:
                    if timeout:
                        result = await asyncio.wait_for(
                            scanner.scan_async(current_content, context),
                            timeout=timeout,
                        )
                    else:
                        result = await scanner.scan_async(current_content, context)

                    result.execution_time_ms = timer.elapsed_ms
                    _sync_latency_metadata(result)
                    logger.scanner_completed(
                        scanner.name,
                        result.passed,
                        timer.elapsed_ms,
                        len(result.violations),
                    )
                    results.append(result)

                    # Early exit on failure
                    if self.early_exit and not result.passed:
                        break

                    # Pass suggested fix to next scanner
                    if result.suggested_fix:
                        current_content = current_content.copy()
                        current_content[content_field] = result.suggested_fix

                except asyncio.TimeoutError:
                    logger.scanner_timeout(scanner.name, timeout or 0)
                    timeout_msg = f"Scanner timed out after {timeout}s"
                    timeout_result = ScanResult(
                        passed=False,
                        scanner_name=scanner.name,
                        scan_type=context.scan_type,
                        violations=[
                            Violation(
                                rule="timeout",
                                severity="critical",
                                message=timeout_msg,
                            )
                        ],
                        execution_time_ms=timer.elapsed_ms,
                        metadata={
                            "error": timeout_msg,
                            "assertion": {"type": scanner.name, "output": ""},
                            "score": 0.0,
                            "reason": timeout_msg,
                            "latency": timer.elapsed_ms,
                            "action": {"type": "ignore"},
                        },
                    )
                    results.append(timeout_result)
                    if self.early_exit:
                        break

                except Exception as e:
                    logger.scanner_error(scanner.name, e)
                    error_result = ScanResult(
                        passed=False,
                        scanner_name=scanner.name,
                        scan_type=context.scan_type,
                        violations=[
                            Violation(
                                rule="scanner_error",
                                severity="critical",
                                message=str(e),
                            )
                        ],
                        execution_time_ms=timer.elapsed_ms,
                        metadata={
                            "error": str(e),
                            "error_type": type(e).__name__,
                            "assertion": {"type": scanner.name, "output": ""},
                            "score": 0.0,
                            "reason": f"Scanner error: {str(e)}",
                            "latency": timer.elapsed_ms,
                            "action": {"type": "ignore"},
                        },
                    )
                    results.append(error_result)
                    if self.early_exit:
                        break

        return results


def get_strategy_for_mode(
    mode: ExecutionMode,
    early_exit: bool = False,
    max_concurrent: int | None = None,
) -> ExecutionStrategy:
    """Get an execution strategy for the given mode.

    Args:
        mode: Execution mode
        early_exit: For pipeline mode, whether to exit on first failure
        max_concurrent: For parallel mode, max concurrent scanners

    Returns:
        Appropriate execution strategy
    """
    if mode == ExecutionMode.PARALLEL:
        return ParallelStrategy(max_concurrent=max_concurrent)
    elif mode == ExecutionMode.SEQUENTIAL:
        return SequentialStrategy()
    elif mode == ExecutionMode.PIPELINE:
        return PipelineStrategy(early_exit=early_exit)
    else:
        raise ValueError(f"Unknown execution mode: {mode}")
