# Copyright 2026 Invoke Labs, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Implement Voker tracking example for Anthropic.

[See Here](https://app.voker.ai/docs/integrations/providers/anthropic).
"""

from __future__ import annotations

from collections.abc import Mapping
from functools import partial
from typing import TYPE_CHECKING, Any, Optional, overload

from voker.client import VokerClient

if TYPE_CHECKING:
    import httpx
    from anthropic import resources as anthropic_resources
    from anthropic import types as anthropic_types
    from anthropic.lib import streaming as anthropic_lib_streaming
    from anthropic.types import beta as anthropic_types_beta
    from anthropic.types import message as anthropic_types_message

try:
    import anthropic
    import httpx
except ImportError:
    raise ModuleNotFoundError

try:
    from anthropic import not_given
except ImportError:
    from anthropic import NOT_GIVEN as not_given


def _rm_not_given(
    data: dict[str, Any | anthropic.NotGiven | anthropic.Omit],
) -> dict[str, Any]:
    return {
        k: v
        for k, v in data.items()
        if not isinstance(v, (anthropic.NotGiven, anthropic.Omit))
    }


class Anthropic(anthropic.Anthropic):
    """Anthropic wrapper that automatically creates llm events in Voker."""

    @overload
    def __init__(
        self,
        voker_client: Optional[VokerClient] = None,
        base_anthropic_client: Optional[anthropic.Anthropic] = None,
        *,
        api_key: Optional[str] = None,
        auth_token: Optional[str] = None,
        base_url: Optional[str | httpx.URL] = None,
        timeout: Optional[float | anthropic.Timeout | anthropic.NotGiven] = not_given,
        max_retries: int = anthropic.DEFAULT_MAX_RETRIES,
        default_headers: Optional[Mapping[str, str]] = None,
        default_query: Optional[Mapping[str, object]] = None,
        # Configure a custom httpx client.
        # We provide a `DefaultHttpxClient` class that you can pass to retain the default values we use for `limits`, `timeout` & `follow_redirects`.
        # See the [httpx documentation](https://www.python-httpx.org/api/#client) for more details.
        http_client: Optional[httpx.Client] = None,
        # Enable or disable schema validation for data returned by the API.
        # When enabled an error APIResponseValidationError is raised
        # if the API responds with invalid data for the expected schema.
        #
        # This parameter may be removed or changed in the future.
        # If you rely on this feature, please open a GitHub issue
        # outlining your use-case to help us decide if it should be
        # part of our public interface in the future.
        _strict_response_validation: bool = False,
    ) -> None: ...

    def __init__(
        self,
        voker_client: Optional[VokerClient] = None,
        base_anthropic_client: Optional[anthropic.Anthropic] = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)

        if voker_client is None:
            voker_client = VokerClient()

        if base_anthropic_client is None:
            base_anthropic_client = self

        # keep for testing access:
        self.__base_anthropic_client = base_anthropic_client
        self.__voker_client = voker_client

        if (
            base_messages := getattr(base_anthropic_client, "messages", None)
        ) is not None:
            self.messages = Messages(voker_client, base_messages)

        if (base_beta := getattr(base_anthropic_client, "beta", None)) is not None:
            self.beta = Beta(voker_client, base_beta)


class Beta:
    """Send calls to the beta endpoints to Voker."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_beta: anthropic_resources.beta.beta.Beta,
    ) -> None:
        self.__base_beta = base_beta

        if (base_messages := getattr(base_beta, "messages", None)) is not None:
            self.messages = BetaMessages(voker_client, base_messages)

    def __getattr__(self, name: str) -> Any:
        """Pass through to the underlying Beta object."""
        return getattr(self.__base_beta, name)


class BetaMessages:
    """Send calls to the messages endpoints to Voker."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_messages: anthropic_resources.beta.messages.messages.Messages,
    ) -> None:
        self.__voker_client = voker_client
        self.__base_messages = base_messages

    def __getattr__(self, name: str) -> Any:
        """Pass through to the underlying Messages object."""
        return getattr(self.__base_messages, name)

    def create(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> anthropic_types_beta.BetaMessage:
        """Send a message creation request to Voker as an llm event.

        See `anthropic.Anthropic.messages.create`.
        """
        to_stream = kwargs.get("stream", False)

        if not to_stream:
            message = self.__base_messages.create(**kwargs)

            self.__voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "anthropic-beta-messages",
                    "inputs": _rm_not_given(kwargs),
                    "output": message.model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

            return message

        api_request = partial(self.__base_messages.create, **kwargs)

        def gen():
            aggregator = _AnthropicBetaMessagesStreamAggregator(
                headers=httpx.Headers().update(kwargs.get("extra_headers")),
                output_format=kwargs.get("response_format", anthropic.not_given),
            )

            for chunk in api_request():
                aggregator.add_event(chunk)
                yield chunk

            self.__voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "anthropic-beta-messages",
                    "inputs": _rm_not_given(kwargs),
                    "output": aggregator.get_result().model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

        return gen()


class Messages:
    """Send calls to the messages endpoints to Voker."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_messages: anthropic_resources.messages.messages.Messages,
    ) -> None:
        self.__voker_client = voker_client
        self.__base_messages = base_messages

    def __getattr__(self, name: str) -> Any:
        """Pass through to the underlying Messages object."""
        return getattr(self.__base_messages, name)

    def create(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> anthropic_types_message.Message:
        """Send a message creation request to Voker as an llm event.

        See `anthropic.Anthropic.messages.create`.
        """
        to_stream = kwargs.get("stream", False)

        if not to_stream:
            message = self.__base_messages.create(**kwargs)

            self.__voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "anthropic-messages",
                    "inputs": _rm_not_given(kwargs),
                    "output": message.model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

            return message

        api_request = self.__base_messages.create(**kwargs)

        def gen():
            aggregator = _AnthropicMessagesStreamAggregator()

            for chunk in api_request:
                aggregator.add_event(chunk)
                yield chunk

            self.__voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "anthropic-messages",
                    "inputs": _rm_not_given(kwargs),
                    # "output": message.model_dump(mode="json"),
                    "output": aggregator.get_result().model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

        return gen()


class AsyncAnthropic(anthropic.AsyncAnthropic):
    """Async Anthropic client wrapper that automatically creates llm events in Voker."""

    @overload
    def __init__(
        self,
        voker_client: Optional[VokerClient] = None,
        base_anthropic_client: Optional[anthropic.AsyncAnthropic] = None,
        *,
        api_key: Optional[str] = None,
        auth_token: Optional[str] = None,
        base_url: Optional[str | httpx.URL] = None,
        timeout: Optional[float | anthropic.Timeout | anthropic.NotGiven] = not_given,
        max_retries: int = anthropic.DEFAULT_MAX_RETRIES,
        default_headers: Optional[Mapping[str, str]] = None,
        default_query: Optional[Mapping[str, object]] = None,
        # Configure a custom httpx client.
        # We provide a `DefaultAsyncHttpxClient` class that you can pass to retain the default values we use for `limits`, `timeout` & `follow_redirects`.
        # See the [httpx documentation](https://www.python-httpx.org/api/#asyncclient) for more details.
        http_client: Optional[httpx.AsyncClient] = None,
        # Enable or disable schema validation for data returned by the API.
        # When enabled an error APIResponseValidationError is raised
        # if the API responds with invalid data for the expected schema.
        #
        # This parameter may be removed or changed in the future.
        # If you rely on this feature, please open a GitHub issue
        # outlining your use-case to help us decide if it should be
        # part of our public interface in the future.
        _strict_response_validation: bool = False,
    ) -> None: ...

    def __init__(
        self,
        voker_client: Optional[VokerClient] = None,
        base_anthropic_client: Optional[anthropic.AsyncAnthropic] = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)

        if voker_client is None:
            voker_client = VokerClient()

        if base_anthropic_client is None:
            base_anthropic_client = self

        # keep for testing access:
        self.__base_anthropic_client = base_anthropic_client
        self.__voker_client = voker_client

        if (
            base_messages := getattr(base_anthropic_client, "messages", None)
        ) is not None:
            self.messages = AsyncMessages(voker_client, base_messages)

        if (base_beta := getattr(base_anthropic_client, "beta", None)) is not None:
            self.beta = AsyncBeta(voker_client, base_beta)


class AsyncBeta:
    """Async Beta client wrapper that automatically creates llm events in Voker."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_beta_client: Optional[anthropic.AsyncBeta] = None,
    ) -> None:
        if (base_messages := getattr(base_beta_client, "messages", None)) is not None:
            self.messages = AsyncBetaMessages(voker_client, base_messages)

    def __getattr__(self, name: str) -> Any:
        """Pass through to the underlying AsyncBeta object."""
        return getattr(self.__base_beta_client, name)


class AsyncBetaMessages:
    """Async Beta Messages client wrapper that automatically creates llm events in Voker."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_messages: anthropic_resources.beta.messages.messages.AsyncMessages,
    ) -> None:
        self.__voker_client = voker_client
        self.__base_messages = base_messages

    def __getattr__(self, name: str) -> Any:
        """Pass through to the underlying AsyncMessages object."""
        return getattr(self.__base_messages, name)

    async def create(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> anthropic_types_message.Message:
        """Send an async message creation request to Voker as an llm event.

        See `anthropic.Anthropic.messages.create`.
        """
        to_stream = kwargs.get("stream", False)

        if not to_stream:
            message = await self.__base_messages.create(**kwargs)

            self.__voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "anthropic-beta-messages",
                    "inputs": _rm_not_given(kwargs),
                    "output": message.model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

            return message

        api_request = self.__base_messages.create(**kwargs)

        async def gen():
            aggregator = _AnthropicBetaMessagesStreamAggregator(
                headers=httpx.Headers().update(kwargs.get("extra_headers")),
                output_format=kwargs.get("response_format", anthropic.not_given),
            )

            async for chunk in await api_request:
                aggregator.add_event(chunk)
                yield chunk

            self.__voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "anthropic-beta-messages",
                    "inputs": _rm_not_given(kwargs),
                    "output": aggregator.get_result().model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

        return gen()


class AsyncMessages:
    """Send async calls to the messages endpoints to Voker as llm events."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_messages: anthropic_resources.messages.messages.AsyncMessages,
    ) -> None:
        self.__voker_client = voker_client
        self.__base_messages = base_messages

    def __getattr__(self, name: str) -> Any:
        """Pass through to the underlying AsyncMessages object."""
        return getattr(self.__base_messages, name)

    async def create(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> anthropic_types_message.Message:
        """Send an async message creation request to Voker as an llm event.

        See `anthropic.Anthropic.messages.create`.
        """
        to_stream = kwargs.get("stream", False)

        if not to_stream:
            message = await self.__base_messages.create(**kwargs)

            self.__voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "anthropic-messages",
                    "inputs": _rm_not_given(kwargs),
                    "output": message.model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

            return message

        api_request = self.__base_messages.create(**kwargs)

        async def gen():
            aggregator = _AnthropicMessagesStreamAggregator()

            async for chunk in await api_request:
                aggregator.add_event(chunk)
                yield chunk

            self.__voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "anthropic-messages",
                    "inputs": _rm_not_given(kwargs),
                    "output": aggregator.get_result().model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

        return gen()


class _AnthropicMessagesStreamAggregator:

    def __init__(self) -> None:
        self.current: anthropic_types.Message | None = None

    def add_event(self, chunk: anthropic_types.RawMessageStreamEvent) -> None:
        from anthropic.lib.streaming._messages import accumulate_event

        self.current = accumulate_event(event=chunk, current_snapshot=self.current)

    def get_result(self) -> anthropic_types.Message:
        assert self.current is not None
        return self.current


class _AnthropicBetaMessagesStreamAggregator:

    def __init__(
        self,
        headers: httpx.Headers,
        output_format: anthropic_types_beta.ResponseFormatT | anthropic.NotGiven,
    ) -> None:
        self.current: anthropic_types_beta.Message | None = None
        self.headers = headers
        self.output_format = output_format

    def add_event(self, chunk: anthropic_types_beta.RawMessageStreamEvent) -> None:
        from anthropic.lib.streaming._beta_messages import accumulate_event

        self.current = accumulate_event(
            event=chunk,
            current_snapshot=self.current,
            request_headers=self.headers,
            output_format=self.output_format,
        )

    def get_result(self) -> anthropic_types_beta.Message:
        assert self.current is not None
        return self.current
