from typing import Any, Protocol, cast

from deconvolute.clients.base import BaseProxy
from deconvolute.detectors.base import BaseDetector
from deconvolute.errors import DeconvoluteError, ThreatDetectedError
from deconvolute.utils.logger import get_logger

logger = get_logger()


class Injector(Protocol):
    def inject(self, content: str) -> tuple[str, Any]: ...


class OpenAIProxy(BaseProxy):
    """
    Synchronous Proxy for the OpenAI client.

    This wrapper serves as a transparent middleware for the OpenAI SDK. It intercepts
    calls to `client.chat.completions.create` to inject security defenses (Canaries)
    and validate outputs (Content Scanning).

    All other calls (e.g. `client.embeddings`, `client.images`, `client.models`) are
    transparently delegated to the underlying client via the BaseProxy mechanism,
    ensuring full compatibility with the original SDK.
    """

    @property
    def chat(self) -> "ChatProxy":
        """
        Access the intercepted 'chat' namespace.

        Returns:
            ChatProxy: A wrapper around the original `client.chat` object that
                enables interception of completion creation calls.
        """
        return ChatProxy(self._client.chat, self._injectors, self._scanners)


class AsyncOpenAIProxy(BaseProxy):
    """
    Asynchronous Proxy for the OpenAI client.

    Mirrors the behavior of `OpenAIProxy` but handles `await` calls for use in
    asyncio event loops. It intercepts `await client.chat.completions.create`.
    """

    @property
    def chat(self) -> "AsyncChatProxy":
        """
        Access the intercepted async 'chat' namespace.

        Returns:
            AsyncChatProxy: A wrapper around the original `client.chat` object.
        """
        return AsyncChatProxy(self._client.chat, self._injectors, self._scanners)


class ChatProxy:
    """
    Helper proxy to navigate the `client.chat` namespace.

    This class exists to bridge the gap between `client` and `client.chat.completions`.
    It forwards all unknown attributes to the real OpenAI chat object.
    """

    def __init__(
        self,
        chat_module: Any,
        injectors: list[BaseDetector],
        scanners: list[BaseDetector],
    ):
        self._chat_module = chat_module
        self._injectors = injectors
        self._scanners = scanners

    def __getattr__(self, name: str) -> Any:
        # Pass through other chat methods (e.g. format)
        return getattr(self._chat_module, name)

    @property
    def completions(self) -> "CompletionsProxy":
        """
        Access the intercepted 'completions' namespace.

        Returns:
            CompletionsProxy: The core interceptor that wraps the `create` method.
        """
        return CompletionsProxy(
            self._chat_module.completions, self._injectors, self._scanners
        )


class CompletionsProxy:
    """
    The core interceptor for Synchronous Chat Completions.

    This class handles the lifecycle of the security checks:
    1. Input: Modifies the 'messages' list (e.g. Canary injection).
    2. Execution: Calls the real OpenAI client.
    3. Output: Validates and sanitizes the response content.
    """

    def __init__(
        self,
        completions_module: Any,
        injectors: list[BaseDetector],
        scanners: list[BaseDetector],
    ):
        self._module = completions_module
        self._injectors = injectors
        self._scanners = scanners

    def create(self, *args: Any, **kwargs: Any) -> Any:
        """
        Wraps the standard `openai.chat.completions.create` method.

        This method intercepts the request to inject security payloads and intercepts
        the response to validate content. If `stream=True` is detected, output
        validation is currently skipped to preserve streaming behavior.

        Args:
            *args: Positional arguments forwarded to OpenAI.
            **kwargs: Keyword arguments (messages, model, tools, etc.) forwarded to
                OpenAI.

        Returns:
            ChatCompletion: The original OpenAI response object, but with its content
            sanitized (e.g. tokens removed) and validated.

        Raises:
            ThreatDetectedError: If a detector (e.g. Language, Canary) identifies a
                threat in any of the generated choices.
            DeconvoluteError: If integrity checks are enabled but the request
                configuration is invalid (e.g. missing a 'system' message).
        """

        # 1. Input (Inject)
        # layer_states holds the secrets (tokens) generated by injectors
        # to be verified later by the scanners.
        layer_states: dict[BaseDetector, Any] = {}

        if self._injectors and "messages" in kwargs:
            self._apply_input_modifiers(kwargs["messages"], layer_states)

        # 2. Execution (Call)
        # Pass-through: If streaming is enabled, we skip output validation
        # because we cannot easily block and scan a generator yet.
        if kwargs.get("stream"):
            logger.warning(
                "Deconvolute: Streaming detected. Output validation (scanners) skipped."
            )
            return self._module.create(*args, **kwargs)

        # Standard blocking call
        response = self._module.create(*args, **kwargs)

        # 3. Output (Check & Clean)
        if self._scanners:
            self._apply_output_validators(response, layer_states)

        return response

    def _apply_input_modifiers(
        self, messages: list[dict[str, str]], layer_states: dict[BaseDetector, str]
    ) -> None:
        """
        Internal helper to execute 'Injectors' on the system prompt.

        It locates the system message, passes its content through all active input
        detectors (e.g. Canary), and updates the message list in-place.
        """
        # Find system message index
        idx = next(
            (i for i, m in enumerate(messages) if m.get("role") == "system"), None
        )

        if idx is None:
            # We strictly require a system prompt if using Integrity checks
            raise DeconvoluteError(
                "Integrity defenses enabled, but no 'system' message found in request. "
                "You must provide a system prompt to use Canary/Integrity checks."
            )

        content = messages[idx]["content"]

        for detector in self._injectors:
            # The detector does the logic; we just handle the plumbing.
            injector = cast(Injector, detector)
            content_with_token, token = injector.inject(content)
            layer_states[detector] = token

        # Update the message list in-place
        messages[idx]["content"] = content_with_token

    def _apply_output_validators(
        self, response: Any, layer_states: dict[BaseDetector, str]
    ) -> None:
        """
        Internal helper to execute 'Scanners' on the response object.

        It iterates through ALL choices in the response (handling `n > 1`).
        If a threat is found in any choice, the entire response is blocked.
        If the response is safe, it cleans artifacts (like Canary tokens) from
        all choices in-place.
        """
        # Iterate over every generated choice (usually 1, but could be n > 1)
        for choice in response.choices:
            message_obj = choice.message
            content = message_obj.content

            # Skip empty content (e.g. pure function calls or tool use)
            if not content:
                continue

            for detector in self._scanners:
                token = layer_states.get(detector)

                # If any choice fails, the entire batch is compromised.
                result = detector.check(content, token=token)

                if result.threat_detected:
                    # We define the component as "Choice X -> Detector Y" for clarity
                    raise ThreatDetectedError(
                        (
                            "Threat detected in choice index "
                            f"{choice.index} by {result.component}"
                        ),
                        result=result,
                    )

                # We must clean Every choice so the user never sees the raw tokens.
                if hasattr(detector, "clean") and token:
                    content = detector.clean(content, token)

            # Mutate the choice in place
            message_obj.content = content


class AsyncChatProxy:
    """
    Helper proxy to navigate the Async `client.chat` namespace.
    """

    def __init__(
        self,
        chat_module: Any,
        injectors: list[BaseDetector],
        scanners: list[BaseDetector],
    ):
        self._chat_module = chat_module
        self._injectors = injectors
        self._scanners = scanners

    def __getattr__(self, name: str) -> Any:
        return getattr(self._chat_module, name)

    @property
    def completions(self) -> "AsyncCompletionsProxy":
        """
        Access the intercepted async 'completions' namespace.
        """
        return AsyncCompletionsProxy(
            self._chat_module.completions, self._injectors, self._scanners
        )


class AsyncCompletionsProxy:
    """
    The core interceptor for Asynchronous Chat Completions.

    Identical in logic to `CompletionsProxy` but uses `await` for execution
    and validation calls.
    """

    def __init__(
        self,
        completions_module: Any,
        injectors: list[BaseDetector],
        scanners: list[BaseDetector],
    ):
        self._module = completions_module
        self._injectors = injectors
        self._scanners = scanners

    async def create(self, *args: Any, **kwargs: Any) -> Any:
        """
        Wraps the standard `await openai.chat.completions.create` method.

        Args:
            *args: Positional arguments forwarded to OpenAI.
            **kwargs: Keyword arguments (messages, model, etc.).

        Returns:
            ChatCompletion: The sanitized OpenAI response object.

        Raises:
            ThreatDetectedError: If a threat is detected in the response.
            DeconvoluteError: If the request configuration is invalid.
        """
        layer_states: dict[BaseDetector, Any] = {}

        if self._injectors and "messages" in kwargs:
            self._apply_input_modifiers(kwargs["messages"], layer_states)

        if kwargs.get("stream"):
            logger.warning(
                "Deconvolute: Streaming detected. Output validation skipped."
            )
            return await self._module.create(*args, **kwargs)

        response = await self._module.create(*args, **kwargs)

        if self._scanners:
            await self._apply_output_validators(response, layer_states)

        return response

    def _apply_input_modifiers(
        self, messages: list[dict[str, str]], layer_states: dict[BaseDetector, str]
    ) -> None:
        """Duplicated helper to avoid inheritance complexity across Sync/Async."""
        idx = next(
            (i for i, m in enumerate(messages) if m.get("role") == "system"), None
        )
        if idx is None:
            raise DeconvoluteError(
                "Integrity defenses enabled, but no 'system' message found."
            )

        content = messages[idx]["content"]
        for detector in self._injectors:
            injector = cast(Injector, detector)
            content_with_token, token = injector.inject(content)
            layer_states[detector] = token
        messages[idx]["content"] = content_with_token

    async def _apply_output_validators(
        self, response: Any, layer_states: dict[BaseDetector, str]
    ) -> None:
        """Async version of output validator."""
        # Iterate over every generated choice (usually 1, but could be n > 1)
        for choice in response.choices:
            message_obj = choice.message
            content = message_obj.content

            # Skip empty content (e.g. pure function calls or tool use)
            if not content:
                continue

            for detector in self._scanners:
                token = layer_states.get(detector)

                # If any choice fails, the entire batch is compromised.
                result = await detector.a_check(content, token=token)

                if result.threat_detected:
                    # We define the component as "Choice X -> Detector Y" for clarity
                    raise ThreatDetectedError(
                        (
                            "Threat detected in choice index "
                            f"{choice.index} by {result.component}"
                        ),
                        result=result,
                    )

                # We must clean Every choice so the user never sees the raw tokens.
                if hasattr(detector, "clean") and token:
                    content = detector.clean(content, token)

            # Mutate the choice in place
            message_obj.content = content
