"""Execution runner for Tru-Guard."""

import asyncio
import time
import threading
from concurrent.futures import ThreadPoolExecutor
from typing import Any

from trusys.actions import ActionResolver, get_action_for_type
from trusys.core.config import GuardConfig
from trusys.exceptions import ScanFailedError
from trusys.execution.strategies import ExecutionStrategy, get_strategy_for_mode
from trusys.scanners.base import BaseScanner
from trusys.scanners.remote import BaseRemoteScanner
from trusys.types import (
    ActionRecord,
    ActionType,
    AggregatedResult,
    ScanContext,
    ScanResult,
    ScanType,
)
from trusys.utils.logging import Timer, get_logger


class ExecutionRunner:
    """Orchestrates scanner execution and action handling."""
    _threadpool: ThreadPoolExecutor | None = None
    _threadpool_lock = threading.Lock()

    def __init__(self, config: GuardConfig) -> None:
        """Initialize the execution runner.

        Args:
            config: Guard configuration
        """
        self.config = config
        self.action_resolver = ActionResolver(strategy=config.conflict_strategy)
        self._strategy: ExecutionStrategy | None = None

    @property
    def strategy(self) -> ExecutionStrategy:
        """Get the execution strategy, creating if needed."""
        if self._strategy is None:
            self._strategy = get_strategy_for_mode(
                self.config.execution_mode,
                early_exit=self.config.early_exit,
                max_concurrent=self.config.max_parallel_scanners,
            )
        return self._strategy

    def update_config(self, config: GuardConfig) -> None:
        """Update the runner configuration.

        Args:
            config: New configuration
        """
        self.config = config
        self.action_resolver = ActionResolver(strategy=config.conflict_strategy)
        self._strategy = None  # Reset strategy to be recreated

    async def run_async(
        self,
        scanners: list[BaseScanner],
        context: ScanContext,
    ) -> AggregatedResult:
        """Run all scanners asynchronously.

        Args:
            scanners: List of scanners to run
            context: Scan context

        Returns:
            Aggregated result from all scanners
        """
        logger = get_logger()
        # Create content dict from context (API expects string, not None)
        content = {
            "prompt": context.prompt if context.prompt is not None else "",
            "response": context.response if context.response is not None else "",
        }

        # Filter to enabled scanners that support the scan type
        active_scanners = [
            s for s in scanners
            if s.enabled and s.supports_scan_type(context.scan_type)
        ]

        if not active_scanners:
            return AggregatedResult(passed=True)

        logger.scan_started(
            context.scan_type.value,
            len(active_scanners),
            self.config.execution_mode.value,
        )

        with Timer() as timer:
            remote_scanners: list[BaseRemoteScanner] = []
            local_scanners: list[BaseScanner] = []

            for scanner in active_scanners:
                if isinstance(scanner, BaseRemoteScanner) and getattr(scanner, "is_remote", False):
                    remote_scanners.append(scanner)
                else:
                    local_scanners.append(scanner)

            remote_task = None
            if remote_scanners:
                import os
                from trusys.core.guard import TruGuard

                api_key = (
                    TruGuard._api_key
                    or getattr(self.config, "api_key", None)
                    or os.getenv("TRUSYS_API_KEY")
                    or None
                )
                application_id = (
                    TruGuard._application_id
                    or getattr(self.config, "application_id", None)
                    or os.getenv("TRUSYS_APPLICATION_ID")
                    or None
                )
                logger.scan_started(
                    f"{context.scan_type.value}_remote",
                    len(remote_scanners),
                    "batch",
                )
                remote_task = self._execute_remote_batch_async(
                    remote_scanners,
                    content,
                    context,
                    api_key,
                    application_id,
                )

            local_task = None
            if local_scanners:
                local_task = self.strategy.execute(
                    local_scanners,
                    content,
                    context,
                    timeout=self.config.timeout,
                )

            if remote_task and local_task:
                remote_results, local_results = await self._gather_results(
                    remote_task,
                    local_task,
                )
            elif remote_task:
                remote_results = await remote_task
                local_results = []
            elif local_task:
                local_results = await local_task
                remote_results = []
            else:
                remote_results = []
                local_results = []

            results = remote_results + local_results

            # Ensure all results have latency synced (final pass)
            from trusys.execution.strategies import _sync_latency_metadata
            for result in results:
                _sync_latency_metadata(result)

            # Process results and apply actions
            try:
                final_content, actions_taken = await self._process_results(
                    results,
                    active_scanners,
                    content,
                    context,
                )
            except ScanFailedError as e:
                e.aggregated_result = AggregatedResult(
                    passed=False,
                    results=results,
                    final_content=None,
                    actions_taken=[],
                    total_execution_time_ms=timer.elapsed_ms,
                )
                raise

        # Determine overall pass/fail
        passed = all(r.passed for r in results)

        total_violations = sum(len(r.violations) for r in results)
        logger.scan_completed(
            passed,
            timer.elapsed_ms,
            len(results),
            total_violations,
        )

        return AggregatedResult(
            passed=passed,
            results=results,
            final_content=final_content,
            actions_taken=actions_taken,
            total_execution_time_ms=timer.elapsed_ms,
        )

    async def _process_results(
        self,
        results: list[ScanResult],
        scanners: list[BaseScanner],
        original_content: dict[str, Any],
        context: ScanContext,
    ) -> tuple[str | None, list[ActionRecord]]:
        """Process scan results and apply actions.

        Args:
            results: Scan results
            scanners: Scanners that were run
            original_content: Original content dict before scanning
            context: Scan context to determine which field to modify

        Returns:
            Tuple of (final_content_string, actions_taken). final_content_string is
            the modified prompt for input scans, or the modified response for output
            scans; or the original content when no modifications were applied.
        """
        logger = get_logger()
        actions_taken: list[ActionRecord] = []
        current_content = original_content.copy()
        
        # Determine which field to modify based on scan type
        content_field = "prompt" if context.scan_type == ScanType.INPUT else "response"
        current_text = current_content.get(content_field, "")

        # Build scanner name to scanner mapping
        scanner_map = {s.name: s for s in scanners}

        # Collect failed results with their actions
        failed_results: list[tuple[str, ActionType, ScanResult]] = []
        for result in results:
            if not result.passed:
                scanner = scanner_map.get(result.scanner_name)
                if scanner:
                    failed_results.append((result.scanner_name, scanner.on_fail, result))

        if not failed_results:
            return None, actions_taken

        # Log error for each failed validator
        for name, action, result in failed_results:
            if result.violations:
                parts = [v.message for v in result.violations[:3]]
                summary = "; ".join(parts)
                if len(result.violations) > 3:
                    summary += f" (+{len(result.violations) - 3} more)"
            else:
                summary = "check failed"
            logger.error(f"Validator failed: {name} - {summary}")

        # Resolve action conflicts (returns list of actions to apply, possibly multiple for PRIORITY without FAIL)
        if len(failed_results) > 1:
            action_info = [(name, action, None) for name, action, _ in failed_results]
            resolved_list = self.action_resolver.resolve(action_info)
            if len(resolved_list) == 1:
                logger.action_conflict(
                    [f"{name}:{action.value}" for name, action, _ in failed_results],
                    resolved_list[0][1].value,
                )
            # else: multiple actions to apply, no conflict log for single winner
        else:
            resolved_list = [
                (failed_results[0][0], failed_results[0][1], None),
            ]

        if not resolved_list:
            return None, actions_taken

        original_content_snapshot = original_content.copy()
        original_text_for_span = original_content_snapshot.get(content_field, "")

        # Enrich violations with matched_text from original content so actions can replace
        # by content instead of character positions (avoids wrong positions after prior actions)
        for scan_result in results:
            for violation in scan_result.violations:
                if violation.span and not violation.metadata.get("matched_text"):
                    start, end = violation.span
                    violation.metadata["matched_text"] = original_text_for_span[start:end]

        # Execute each resolved action in order (e.g. all FIX then all FILTER when no FAIL)
        # Track whether current_text has diverged from the original so we can
        # invalidate stale span positions on subsequent actions.
        text_modified_by_prior_action = False

        for resolved_name, resolved_action, _ in resolved_list:
            resolved_result = next(
                (r for name, _, r in failed_results if name == resolved_name),
                failed_results[0][2],
            )

            # When a prior action changed the text length, span positions from the
            # original scan are stale (they point to the wrong characters in the
            # now-modified text).  Clear them so the action layer falls back to
            # matched_text-based replacement, which is content-addressed and still
            # valid regardless of prior edits.
            if text_modified_by_prior_action:
                for violation in resolved_result.violations:
                    if violation.span is not None and violation.metadata.get("matched_text"):
                        violation.span = None

            try:
                action = get_action_for_type(resolved_action)
                modified_text, should_continue = action.execute(
                    resolved_result,
                    current_text,
                )

                logger.action_taken(
                    resolved_action.value,
                    resolved_name,
                    content_modified=modified_text is not None,
                )

                original_text = current_text
                if modified_text is not None:
                    current_content[content_field] = modified_text
                    current_text = modified_text
                    text_modified_by_prior_action = True

                actions_taken.append(
                    ActionRecord(
                        action_type=resolved_action,
                        scanner_name=resolved_name,
                        original_content=original_text,
                        modified_content=modified_text if modified_text is not None else original_text,
                        timestamp=time.time(),
                    )
                )

            except ScanFailedError:
                # Re-raise fail action exceptions
                raise

        # Return the modified field's string, or original content when no modifications applied
        if current_content != original_content_snapshot:
            return current_content.get(content_field, ""), actions_taken
        return original_content_snapshot.get(content_field, ""), actions_taken

    def run(
        self,
        scanners: list[BaseScanner],
        context: ScanContext,
    ) -> AggregatedResult:
        """Run all scanners synchronously.

        Optimized to avoid asyncio overhead when only remote scanners are used.

        Args:
            scanners: List of scanners to run
            context: Scan context

        Returns:
            Aggregated result from all scanners
        """
        import os

        from trusys.core.guard import TruGuard
        from trusys.execution.strategies import _sync_latency_metadata

        logger = get_logger()

        # Create content dict from context
        content = {
            "prompt": context.prompt if context.prompt is not None else "",
            "response": context.response if context.response is not None else "",
        }

        # Filter to enabled scanners that support the scan type
        active_scanners = [
            s for s in scanners
            if s.enabled and s.supports_scan_type(context.scan_type)
        ]

        if not active_scanners:
            return AggregatedResult(passed=True)

        # Separate remote scanners from local scanners
        remote_scanners: list[BaseRemoteScanner] = []
        local_scanners: list[BaseScanner] = []

        for scanner in active_scanners:
            if isinstance(scanner, BaseRemoteScanner) and getattr(scanner, "is_remote", False):
                remote_scanners.append(scanner)
            else:
                local_scanners.append(scanner)

        # If we have local scanners, fall back to async path (strategies are async)
        if local_scanners:
            import asyncio
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = None

            if loop and loop.is_running():
                pool = self._get_threadpool()
                future = pool.submit(asyncio.run, self.run_async(scanners, context))
                return future.result()
            else:
                return asyncio.run(self.run_async(scanners, context))

        # Fast path: only remote scanners - no asyncio needed
        logger.scan_started(
            context.scan_type.value,
            len(active_scanners),
            self.config.execution_mode.value,
        )

        with Timer() as timer:
            # Get API credentials
            api_key = (
                TruGuard._api_key
                or getattr(self.config, "api_key", None)
                or os.getenv("TRUSYS_API_KEY")
                or None
            )
            application_id = (
                TruGuard._application_id
                or getattr(self.config, "application_id", None)
                or os.getenv("TRUSYS_APPLICATION_ID")
                or None
            )

            # Execute remote scanners directly (sync, with connection pooling)
            results = self._execute_remote_batch(
                remote_scanners,
                content,
                context,
                api_key,
                application_id,
            )

            # Sync latency metadata
            for result in results:
                _sync_latency_metadata(result)

            # Process results and apply actions (sync version)
            try:
                final_content, actions_taken = self._process_results_sync(
                    results,
                    active_scanners,
                    content,
                    context,
                )
            except ScanFailedError as e:
                e.aggregated_result = AggregatedResult(
                    passed=False,
                    results=results,
                    final_content=None,
                    actions_taken=[],
                    total_execution_time_ms=timer.elapsed_ms,
                )
                raise

        # Determine overall pass/fail
        passed = all(r.passed for r in results)

        total_violations = sum(len(r.violations) for r in results)
        logger.scan_completed(
            passed,
            timer.elapsed_ms,
            len(results),
            total_violations,
        )

        return AggregatedResult(
            passed=passed,
            results=results,
            final_content=final_content,
            actions_taken=actions_taken,
            total_execution_time_ms=timer.elapsed_ms,
        )

    def _process_results_sync(
        self,
        results: list[ScanResult],
        scanners: list[BaseScanner],
        original_content: dict[str, Any],
        context: ScanContext,
    ) -> tuple[str | None, list[ActionRecord]]:
        """Sync version of _process_results for the fast path."""
        import time

        from trusys.actions import get_action_for_type

        logger = get_logger()
        actions_taken: list[ActionRecord] = []
        current_content = original_content.copy()

        # Determine which field to modify based on scan type
        content_field = "prompt" if context.scan_type == ScanType.INPUT else "response"
        current_text = current_content.get(content_field, "")

        # Build scanner name to scanner mapping
        scanner_map = {s.name: s for s in scanners}

        # Collect failed results with their actions
        failed_results: list[tuple[str, ActionType, ScanResult]] = []
        for result in results:
            if not result.passed:
                scanner = scanner_map.get(result.scanner_name)
                if scanner:
                    failed_results.append((result.scanner_name, scanner.on_fail, result))

        if not failed_results:
            return None, actions_taken

        # Log error for each failed validator
        for name, action, result in failed_results:
            if result.violations:
                parts = [v.message for v in result.violations[:3]]
                summary = "; ".join(parts)
                if len(result.violations) > 3:
                    summary += f" (+{len(result.violations) - 3} more)"
            else:
                summary = "check failed"
            logger.error(f"Validator failed: {name} - {summary}")

        # Resolve action conflicts
        if len(failed_results) > 1:
            action_info = [(name, action, None) for name, action, _ in failed_results]
            resolved_list = self.action_resolver.resolve(action_info)
            if len(resolved_list) == 1:
                logger.action_conflict(
                    [f"{name}:{action.value}" for name, action, _ in failed_results],
                    resolved_list[0][1].value,
                )
        else:
            resolved_list = [
                (failed_results[0][0], failed_results[0][1], None),
            ]

        if not resolved_list:
            return None, actions_taken

        original_content_snapshot = original_content.copy()
        original_text_for_span = original_content_snapshot.get(content_field, "")

        # Enrich violations with matched_text
        for scan_result in results:
            for violation in scan_result.violations:
                if violation.span and not violation.metadata.get("matched_text"):
                    start, end = violation.span
                    violation.metadata["matched_text"] = original_text_for_span[start:end]

        # Execute each resolved action
        # Track whether current_text has diverged from the original so we can
        # invalidate stale span positions on subsequent actions.
        text_modified_by_prior_action = False

        for resolved_name, resolved_action, _ in resolved_list:
            resolved_result = next(
                (r for name, _, r in failed_results if name == resolved_name),
                failed_results[0][2],
            )

            # When a prior action changed the text length, span positions from the
            # original scan are stale.  Clear them so the action layer falls back to
            # matched_text-based replacement.
            if text_modified_by_prior_action:
                for violation in resolved_result.violations:
                    if violation.span is not None and violation.metadata.get("matched_text"):
                        violation.span = None

            try:
                action = get_action_for_type(resolved_action)
                modified_text, should_continue = action.execute(
                    resolved_result,
                    current_text,
                )

                logger.action_taken(
                    resolved_action.value,
                    resolved_name,
                    content_modified=modified_text is not None,
                )

                original_text = current_text
                if modified_text is not None:
                    current_content[content_field] = modified_text
                    current_text = modified_text
                    text_modified_by_prior_action = True

                actions_taken.append(
                    ActionRecord(
                        action_type=resolved_action,
                        scanner_name=resolved_name,
                        original_content=original_text,
                        modified_content=modified_text if modified_text is not None else original_text,
                        timestamp=time.time(),
                    )
                )

            except ScanFailedError:
                raise

        if current_content != original_content_snapshot:
            return current_content.get(content_field, ""), actions_taken
        return original_content_snapshot.get(content_field, ""), actions_taken

    def _execute_remote_batch(
        self,
        remote_scanners: list[BaseRemoteScanner],
        content: dict[str, Any],
        context: ScanContext,
        api_key: str | None = None,
        application_id: str | None = None,
    ) -> list[ScanResult]:
        """Execute remote scanners in a single batch API call.

        This is a synchronous method using connection pooling for low latency.

        Args:
            remote_scanners: List of remote scanner instances
            content: Content dictionary to scan
            context: Scan context
            api_key: API key for authentication
            application_id: Application ID (required by some remote APIs)

        Returns:
            List of ScanResult objects
        """
        logger = get_logger()
        logger.scan_started(
            f"{context.scan_type.value}_remote_batch",
            len(remote_scanners),
            "batch",
        )

        with Timer() as timer:
            try:
                results, api_call_ms = BaseRemoteScanner.execute_batch(
                    remote_scanners,
                    content,
                    context,
                    api_key,
                    application_id,
                )
                # Use remote API call latency (not overall timer) when API didn't provide per-result latency
                for result in results:
                    if result.execution_time_ms == 0.0 and api_call_ms > 0:
                        result.execution_time_ms = api_call_ms / len(remote_scanners)
                        result.metadata["latency"] = result.execution_time_ms
                    logger.scanner_completed(
                        result.scanner_name,
                        result.passed,
                        result.execution_time_ms,
                        len(result.violations),
                    )
            except Exception as e:
                logger.scanner_error("remote_batch", e)
                # Single error result: empty componentResults, error in error column (like local scanners)
                results = [
                    ScanResult(
                        passed=False,
                        scanner_name="remote_batch",
                        scan_type=context.scan_type,
                        violations=[],
                        execution_time_ms=timer.elapsed_ms,
                        metadata={"error": str(e), "latency": timer.elapsed_ms},
                    )
                ]

        return results

    async def _execute_remote_batch_async(
        self,
        remote_scanners: list[BaseRemoteScanner],
        content: dict[str, Any],
        context: ScanContext,
        api_key: str | None = None,
        application_id: str | None = None,
    ) -> list[ScanResult]:
        """Execute remote scanners in a single async batch API call."""
        logger = get_logger()
        logger.scan_started(
            f"{context.scan_type.value}_remote_batch",
            len(remote_scanners),
            "batch",
        )

        with Timer() as timer:
            try:
                results, api_call_ms = await BaseRemoteScanner.execute_batch_async(
                    remote_scanners,
                    content,
                    context,
                    api_key,
                    application_id,
                )
                for result in results:
                    if result.execution_time_ms == 0.0 and api_call_ms > 0:
                        result.execution_time_ms = api_call_ms / len(remote_scanners)
                        result.metadata["latency"] = result.execution_time_ms
                    logger.scanner_completed(
                        result.scanner_name,
                        result.passed,
                        result.execution_time_ms,
                        len(result.violations),
                    )
            except Exception as e:
                logger.scanner_error("remote_batch", e)
                # Single error result: empty componentResults, error in error column (like local scanners)
                results = [
                    ScanResult(
                        passed=False,
                        scanner_name="remote_batch",
                        scan_type=context.scan_type,
                        violations=[],
                        execution_time_ms=timer.elapsed_ms,
                        metadata={"error": str(e), "latency": timer.elapsed_ms},
                    )
                ]

        return results

    @classmethod
    def _get_threadpool(cls) -> ThreadPoolExecutor:
        if cls._threadpool is None:
            with cls._threadpool_lock:
                if cls._threadpool is None:
                    cls._threadpool = ThreadPoolExecutor()
        return cls._threadpool

    @staticmethod
    async def _gather_results(
        remote_task,
        local_task,
    ) -> tuple[list[ScanResult], list[ScanResult]]:
        remote_results, local_results = await asyncio.gather(remote_task, local_task)
        return remote_results, local_results
