"""Wrapper classes for LLM provider clients."""

from __future__ import annotations

from collections.abc import AsyncGenerator, Iterator
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from tokpulse.tracker import TokpulseTracker


def _get_thinking_tokens(usage: Any) -> int | None:
    """Extract thinking tokens from a usage object.

    Prefers ``thinking_input_tokens`` (newer Anthropic field) over
    ``thinking_tokens`` (legacy field).  Returns ``None`` when neither
    is present.  Preserves zero values — a zero explicitly set in the
    API response is meaningful and must not be silently dropped.
    """
    t1: int | None = getattr(usage, "thinking_input_tokens", None)
    if t1 is not None:
        return t1
    result: int | None = getattr(usage, "thinking_tokens", None)
    return result


# ---------------------------------------------------------------------------
# Context manager wrappers for Anthropic MessageStreamManager
# ---------------------------------------------------------------------------


class _TrackedStreamManager:
    """Sync context manager wrapping Anthropic ``MessageStreamManager`` to track usage on exit."""

    def __init__(self, stream_manager: Any, tracker: TokpulseTracker, model: str) -> None:
        self._stream_manager = stream_manager
        self._tracker = tracker
        self._model = model
        self._stream: Any = None

    def __enter__(self) -> Any:
        self._stream = self._stream_manager.__enter__()
        return self._stream

    def __exit__(self, *args: Any) -> Any:
        result = self._stream_manager.__exit__(*args)
        if self._stream is not None:
            try:
                final = self._stream.get_final_message()
                usage = final.usage
                self._tracker.track(
                    provider="anthropic",
                    model=getattr(final, "model", None) or self._model,
                    input_tokens=getattr(usage, "input_tokens", 0),
                    output_tokens=getattr(usage, "output_tokens", 0),
                    cached_read_tokens=getattr(usage, "cache_read_input_tokens", None),
                    cached_write_tokens=getattr(usage, "cache_creation_input_tokens", None),
                    thinking_tokens=_get_thinking_tokens(usage),
                )
            except Exception:
                pass
        return result


class _AsyncTrackedStreamManager:
    """Async context manager wrapping ``AsyncMessageStreamManager`` to track usage on exit."""

    def __init__(self, stream_manager: Any, tracker: TokpulseTracker, model: str) -> None:
        self._stream_manager = stream_manager
        self._tracker = tracker
        self._model = model
        self._stream: Any = None

    async def __aenter__(self) -> Any:
        self._stream = await self._stream_manager.__aenter__()
        return self._stream

    async def __aexit__(self, *args: Any) -> Any:
        result = await self._stream_manager.__aexit__(*args)
        if self._stream is not None:
            try:
                final = self._stream.get_final_message()
                usage = final.usage
                self._tracker.track(
                    provider="anthropic",
                    model=getattr(final, "model", None) or self._model,
                    input_tokens=getattr(usage, "input_tokens", 0),
                    output_tokens=getattr(usage, "output_tokens", 0),
                    cached_read_tokens=getattr(usage, "cache_read_input_tokens", None),
                    cached_write_tokens=getattr(usage, "cache_creation_input_tokens", None),
                    thinking_tokens=_get_thinking_tokens(usage),
                )
            except Exception:
                pass
        return result


# ---------------------------------------------------------------------------
# Synchronous wrappers
# ---------------------------------------------------------------------------


class OpenAIWrapper:
    """
    Wrapper for OpenAI client that automatically tracks usage.

    The wrapper transparently proxies all calls to the underlying client
    while capturing token usage from responses.
    """

    def __init__(self, client: Any, tracker: TokpulseTracker):
        self._client = client
        self._tracker = tracker

    @property
    def chat(self) -> OpenAIChatWrapper:
        """Access the chat completions API."""
        return OpenAIChatWrapper(self._client.chat, self._tracker)

    def __getattr__(self, name: str) -> Any:
        """Proxy all other attributes to the underlying client."""
        return getattr(self._client, name)


class OpenAIChatWrapper:
    """Wrapper for OpenAI chat API."""

    def __init__(self, chat: Any, tracker: TokpulseTracker):
        self._chat = chat
        self._tracker = tracker

    @property
    def completions(self) -> OpenAICompletionsWrapper:
        """Access the completions API."""
        return OpenAICompletionsWrapper(self._chat.completions, self._tracker)

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


class OpenAICompletionsWrapper:
    """Wrapper for OpenAI chat completions API."""

    def __init__(self, completions: Any, tracker: TokpulseTracker):
        self._completions = completions
        self._tracker = tracker

    def create(self, **kwargs: Any) -> Any:
        """Create a chat completion and track usage."""
        stream = kwargs.get("stream", False)
        response = self._completions.create(**kwargs)

        if stream:
            return self._wrap_stream(response, kwargs.get("model", "unknown"))

        self._track_response(response)
        return response

    def _track_response(self, response: Any) -> None:
        """Track token usage from a non-streaming response."""
        if not hasattr(response, "usage") or response.usage is None:
            return

        usage = response.usage
        model = getattr(response, "model", "unknown")

        cached_tokens = None
        if hasattr(usage, "prompt_tokens_details"):
            details = usage.prompt_tokens_details
            if details and hasattr(details, "cached_tokens"):
                cached_tokens = details.cached_tokens

        self._tracker.track(
            provider="openai",
            model=model,
            input_tokens=usage.prompt_tokens,
            output_tokens=usage.completion_tokens,
            cached_input_tokens=cached_tokens,
        )

    def _wrap_stream(self, stream: Any, model: str) -> Iterator[Any]:
        """Wrap a streaming response to track usage on completion or close."""
        input_tokens = 0
        output_tokens = 0
        cached_tokens = None

        try:
            for chunk in stream:
                if hasattr(chunk, "usage") and chunk.usage is not None:
                    input_tokens = chunk.usage.prompt_tokens
                    output_tokens = chunk.usage.completion_tokens

                    if hasattr(chunk.usage, "prompt_tokens_details"):
                        details = chunk.usage.prompt_tokens_details
                        if details and hasattr(details, "cached_tokens"):
                            cached_tokens = details.cached_tokens

                yield chunk
        finally:
            if input_tokens > 0 or output_tokens > 0:
                self._tracker.track(
                    provider="openai",
                    model=model,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    cached_input_tokens=cached_tokens,
                )

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


class AnthropicWrapper:
    """
    Wrapper for Anthropic client that automatically tracks usage.

    Handles Anthropic-specific features like cache read/write tokens
    and extended thinking tokens.
    """

    def __init__(self, client: Any, tracker: TokpulseTracker):
        self._client = client
        self._tracker = tracker

    @property
    def messages(self) -> AnthropicMessagesWrapper:
        """Access the messages API."""
        return AnthropicMessagesWrapper(self._client.messages, self._tracker)

    @property
    def beta(self) -> AnthropicBetaWrapper:
        """Access the beta API namespace with tracking."""
        return AnthropicBetaWrapper(self._client.beta, self._tracker)

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


class AnthropicMessagesWrapper:
    """Wrapper for Anthropic messages API."""

    def __init__(self, messages: Any, tracker: TokpulseTracker):
        self._messages = messages
        self._tracker = tracker

    def create(self, **kwargs: Any) -> Any:
        """Create a message and track usage."""
        stream = kwargs.get("stream", False)
        response = self._messages.create(**kwargs)

        if stream:
            return self._wrap_stream(response, kwargs.get("model", "unknown"))

        self._track_response(response)
        return response

    def stream(self, **kwargs: Any) -> _TrackedStreamManager:
        """Create a streaming message via the context manager API and track usage on exit.

        Usage::

            with client.messages.stream(model=..., ...) as s:
                for text in s.text_stream:
                    print(text)
        """
        return _TrackedStreamManager(
            self._messages.stream(**kwargs),
            self._tracker,
            kwargs.get("model", "unknown"),
        )

    def _track_response(self, response: Any) -> None:
        """Track token usage from a non-streaming response."""
        if not hasattr(response, "usage") or response.usage is None:
            return

        usage = response.usage
        model = getattr(response, "model", "unknown")

        self._tracker.track(
            provider="anthropic",
            model=model,
            input_tokens=getattr(usage, "input_tokens", 0),
            output_tokens=getattr(usage, "output_tokens", 0),
            cached_read_tokens=getattr(usage, "cache_read_input_tokens", None),
            cached_write_tokens=getattr(usage, "cache_creation_input_tokens", None),
            thinking_tokens=_get_thinking_tokens(usage),
        )

    def _wrap_stream(self, stream: Any, model: str) -> Iterator[Any]:
        """Wrap a streaming response to track usage on completion or close."""
        input_tokens = 0
        output_tokens = 0
        cache_read = None
        cache_write = None
        thinking = None

        try:
            for event in stream:
                if hasattr(event, "type"):
                    if event.type == "message_start" and hasattr(event, "message"):
                        msg = event.message
                        if hasattr(msg, "usage"):
                            input_tokens = getattr(msg.usage, "input_tokens", 0)
                            cache_read = getattr(msg.usage, "cache_read_input_tokens", None)
                            cache_write = getattr(msg.usage, "cache_creation_input_tokens", None)
                            thinking = _get_thinking_tokens(msg.usage)

                    elif event.type == "message_delta" and hasattr(event, "usage"):
                        output_tokens = getattr(event.usage, "output_tokens", 0)

                yield event
        finally:
            if input_tokens > 0 or output_tokens > 0:
                self._tracker.track(
                    provider="anthropic",
                    model=model,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    cached_read_tokens=cache_read,
                    cached_write_tokens=cache_write,
                    thinking_tokens=thinking,
                )

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


class AnthropicBetaWrapper:
    """Wrapper for the Anthropic ``client.beta`` namespace that intercepts ``messages`` calls."""

    def __init__(self, beta: Any, tracker: TokpulseTracker) -> None:
        self._beta = beta
        self._tracker = tracker

    @property
    def messages(self) -> AnthropicMessagesWrapper:
        """Access the beta messages API with tracking."""
        return AnthropicMessagesWrapper(self._beta.messages, self._tracker)

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


class GoogleWrapper:
    """Wrapper for Google GenerativeAI model that automatically tracks usage."""

    def __init__(self, model: Any, tracker: TokpulseTracker):
        self._model = model
        self._tracker = tracker
        self._model_name = getattr(model, "model_name", "unknown")

    def generate_content(self, *args: Any, **kwargs: Any) -> Any:
        """Generate content and track usage."""
        stream = kwargs.get("stream", False)
        response = self._model.generate_content(*args, **kwargs)

        if stream:
            return self._wrap_stream(response)

        self._track_response(response)
        return response

    def _track_response(self, response: Any) -> None:
        """Track token usage from a non-streaming response."""
        if not hasattr(response, "usage_metadata"):
            return

        usage = response.usage_metadata

        cached_tokens = None
        if hasattr(usage, "cached_content_token_count"):
            cached_tokens = usage.cached_content_token_count

        self._tracker.track(
            provider="google",
            model=self._model_name,
            input_tokens=getattr(usage, "prompt_token_count", 0),
            output_tokens=getattr(usage, "candidates_token_count", 0),
            cached_input_tokens=cached_tokens,
        )

    def _wrap_stream(self, stream: Any) -> Iterator[Any]:
        """Wrap a streaming response to track usage on completion or close."""
        input_tokens = 0
        output_tokens = 0
        cached_tokens = None

        try:
            for chunk in stream:
                if hasattr(chunk, "usage_metadata"):
                    usage = chunk.usage_metadata
                    input_tokens = getattr(usage, "prompt_token_count", input_tokens)
                    output_tokens = getattr(usage, "candidates_token_count", output_tokens)

                    if hasattr(usage, "cached_content_token_count"):
                        cached_tokens = usage.cached_content_token_count

                yield chunk
        finally:
            if input_tokens > 0 or output_tokens > 0:
                self._tracker.track(
                    provider="google",
                    model=self._model_name,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    cached_input_tokens=cached_tokens,
                )

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


# ---------------------------------------------------------------------------
# Asynchronous wrappers
# ---------------------------------------------------------------------------


class AsyncOpenAIWrapper:
    """
    Wrapper for async OpenAI client that automatically tracks usage.

    Use with ``AsyncOpenAI`` from the ``openai`` package.

    Example::

        tracker = TokpulseTracker(api_key="lct_xxx")
        client = tracker.openai_async(AsyncOpenAI())
        response = await client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": "Hello!"}],
        )
        await tracker.aflush()
    """

    def __init__(self, client: Any, tracker: TokpulseTracker):
        self._client = client
        self._tracker = tracker

    @property
    def chat(self) -> AsyncOpenAIChatWrapper:
        """Access the chat completions API."""
        return AsyncOpenAIChatWrapper(self._client.chat, self._tracker)

    def __getattr__(self, name: str) -> Any:
        """Proxy all other attributes to the underlying client."""
        return getattr(self._client, name)


class AsyncOpenAIChatWrapper:
    """Wrapper for async OpenAI chat API."""

    def __init__(self, chat: Any, tracker: TokpulseTracker):
        self._chat = chat
        self._tracker = tracker

    @property
    def completions(self) -> AsyncOpenAICompletionsWrapper:
        """Access the completions API."""
        return AsyncOpenAICompletionsWrapper(self._chat.completions, self._tracker)

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


class AsyncOpenAICompletionsWrapper:
    """Wrapper for async OpenAI chat completions API."""

    def __init__(self, completions: Any, tracker: TokpulseTracker):
        self._completions = completions
        self._tracker = tracker

    async def create(self, **kwargs: Any) -> Any:
        """Create a chat completion and track usage."""
        stream = kwargs.get("stream", False)
        response = await self._completions.create(**kwargs)

        if stream:
            return self._wrap_stream(response, kwargs.get("model", "unknown"))

        self._track_response(response)
        return response

    def _track_response(self, response: Any) -> None:
        """Track token usage from a non-streaming response."""
        if not hasattr(response, "usage") or response.usage is None:
            return

        usage = response.usage
        model = getattr(response, "model", "unknown")

        cached_tokens = None
        if hasattr(usage, "prompt_tokens_details"):
            details = usage.prompt_tokens_details
            if details and hasattr(details, "cached_tokens"):
                cached_tokens = details.cached_tokens

        self._tracker.track(
            provider="openai",
            model=model,
            input_tokens=usage.prompt_tokens,
            output_tokens=usage.completion_tokens,
            cached_input_tokens=cached_tokens,
        )

    async def _wrap_stream(self, stream: Any, model: str) -> AsyncGenerator[Any, None]:
        """Wrap a streaming response to track usage on completion or close."""
        input_tokens = 0
        output_tokens = 0
        cached_tokens = None

        try:
            async for chunk in stream:
                if hasattr(chunk, "usage") and chunk.usage is not None:
                    input_tokens = chunk.usage.prompt_tokens
                    output_tokens = chunk.usage.completion_tokens

                    if hasattr(chunk.usage, "prompt_tokens_details"):
                        details = chunk.usage.prompt_tokens_details
                        if details and hasattr(details, "cached_tokens"):
                            cached_tokens = details.cached_tokens

                yield chunk
        finally:
            if input_tokens > 0 or output_tokens > 0:
                self._tracker.track(
                    provider="openai",
                    model=model,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    cached_input_tokens=cached_tokens,
                )

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


class AsyncAnthropicWrapper:
    """
    Wrapper for async Anthropic client that automatically tracks usage.

    Use with ``AsyncAnthropic`` from the ``anthropic`` package.

    Example::

        tracker = TokpulseTracker(api_key="lct_xxx")
        client = tracker.anthropic_async(AsyncAnthropic())
        message = await client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=1024,
            messages=[{"role": "user", "content": "Hello!"}],
        )
        await tracker.aflush()
    """

    def __init__(self, client: Any, tracker: TokpulseTracker):
        self._client = client
        self._tracker = tracker

    @property
    def messages(self) -> AsyncAnthropicMessagesWrapper:
        """Access the messages API."""
        return AsyncAnthropicMessagesWrapper(self._client.messages, self._tracker)

    @property
    def beta(self) -> AsyncAnthropicBetaWrapper:
        """Access the beta API namespace with tracking."""
        return AsyncAnthropicBetaWrapper(self._client.beta, self._tracker)

    def __getattr__(self, name: str) -> Any:
        """Proxy all other attributes to the underlying client."""
        return getattr(self._client, name)


class AsyncAnthropicMessagesWrapper:
    """Wrapper for async Anthropic messages API."""

    def __init__(self, messages: Any, tracker: TokpulseTracker):
        self._messages = messages
        self._tracker = tracker

    async def create(self, **kwargs: Any) -> Any:
        """Create a message and track usage."""
        stream = kwargs.get("stream", False)
        response = await self._messages.create(**kwargs)

        if stream:
            return self._wrap_stream(response, kwargs.get("model", "unknown"))

        self._track_response(response)
        return response

    def stream(self, **kwargs: Any) -> _AsyncTrackedStreamManager:
        """Create a streaming message via the async context manager API and track usage on exit.

        Usage::

            async with client.messages.stream(model=..., ...) as s:
                async for text in s.text_stream:
                    print(text)
        """
        return _AsyncTrackedStreamManager(
            self._messages.stream(**kwargs),
            self._tracker,
            kwargs.get("model", "unknown"),
        )

    def _track_response(self, response: Any) -> None:
        """Track token usage from a non-streaming response."""
        if not hasattr(response, "usage") or response.usage is None:
            return

        usage = response.usage
        model = getattr(response, "model", "unknown")

        self._tracker.track(
            provider="anthropic",
            model=model,
            input_tokens=getattr(usage, "input_tokens", 0),
            output_tokens=getattr(usage, "output_tokens", 0),
            cached_read_tokens=getattr(usage, "cache_read_input_tokens", None),
            cached_write_tokens=getattr(usage, "cache_creation_input_tokens", None),
            thinking_tokens=_get_thinking_tokens(usage),
        )

    async def _wrap_stream(self, stream: Any, model: str) -> AsyncGenerator[Any, None]:
        """Wrap a streaming response to track usage on completion or close."""
        input_tokens = 0
        output_tokens = 0
        cache_read = None
        cache_write = None
        thinking = None

        try:
            async for event in stream:
                if hasattr(event, "type"):
                    if event.type == "message_start" and hasattr(event, "message"):
                        msg = event.message
                        if hasattr(msg, "usage"):
                            input_tokens = getattr(msg.usage, "input_tokens", 0)
                            cache_read = getattr(msg.usage, "cache_read_input_tokens", None)
                            cache_write = getattr(msg.usage, "cache_creation_input_tokens", None)
                            thinking = _get_thinking_tokens(msg.usage)

                    elif event.type == "message_delta" and hasattr(event, "usage"):
                        output_tokens = getattr(event.usage, "output_tokens", 0)

                yield event
        finally:
            if input_tokens > 0 or output_tokens > 0:
                self._tracker.track(
                    provider="anthropic",
                    model=model,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    cached_read_tokens=cache_read,
                    cached_write_tokens=cache_write,
                    thinking_tokens=thinking,
                )

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


class AsyncAnthropicBetaWrapper:
    """Async wrapper for the Anthropic ``client.beta`` namespace, intercepting ``messages``."""

    def __init__(self, beta: Any, tracker: TokpulseTracker) -> None:
        self._beta = beta
        self._tracker = tracker

    @property
    def messages(self) -> AsyncAnthropicMessagesWrapper:
        """Access the async beta messages API with tracking."""
        return AsyncAnthropicMessagesWrapper(self._beta.messages, self._tracker)

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