"""
Paved wrapper for Anthropic with policy enforcement.

Two modes:
  - **Gateway mode** (default when paved_api_key is provided):
    Routes all requests through the Paved gateway's Anthropic Messages endpoint.
    No client-side Anthropic 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 Anthropic directly.
"""
import os
from typing import Any, Dict, List, Optional, Union, Iterator, AsyncIterator

try:
    from anthropic import Anthropic as OriginalAnthropic
    from anthropic import AsyncAnthropic as OriginalAsyncAnthropic
    from anthropic import Stream, AsyncStream
    ANTHROPIC_AVAILABLE = True
except ImportError:
    ANTHROPIC_AVAILABLE = False
    OriginalAnthropic = object
    OriginalAsyncAnthropic = object

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 Anthropic api_key."""
    key = paved_api_key or os.getenv("PAVED_API_KEY")
    return bool(key) and "api_key" not in kwargs


class PavedMessages:
    """Wrapper for Anthropic messages API with policy enforcement (SDK-check mode)."""

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

    def _extract_user_content(self, messages: List[Dict[str, Any]]) -> str:
        user_content = []
        for msg in messages:
            if msg.get("role") == "user":
                content = msg.get("content", "")
                if isinstance(content, str):
                    user_content.append(content)
                elif isinstance(content, list):
                    for block in content:
                        if isinstance(block, dict) and block.get("type") == "text":
                            user_content.append(block.get("text", ""))
        return " ".join(user_content) if user_content else ""

    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"),
            "top_k": kwargs.get("top_k"),
            "tools": kwargs.get("tools"),
            "tool_choice": kwargs.get("tool_choice"),
            "system": kwargs.get("system"),
            "framework": "anthropic"
        }

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

    def stream(self, **kwargs) -> Union[Stream, Iterator]:
        payload = self._build_policy_payload(**kwargs)
        self._paved.check("llm", payload)
        return self._original.create(stream=True, **kwargs)


class AsyncPavedMessages:
    """Async wrapper for Anthropic messages API with policy enforcement (SDK-check mode)."""

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

    def _extract_user_content(self, messages: List[Dict[str, Any]]) -> str:
        user_content = []
        for msg in messages:
            if msg.get("role") == "user":
                content = msg.get("content", "")
                if isinstance(content, str):
                    user_content.append(content)
                elif isinstance(content, list):
                    for block in content:
                        if isinstance(block, dict) and block.get("type") == "text":
                            user_content.append(block.get("text", ""))
        return " ".join(user_content) if user_content else ""

    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"),
            "top_k": kwargs.get("top_k"),
            "tools": kwargs.get("tools"),
            "tool_choice": kwargs.get("tool_choice"),
            "system": kwargs.get("system"),
            "framework": "anthropic"
        }

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

    async def stream(self, **kwargs) -> Union[AsyncStream, AsyncIterator]:
        payload = self._build_policy_payload(**kwargs)
        self._paved.check("llm", payload)
        return await self._original.create(stream=True, **kwargs)


class Anthropic(OriginalAnthropic):
    """
    Paved drop-in replacement for Anthropic.

    Gateway mode (default — no Anthropic key needed):
        client = Anthropic(paved_api_key="pvd_live_...")
        response = client.messages.create(
            model="claude-sonnet-4-5-20250929",
            max_tokens=1024,
            messages=[{"role": "user", "content": "Hello!"}]
        )

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

    def __init__(
        self,
        agent_id: str = "anthropic-agent",
        paved_policies: Optional[List[str]] = None,
        paved_policy: Optional[Dict[str, Any]] = None,
        paved_api_key: Optional[str] = None,
        paved_url: Optional[str] = None,
        _paved: Optional[Agent] = None,
        **kwargs
    ):
        if not ANTHROPIC_AVAILABLE:
            raise ImportError(
                "anthropic package is required. Install with: pip install anthropic"
            )

        self._gateway_mode = _is_gateway_mode(paved_api_key, kwargs)

        if self._gateway_mode:
            # Gateway mode: route through Paved's Anthropic Messages endpoint
            resolved_url = _resolve_paved_url(paved_url)
            # The gateway Anthropic endpoint is at /gateway/v1/messages
            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 messages with policy enforcement
            if _paved:
                self._paved = _paved
            else:
                self._paved = Agent(
                    agent_id=agent_id,
                    policies=paved_policies or [],
                    policy_config=paved_policy,
                    api_key=paved_api_key,
                    base_url=_resolve_paved_url(paved_url)
                )
            self.messages = PavedMessages(super().messages, self._paved)


class AsyncAnthropic(OriginalAsyncAnthropic):
    """Async Paved drop-in replacement for Anthropic. Same gateway/SDK-check mode logic."""

    def __init__(
        self,
        agent_id: str = "anthropic-agent",
        paved_policies: Optional[List[str]] = None,
        paved_policy: Optional[Dict[str, Any]] = None,
        paved_api_key: Optional[str] = None,
        paved_url: Optional[str] = None,
        _paved: Optional[Agent] = None,
        **kwargs
    ):
        if not ANTHROPIC_AVAILABLE:
            raise ImportError(
                "anthropic package is required. Install with: pip install anthropic"
            )

        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:
            if _paved:
                self._paved = _paved
            else:
                self._paved = Agent(
                    agent_id=agent_id,
                    policies=paved_policies or [],
                    policy_config=paved_policy,
                    api_key=paved_api_key,
                    base_url=_resolve_paved_url(paved_url)
                )
            self.messages = AsyncPavedMessages(super().messages, self._paved)
