from contextlib import nullcontext
from typing import Any, Dict, Optional
import importlib.util
from concurrent.futures import Future, ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
from threading import BoundedSemaphore
from loguru import logger as _logger
from brimley.core.models import BrimleyFunction, PythonFunction, SqlFunction, TemplateFunction, ApiFunction, CliFunction
from brimley.core.context import BrimleyContext
from brimley.execution.python_runner import PythonRunner
from brimley.execution.sql_runner import SqlRunner
from brimley.execution.jinja_runner import JinjaRunner
from brimley.execution.api_runner import ApiRunner
from brimley.execution.cli_runner import CliRunner
from brimley.utils.diagnostics import BrimleyExecutionError
from brimley.infrastructure.logging import (
    get_or_create_correlation_id,
    set_external_trace_id,
)

class Dispatcher:
    """
    Routes execution to the appropriate runner based on function type.
    """
    def __init__(self):
        self.python_runner = PythonRunner()
        self.sql_runner = SqlRunner()
        self.jinja_runner = JinjaRunner()
        self.api_runner = ApiRunner()
        self.cli_runner = CliRunner()
        self._executor: Optional[ThreadPoolExecutor] = None
        self._inflight_slots: Optional[BoundedSemaphore] = None
        self._runtime_signature: Optional[tuple[int, int, str]] = None

    def _ensure_runtime_controls(self, context: BrimleyContext) -> None:
        execution = context.execution
        runtime_signature = (
            execution.thread_pool_size,
            execution.queue.max_size,
            execution.queue.on_full,
        )
        if self._executor is not None and self._inflight_slots is not None and self._runtime_signature == runtime_signature:
            return

        if self._executor is not None:
            self._executor.shutdown(wait=False, cancel_futures=True)

        total_slots = execution.thread_pool_size + execution.queue.max_size
        self._executor = ThreadPoolExecutor(max_workers=execution.thread_pool_size)
        self._inflight_slots = BoundedSemaphore(value=total_slots)
        self._runtime_signature = runtime_signature

    def _resolve_timeout_seconds(self, func: BrimleyFunction, context: BrimleyContext) -> float:
        if func.timeout_seconds is not None:
            return float(func.timeout_seconds)
        return float(context.execution.timeout_seconds)

    def _dispatch_sync_call(
        self,
        func: BrimleyFunction,
        args: Dict[str, Any],
        context: BrimleyContext,
        runtime_injections: Optional[Dict[str, Any]],
        request_ctx: Optional[Any] = None,
    ) -> Any:
        if func.type == "python_function":
            if isinstance(func, PythonFunction):
                return self.python_runner.run(func, args, context, runtime_injections=runtime_injections, request_ctx=request_ctx)
        elif func.type == "sql_function":
            if isinstance(func, SqlFunction):
                return self.sql_runner.run(func, args, context)
        elif func.type == "template_function":
            if isinstance(func, TemplateFunction):
                return self.jinja_runner.run(func, args, context)
        elif func.type == "api_function":
            # NOTE(v0.9): stub intercept point for MockRegistry — do not remove.
            if isinstance(func, ApiFunction):
                return self.api_runner.run(func, args, context)
        elif func.type == "cli_function":
            # NOTE(v0.9): stub intercept point for MockRegistry — do not remove.
            if isinstance(func, CliFunction):
                return self.cli_runner.run(func, args, context)

        raise NotImplementedError(f"No runner for function type: {func.type} ({type(func)})")

    def _has_fastmcp_runtime_injection(
        self,
        runtime_injections: Optional[Dict[str, Any]],
    ) -> bool:
        if not runtime_injections:
            return False

        for key in (
            "mcp_context",
            "mcp",
            "ctx",
            "context",
            "fastmcp_context",
            "mcp.server.fastmcp.Context",
        ):
            if key in runtime_injections and runtime_injections[key] is not None:
                return True

        return False

    def _screen_for_prompt_injection(self, args: Dict[str, Any], func_name: str) -> None:
        """
        Optional llm-guard PromptInjection screening hook (v0.7).

        Enabled via ``security.prompt_injection_screening: true`` in ``brimley.yaml``.
        Requires the optional ``security`` extra: ``poetry install --extras security``.

        If ``llm-guard`` is not installed, logs a warning and skips screening
        gracefully — the hook is the structural commitment; the dependency is opt-in.
        """
        if importlib.util.find_spec("llm_guard") is None:
            _logger.warning(
                "Prompt injection screening is enabled but 'llm-guard' is not installed. "
                "Install with: poetry install --extras security. Skipping screening."
            )
            return

        try:
            from llm_guard.input_scanners import PromptInjection  # type: ignore[import]

            scanner = PromptInjection()
            for key, value in args.items():
                if not isinstance(value, str):
                    continue
                sanitized, is_valid, risk_score = scanner.scan("", value)
                if not is_valid:
                    _logger.warning(
                        "Prompt injection detected in argument '{}' for function '{}' "
                        "(risk score: {:.2f}). Call blocked.",
                        key,
                        func_name,
                        risk_score,
                    )
                    raise BrimleyExecutionError(
                        message=(
                            f"Prompt injection detected in argument '{key}'. "
                            "Call blocked by security screening."
                        ),
                        func_name=func_name,
                    )
        except BrimleyExecutionError:
            raise
        except Exception as exc:
            _logger.warning(
                "Prompt injection screening failed for function '{}': {}. Skipping.",
                func_name,
                exc,
            )

    def run(
        self,
        func: BrimleyFunction,
        args: Dict[str, Any],
        context: BrimleyContext,
        runtime_injections: Optional[Dict[str, Any]] = None,
    ) -> Any:
        # Establish / inherit correlation ID for this dispatch.
        get_or_create_correlation_id()

        # Propagate external trace ID from FastMCP request context when present.
        if runtime_injections:
            mcp_ctx = runtime_injections.get("mcp_context") or runtime_injections.get("mcp")
            if mcp_ctx is not None:
                request_id = getattr(mcp_ctx, "request_id", None)
                if request_id:
                    set_external_trace_id(str(request_id))

        _logger.trace("Dispatching function '{}' (type={})", func.name, func.type)

        # llm-guard PromptInjection screening (v0.7 hook — configurable, opt-in).
        # Enable in brimley.yaml: security.prompt_injection_screening: true
        # Install the optional extra: poetry install --extras security
        security_cfg = getattr(context.config, "security", None) or {}
        if isinstance(security_cfg, dict) and security_cfg.get("prompt_injection_screening", False):
            self._screen_for_prompt_injection(args, func.name)

        if func.type == "python_function" and self._has_fastmcp_runtime_injection(runtime_injections):
            scope_ctx = (
                context.container.request_scope(context)
                if context.container is not None
                else nullcontext(None)
            )
            with scope_ctx as request_ctx:
                try:
                    result = self._dispatch_sync_call(func, args, context, runtime_injections, request_ctx)
                    _logger.trace("Function '{}' completed (type={})", func.name, func.type)
                    return result
                except Exception as exc:
                    _logger.trace("Function '{}' failed (type={}): {}", func.name, func.type, exc)
                    raise

        self._ensure_runtime_controls(context)
        timeout_seconds = self._resolve_timeout_seconds(func, context)
        queue_strategy = context.execution.queue.on_full

        assert self._inflight_slots is not None
        assert self._executor is not None

        if queue_strategy == "block":
            acquired = self._inflight_slots.acquire(timeout=timeout_seconds)
        else:
            acquired = self._inflight_slots.acquire(blocking=False)

        if not acquired:
            raise BrimleyExecutionError(
                message=(
                    f"Execution queue is full for function '{func.name}' (on_full={queue_strategy})."
                ),
                func_name=func.name,
            )

        scope_ctx = (
            context.container.request_scope(context)
            if context.container is not None
            else nullcontext(None)
        )
        with scope_ctx as request_ctx:
            try:
                future: Future = self._executor.submit(
                    self._dispatch_sync_call,
                    func,
                    args,
                    context,
                    runtime_injections,
                    request_ctx,
                )
                try:
                    result = future.result(timeout=timeout_seconds)
                    _logger.trace("Function '{}' completed (type={})", func.name, func.type)
                    return result
                except FuturesTimeoutError as exc:
                    future.cancel()
                    _logger.trace("Function '{}' timed out after {}s (type={})", func.name, timeout_seconds, func.type)
                    raise BrimleyExecutionError(
                        message=f"Execution timed out after {timeout_seconds:.3f}s.",
                        func_name=func.name,
                    ) from exc
                except Exception as exc:
                    _logger.trace("Function '{}' failed (type={}): {}", func.name, func.type, exc)
                    raise
            finally:
                self._inflight_slots.release()
