"""
Paved OpenAI API Wrapper.

Two modes:
  - **Gateway mode** (default when paved_api_key is provided):
    Routes all requests through the Paved gateway which handles policy checks,
    cost tracking, and audit logging server-side. No client-side OpenAI key needed.

  - **SDK-check mode** (legacy, when api_key is provided but no paved_api_key):
    Client-side policy checks via Guard API, then forwards to OpenAI directly.
"""
import os
from typing import Any, Dict, List, Optional, Union, Iterator, AsyncIterator
from openai import OpenAI as OriginalOpenAI
from openai import AsyncOpenAI as OriginalAsyncOpenAI
from openai import Stream, AsyncStream

from pvd.sdk import Agent

# Default Paved platform URL
_DEFAULT_PAVED_URL = "https://api.hipaved.com/v1"


def _resolve_paved_url(paved_url: Optional[str] = None) -> str:
    return paved_url or os.getenv("PAVED_URL", _DEFAULT_PAVED_URL)


def _is_gateway_mode(paved_api_key: Optional[str], kwargs: dict) -> bool:
    """Gateway mode when paved_api_key is provided and no explicit OpenAI api_key."""
    key = paved_api_key or os.getenv("PAVED_API_KEY")
    has_openai_key = "api_key" in kwargs or os.getenv("OPENAI_API_KEY")
    # Gateway mode if we have a paved key and user didn't explicitly provide an OpenAI key
    return bool(key) and "api_key" not in kwargs


class PavedOpenAIWrapper:
    """Mixin for OpenAI wrappers to add policy enforcement."""

    def __init__(
        self,
        agent_id: str = "openai-agent",
        paved_api_key: str = None,
        paved_url: str = None,
        paved_policies: List[str] = None,
        paved_policy: Dict[str, Any] = None,
        **kwargs
    ):
        # Initialize Paved agent for SDK-check mode
        self._paved = Agent(
            agent_id=agent_id,
            policies=paved_policies or [],
            policy_config=paved_policy,
            base_url=_resolve_paved_url(paved_url),
            api_key=paved_api_key or os.getenv("PAVED_API_KEY")
        )


class PavedOpenAICompletions:
    """Wrapper for OpenAI completions with policy enforcement (SDK-check mode)."""

    def __init__(self, original_completions, paved_agent):
        self._original = original_completions
        self._paved = paved_agent

    def _extract_user_content(self, messages: List[Dict]) -> str:
        return "\n".join([
            m.get("content", "") for m in messages
            if isinstance(m, dict) and m.get("role") == "user"
        ])

    def _build_policy_payload(self, **kwargs) -> Dict[str, Any]:
        messages = kwargs.get("messages", [])
        return {
            "text": self._extract_user_content(messages),
            "messages": messages,
            "model": kwargs.get("model", "unknown"),
            "max_tokens": kwargs.get("max_tokens"),
            "temperature": kwargs.get("temperature"),
            "top_p": kwargs.get("top_p"),
            "frequency_penalty": kwargs.get("frequency_penalty"),
            "presence_penalty": kwargs.get("presence_penalty"),
            "tool_choice": kwargs.get("tool_choice"),
            "tools": kwargs.get("tools"),
            "response_format": kwargs.get("response_format"),
            "stream": kwargs.get("stream", False),
            "framework": "openai"
        }

    def create(self, **kwargs):
        payload = self._build_policy_payload(**kwargs)
        self._paved.check("llm", payload)
        return self._original.create(**kwargs)


class PavedOpenAIChat:
    def __init__(self, original_chat, paved_agent):
        self.completions = PavedOpenAICompletions(original_chat.completions, paved_agent)


class PavedOpenAIEmbeddings:
    def __init__(self, original_embeddings, paved_agent):
        self._original = original_embeddings
        self._paved = paved_agent

    def create(self, **kwargs):
        input_text = kwargs.get("input", "")
        model = kwargs.get("model", "unknown")
        if isinstance(input_text, list):
            text_content = " ".join(input_text)
        else:
            text_content = input_text
        payload = {
            "text": text_content,
            "input": input_text,
            "model": model,
            "encoding_format": kwargs.get("encoding_format"),
            "framework": "openai",
            "action": "embedding"
        }
        self._paved.check("embedding", payload)
        return self._original.create(**kwargs)


class OpenAI(OriginalOpenAI, PavedOpenAIWrapper):
    """
    Paved drop-in replacement for OpenAI.

    Gateway mode (default — no OpenAI key needed):
        client = OpenAI(paved_api_key="pvd_live_...")
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": "Hello!"}]
        )

    SDK-check mode (provide your own OpenAI key):
        client = OpenAI(api_key="sk-...", paved_api_key="pvd_live_...")
    """

    def __init__(
        self,
        agent_id: str = "openai-agent",
        paved_api_key: str = None,
        paved_url: str = None,
        paved_policies: List[str] = None,
        paved_policy: Dict[str, Any] = None,
        _paved: Optional[Agent] = None,
        **kwargs
    ):
        self._gateway_mode = _is_gateway_mode(paved_api_key, kwargs)

        if self._gateway_mode:
            # Gateway mode: route through Paved, use paved key as Bearer token
            resolved_url = _resolve_paved_url(paved_url)
            gateway_url = f"{resolved_url.rstrip('/')}/gateway"
            resolved_key = paved_api_key or os.getenv("PAVED_API_KEY")
            kwargs["base_url"] = gateway_url
            kwargs["api_key"] = resolved_key

        super().__init__(**kwargs)

        if not self._gateway_mode:
            # SDK-check mode: wrap with policy enforcement
            PavedOpenAIWrapper.__init__(
                self,
                agent_id=agent_id,
                paved_api_key=paved_api_key,
                paved_url=paved_url,
                paved_policies=paved_policies,
                paved_policy=paved_policy
            )
            if _paved:
                self._paved = _paved
            self.chat = PavedOpenAIChat(super().chat, self._paved)
            self.embeddings = PavedOpenAIEmbeddings(super().embeddings, self._paved)
        # In gateway mode, no wrapping needed — the gateway does all policy checks


class AsyncPavedOpenAICompletions:
    def __init__(self, original_completions, paved_agent):
        self._original = original_completions
        self._paved = paved_agent

    def _extract_user_content(self, messages: List[Dict]) -> str:
        return "\n".join([
            m.get("content", "") for m in messages
            if isinstance(m, dict) and m.get("role") == "user"
        ])

    def _build_policy_payload(self, **kwargs) -> Dict[str, Any]:
        messages = kwargs.get("messages", [])
        return {
            "text": self._extract_user_content(messages),
            "messages": messages,
            "model": kwargs.get("model", "unknown"),
            "max_tokens": kwargs.get("max_tokens"),
            "temperature": kwargs.get("temperature"),
            "top_p": kwargs.get("top_p"),
            "frequency_penalty": kwargs.get("frequency_penalty"),
            "presence_penalty": kwargs.get("presence_penalty"),
            "tool_choice": kwargs.get("tool_choice"),
            "tools": kwargs.get("tools"),
            "response_format": kwargs.get("response_format"),
            "stream": kwargs.get("stream", False),
            "framework": "openai"
        }

    async def create(self, **kwargs):
        payload = self._build_policy_payload(**kwargs)
        self._paved.check("llm", payload)
        return await self._original.create(**kwargs)


class AsyncPavedOpenAIChat:
    def __init__(self, original_chat, paved_agent):
        self.completions = AsyncPavedOpenAICompletions(original_chat.completions, paved_agent)


class AsyncPavedOpenAIEmbeddings:
    def __init__(self, original_embeddings, paved_agent):
        self._original = original_embeddings
        self._paved = paved_agent

    async def create(self, **kwargs):
        input_text = kwargs.get("input", "")
        model = kwargs.get("model", "unknown")
        if isinstance(input_text, list):
            text_content = " ".join(input_text)
        else:
            text_content = input_text
        payload = {
            "text": text_content,
            "input": input_text,
            "model": model,
            "encoding_format": kwargs.get("encoding_format"),
            "framework": "openai",
            "action": "embedding"
        }
        self._paved.check("embedding", payload)
        return await self._original.create(**kwargs)


class AsyncOpenAI(OriginalAsyncOpenAI, PavedOpenAIWrapper):
    """Async Paved drop-in replacement for OpenAI. Same gateway/SDK-check mode logic."""

    def __init__(
        self,
        agent_id: str = "openai-agent",
        paved_api_key: str = None,
        paved_url: str = None,
        paved_policies: List[str] = None,
        paved_policy: Dict[str, Any] = None,
        _paved: Optional[Agent] = None,
        **kwargs
    ):
        self._gateway_mode = _is_gateway_mode(paved_api_key, kwargs)

        if self._gateway_mode:
            resolved_url = _resolve_paved_url(paved_url)
            gateway_url = f"{resolved_url.rstrip('/')}/gateway"
            resolved_key = paved_api_key or os.getenv("PAVED_API_KEY")
            kwargs["base_url"] = gateway_url
            kwargs["api_key"] = resolved_key

        super().__init__(**kwargs)

        if not self._gateway_mode:
            PavedOpenAIWrapper.__init__(
                self,
                agent_id=agent_id,
                paved_api_key=paved_api_key,
                paved_url=paved_url,
                paved_policies=paved_policies,
                paved_policy=paved_policy
            )
            if _paved:
                self._paved = _paved
            self.chat = AsyncPavedOpenAIChat(super().chat, self._paved)
            self.embeddings = AsyncPavedOpenAIEmbeddings(super().embeddings, self._paved)
