"""
OpenAI Client Wrapper for Vaikora Security.

Wraps any OpenAI-compatible client (OpenAI, Azure OpenAI, OpenRouter, etc.)
with Vaikora policy enforcement, input validation, and audit logging.

Usage:
    from openai import OpenAI
    from vaikora import wrap_openai

    client = OpenAI(api_key="sk-...")
    # or OpenRouter:
    # client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key="or-...")

    secured = wrap_openai(client, api_key="vk_...", agent_id="your-agent-id")
    response = secured.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": "Hello!"}]
    )
"""
import logging
from typing import Any, Optional

logger = logging.getLogger("vaikora.openai_wrapper")


class _WrappedChatCompletions:
    """Wraps chat.completions to intercept create() calls."""

    def __init__(self, original_completions: Any, wrapper: "VaikoraOpenAIWrapper"):
        self._original = original_completions
        self._wrapper = wrapper

    def create(self, **kwargs: Any) -> Any:
        """Intercept chat.completions.create() with policy enforcement."""
        messages = kwargs.get("messages", [])

        # Extract the latest user message for policy evaluation
        user_message = ""
        for msg in reversed(messages):
            if msg.get("role") == "user":
                user_message = msg.get("content", "")
                break

        # Evaluate against Vaikora policies
        if user_message and self._wrapper._should_evaluate:
            try:
                import httpx
                eval_response = httpx.post(
                    f"{self._wrapper._base_url}/api/v1/actions/evaluate",
                    headers={
                        "Authorization": f"Bearer {self._wrapper._api_key}",
                        "Content-Type": "application/json",
                    },
                    json={
                        "agent_key": self._wrapper._agent_key or "",
                        "action_type": "llm.chat.completion",
                        "resource_type": "openai",
                        "resource_id": kwargs.get("model", "unknown"),
                        "payload": {
                            "message": user_message,
                            "content": user_message,
                            "model": kwargs.get("model", ""),
                        },
                    },
                    timeout=10.0,
                )
                if eval_response.status_code == 200:
                    result = eval_response.json()
                    if result.get("decision") == "block" or result.get("status") == "blocked":
                        from vaikora.exceptions import PolicyViolationError
                        raise PolicyViolationError(
                            message="Request blocked by Vaikora security policy.",
                            policy_id=result.get("policy_id"),
                        )
                    logger.debug("Vaikora policy check passed")
                else:
                    logger.warning(f"Vaikora policy check returned {eval_response.status_code}, allowing by default")
            except ImportError:
                logger.error("httpx is required for Vaikora wrapper. Install with: pip install httpx")
            except Exception as e:
                if "PolicyViolationError" in type(e).__name__:
                    raise
                logger.warning(f"Vaikora policy check failed ({e}), allowing by default")

        # Forward to original client
        return self._original.create(**kwargs)


class _WrappedChat:
    """Wraps client.chat to provide wrapped completions."""

    def __init__(self, original_chat: Any, wrapper: "VaikoraOpenAIWrapper"):
        self.completions = _WrappedChatCompletions(
            original_chat.completions, wrapper
        )


class VaikoraOpenAIWrapper:
    """
    Wraps any OpenAI-compatible client with Vaikora security.

    Works with:
    - OpenAI (`openai.OpenAI`)
    - Azure OpenAI (`openai.AzureOpenAI`)
    - OpenRouter (OpenAI client with custom base_url)
    - Any OpenAI-compatible provider

    All `chat.completions.create()` calls are intercepted and evaluated
    against your Vaikora security policies before being forwarded.
    """

    def __init__(
        self,
        client: Any,
        api_key: str,
        agent_key: Optional[str] = None,
        base_url: str = "https://api.vaikora.com",
        evaluate_policies: bool = True,
    ):
        """
        Wrap an OpenAI-compatible client with Vaikora security.

        Args:
            client: An OpenAI-compatible client instance
            api_key: Your Vaikora API key (Bearer token)
            agent_key: Your Vaikora agent key for policy evaluation
            base_url: Vaikora API base URL
            evaluate_policies: Whether to evaluate policies (default: True)
        """
        self._client = client
        self._api_key = api_key
        self._agent_key = agent_key
        self._base_url = base_url.rstrip("/")
        self._should_evaluate = evaluate_policies

        # Wrap the chat interface
        self.chat = _WrappedChat(client.chat, self)

    def __getattr__(self, name: str) -> Any:
        """Forward all other attribute access to the original client."""
        if name in ("chat", "_client", "_api_key", "_agent_key", "_base_url", "_should_evaluate"):
            return object.__getattribute__(self, name)
        return getattr(self._client, name)


def wrap_openai(
    client: Any,
    api_key: str,
    agent_key: Optional[str] = None,
    base_url: str = "https://api.vaikora.com",
) -> VaikoraOpenAIWrapper:
    """
    Wrap an OpenAI-compatible client with Vaikora security.

    This works with OpenAI, Azure OpenAI, OpenRouter, or any
    OpenAI-compatible provider.

    Args:
        client: An OpenAI client instance (OpenAI, AzureOpenAI, etc.)
        api_key: Your Vaikora API key
        agent_key: Your Vaikora agent key for policy evaluation
        base_url: Vaikora API base URL (default: https://api.vaikora.com)

    Returns:
        A wrapped client that enforces Vaikora policies

    Example:
        from openai import OpenAI
        from vaikora import wrap_openai

        # Standard OpenAI
        client = wrap_openai(OpenAI(), api_key="vk_xxx")

        # OpenRouter
        client = wrap_openai(
            OpenAI(base_url="https://openrouter.ai/api/v1", api_key="or-xxx"),
            api_key="vk_xxx",
            agent_key="agent_xxx",
        )

        # Azure OpenAI
        from openai import AzureOpenAI
        client = wrap_openai(
            AzureOpenAI(azure_endpoint="...", api_key="..."),
            api_key="vk_xxx",
        )

        # Use normally - policies are enforced automatically
        response = client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": "Hello!"}]
        )
    """
    return VaikoraOpenAIWrapper(
        client=client,
        api_key=api_key,
        agent_key=agent_key,
        base_url=base_url,
    )
