"""LLM Middleware Wrapper.

Drop-in guarded wrapper around OpenAI and Anthropic APIs that
automatically scans, sanitizes, logs, and optionally blocks prompts.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Literal, Optional

from contextguard.guard import Guard, ScanResult


@dataclass
class GuardedResponse:
    """Response from a guarded LLM call."""

    content: str
    scan_result: ScanResult
    blocked: bool = False
    raw_response: Optional[Any] = None

    def __repr__(self) -> str:
        status = "BLOCKED" if self.blocked else "OK"
        return (
            f"GuardedResponse(status={status}, "
            f"risk={self.scan_result.risk_score}, "
            f"content_length={len(self.content)})"
        )


class GuardedLLM:
    """Guarded wrapper around LLM provider APIs.

    Parameters
    ----------
    provider : str
        ``"openai"`` or ``"anthropic"``.
    api_key : str
        API key for the provider.
    model : str, optional
        Model name (defaults to provider's default).
    guard : Guard, optional
        Pre-configured Guard instance. If None, a default one is created.
    auto_sanitize : bool
        Automatically sanitize prompts before sending (default True).
    scan_output : bool
        Scan LLM output for sensitive data (default False).
    block_on_high_risk : bool
        Block requests with high risk score (default True).

    Examples
    --------
    >>> llm = GuardedLLM(provider="openai", api_key="sk-...")
    >>> resp = llm.chat("Summarize the quarterly report")
    >>> print(resp.content)
    """

    def __init__(
        self,
        provider: Literal["openai", "anthropic"] = "openai",
        api_key: str = "",
        model: Optional[str] = None,
        guard: Optional[Guard] = None,
        auto_sanitize: bool = True,
        scan_output: bool = False,
        block_on_high_risk: bool = True,
        block_threshold: float = 0.9,
    ) -> None:
        self.provider = provider
        self.api_key = api_key
        self.model = model
        self.auto_sanitize = auto_sanitize
        self.scan_output = scan_output
        self.block_on_high_risk = block_on_high_risk

        self.guard = guard or Guard(
            block_threshold=block_threshold,
            enable_logging=True,
        )

        self._client: Any = None
        self._init_client()

    def _init_client(self) -> None:
        """Lazy-initialise the provider client."""
        if self.provider == "openai":
            try:
                import openai

                self._client = openai.OpenAI(api_key=self.api_key)
                if not self.model:
                    self.model = "gpt-3.5-turbo"
            except ImportError:
                pass
        elif self.provider == "anthropic":
            try:
                import anthropic

                self._client = anthropic.Anthropic(api_key=self.api_key)
                if not self.model:
                    self.model = "claude-3-haiku-20240307"
            except ImportError:
                pass

    # ------------------------------------------------------------------
    # Core methods
    # ------------------------------------------------------------------

    def _process_input(self, prompt: str) -> tuple[str, ScanResult]:
        """Scan and optionally sanitize the input prompt."""
        scan_result = self.guard.scan(prompt)
        processed = prompt
        if self.auto_sanitize and (
            scan_result.secrets.is_detected or scan_result.pii.is_detected
        ):
            processed = self.guard.sanitize(prompt)
        return processed, scan_result

    def _process_output(self, output: str) -> str:
        """Optionally scan and sanitize LLM output."""
        if self.scan_output:
            return self.guard.sanitize(output)
        return output

    def chat(
        self,
        prompt: str,
        system: Optional[str] = None,
        **kwargs: Any,
    ) -> GuardedResponse:
        """Send a chat message through the guard pipeline.

        Flow: Scan → Sanitize → Log → Send → (Scan Output) → Return
        """
        processed, scan_result = self._process_input(prompt)

        # Block if risk is too high
        if self.block_on_high_risk and scan_result.blocked:
            return GuardedResponse(
                content="[BLOCKED] Request blocked due to high security risk.",
                scan_result=scan_result,
                blocked=True,
            )

        # Call the provider
        content = ""
        raw_response = None

        if self._client is None:
            return GuardedResponse(
                content=f"[ERROR] {self.provider} client not available. "
                f"Install with: pip install contextguard[{self.provider}]",
                scan_result=scan_result,
                blocked=False,
            )

        try:
            if self.provider == "openai":
                messages = []
                if system:
                    messages.append({"role": "system", "content": system})
                messages.append({"role": "user", "content": processed})
                raw_response = self._client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    **kwargs,
                )
                content = raw_response.choices[0].message.content or ""

            elif self.provider == "anthropic":
                raw_response = self._client.messages.create(
                    model=self.model,
                    max_tokens=kwargs.pop("max_tokens", 1024),
                    system=system or "",
                    messages=[{"role": "user", "content": processed}],
                    **kwargs,
                )
                content = raw_response.content[0].text if raw_response.content else ""

        except Exception as e:
            content = f"[ERROR] Provider call failed: {e}"

        # Optionally scan output
        content = self._process_output(content)

        return GuardedResponse(
            content=content,
            scan_result=scan_result,
            blocked=False,
            raw_response=raw_response,
        )

    def complete(
        self,
        prompt: str,
        **kwargs: Any,
    ) -> GuardedResponse:
        """Alias for chat() without a system message."""
        return self.chat(prompt, **kwargs)
