# 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 OpenAI.

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

from __future__ import annotations

from collections.abc import Mapping
from functools import partial
from typing import (
    TYPE_CHECKING,
    Any,
    Awaitable,
    Callable,
    Generator,
    Optional,
    overload,
)

from voker.client import VokerClient

if TYPE_CHECKING:
    import httpx
    from openai import _streaming as openai_streaming
    from openai import resources as openai_resources
    from openai.resources import responses as openai_responses
    from openai.types import chat as openai_types_chat
    from openai.types import responses as openai_types_responses


try:
    import openai
except ImportError:
    raise ModuleNotFoundError


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


class OpenAI(openai.OpenAI):
    """OpenAI wrapper that automatically creates llm events in Voker."""

    @overload
    def __init__(
        self,
        voker_client: Optional[VokerClient] = None,
        base_openai_client: Optional[openai.OpenAI] = None,
        *,
        api_key: Optional[str | Callable[[], str]] = None,
        organization: Optional[str] = None,
        project: Optional[str] = None,
        webhook_secret: Optional[str] = None,
        base_url: Optional[str | httpx.URL] = None,
        websocket_base_url: Optional[str | httpx.URL] = None,
        timeout: Optional[float | openai.Timeout | openai.NotGiven] = openai.not_given,
        max_retries: int = openai.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,
        **kwargs,
    ) -> None: ...

    def __init__(
        self,
        voker_client: Optional[VokerClient] = None,
        base_openai_client: Optional[openai.OpenAI] = None,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        if voker_client is None:
            voker_client = VokerClient()

        if base_openai_client is None:
            base_openai_client = self

        # keep for testing access:
        self.__base_openai_client = base_openai_client
        self.__voker_client = voker_client

        if (base_chat := getattr(base_openai_client, "chat", None)) is not None:
            self.chat = Chat(voker_client, base_chat)

        if (
            base_responses := getattr(base_openai_client, "responses", None)
        ) is not None:
            self.responses = Responses(voker_client, base_responses)


class Chat:
    """Send calls to the chat endpoints to Voker."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_chat: openai_resources.chat.chat.Chat,
    ) -> None:
        self.__base_chat = base_chat

        if (base_completions := getattr(base_chat, "completions", None)) is not None:
            self.completions = ChatCompletions(voker_client, base_completions)

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


class ChatCompletions:
    """Send calls to the chat completion endpoints to Voker as llm events."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_chat_completions: openai_resources.chat.completions.completions.Completions,
    ) -> None:
        self._voker_client = voker_client
        self.__base_chat_completions = base_chat_completions

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

    # TODO: how do i make this show up in intellisense properly?
    def create(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> (
        openai_types_chat.chat_completion.ChatCompletion
        | openai_streaming.Stream[
            openai_types_chat.chat_completion_chunk.ChatCompletionChunk
        ]
    ):
        """Send a chat completion request to Voker as an llm event.

        See `openai.OpenAI.chat.completions.create`
        """
        # TODO: add our other inputs in like person_id, ....
        # TODO: start time, end time, etc... # retries...
        to_stream = kwargs.get("stream", False)

        if not to_stream:
            response = self.__base_chat_completions.create(**kwargs)

            self._voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "openai-chat-completions",
                    "inputs": _rm_not_given(kwargs),
                    "output": response.model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

            return response

        def gen() -> (
            Generator[
                openai_types_chat.chat_completion_chunk.ChatCompletionChunk, None, None
            ]
        ):
            api_request: partial[
                openai_streaming.Stream[
                    openai_types_chat.chat_completion_chunk.ChatCompletionChunk
                ]
            ] = partial(self.__base_chat_completions.create, **kwargs)

            manager = openai.lib.streaming.chat.ChatCompletionStreamManager(
                api_request,
                response_format=kwargs.get("response_format", openai._types.omit),
                input_tools=kwargs.get("tools", openai._types.omit),
            )

            with manager as stream:
                for chunk in stream:
                    c = getattr(chunk, "chunk", None)

                    if c is None or len(c.choices) == 0:
                        continue

                    yield c

                response = stream.get_final_completion()

            self._voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "openai-chat-completions",
                    "inputs": _rm_not_given(kwargs),
                    "output": response.model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

        return gen()  # type: ignore


class Responses:
    """Send calls to the responses API to Voker."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_responses: openai_responses.responses.Responses,
    ) -> None:
        self._voker_client = voker_client
        self.__base_responses = base_responses

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

    def create(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> (
        openai_types_responses.response.Response
        | openai_streaming.Stream[
            openai_types_responses.response_stream_event.ResponseStreamEvent
        ]
    ):
        """Send a response request to Voker as an llm event.

        See `openai.OpenAI.responses.create`
        """
        # TODO: add our other inputs in like person_id, ....
        # TODO: start time, end time, etc... # retries...
        # TODO: instructions, conversation, previous_response_id, prompt, etc...
        to_stream = kwargs.get("stream", False)

        if not to_stream:
            response = self.__base_responses.create(**kwargs)

            self._voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "openai-responses",
                    "inputs": _rm_not_given(kwargs),
                    "output": response.model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

            return response

        from openai.lib.streaming.responses._responses import ResponseStreamManager

        def gen() -> Generator[
            openai_types_responses.response_stream_event.ResponseStreamEvent,
            None,
            None,
        ]:
            api_request: partial[
                openai_streaming.Stream[
                    openai_types_responses.response_stream_event.ResponseStreamEvent
                ]
            ] = partial(self.__base_responses.create, **kwargs)

            manager = ResponseStreamManager(
                api_request,
                text_format=kwargs.get("text_format", openai._types.omit),
                input_tools=kwargs.get("tools", openai._types.omit),
                starting_after=kwargs.get("starting_after", None),
            )

            with manager as stream:
                for chunk in stream:
                    yield chunk

                response = stream.get_final_response()

            self._voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "openai-responses",
                    "inputs": _rm_not_given(kwargs),
                    "output": response.model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

        return gen()  # type: ignore


class AsyncOpenAI(openai.AsyncOpenAI):
    """OpenAI wrapper that automatically creates llm events in Voker."""

    @overload
    def __init__(
        self,
        voker_client: Optional[VokerClient] = None,
        base_openai_client: Optional[openai.AsyncOpenAI] = None,
        *,
        api_key: Optional[str | Callable[[], Awaitable[str]]] = None,
        organization: Optional[str] = None,
        project: Optional[str] = None,
        webhook_secret: Optional[str] = None,
        base_url: Optional[str | httpx.URL] = None,
        websocket_base_url: Optional[str | httpx.URL] = None,
        timeout: Optional[float | openai.Timeout | openai.NotGiven] = openai.not_given,
        max_retries: int = openai.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.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,
        **kwargs,
    ) -> None: ...

    def __init__(
        self,
        voker_client: Optional[VokerClient] = None,
        base_openai_client: Optional[openai.AsyncOpenAI] = None,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        if voker_client is None:
            voker_client = VokerClient()

        if base_openai_client is None:
            base_openai_client = self

        # keep for testing access:
        self.__base_openai_client = base_openai_client
        self.__voker_client = voker_client

        if (base_chat := getattr(base_openai_client, "chat", None)) is not None:
            self.chat = AsyncChat(voker_client, base_chat)

        if (
            base_responses := getattr(base_openai_client, "responses", None)
        ) is not None:
            self.responses = AsyncResponses(voker_client, base_responses)


class AsyncChat:
    """Send calls to the chat endpoints to Voker as llm events."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_chat: openai_resources.chat.chat.AsyncChat,
    ) -> None:
        self.__base_chat = base_chat

        if (base_completions := getattr(base_chat, "completions", None)) is not None:
            self.completions = AsyncChatCompletions(voker_client, base_completions)

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


class AsyncChatCompletions:
    """Send calls to the chat completion endpoints to Voker as llm events."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_chat_completions: openai_resources.chat.completions.completions.AsyncCompletions,
    ) -> None:
        self._voker_client = voker_client
        self.__base_chat_completions = base_chat_completions

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

    async def create(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> (
        openai_types_chat.chat_completion.ChatCompletion
        | openai_streaming.AsyncStream[
            openai_types_chat.chat_completion_chunk.ChatCompletionChunk
        ]
    ):
        """Send a chat completion request to Voker as an llm event.

        See `openai.OpenAI.chat.completions.create`
        """
        # TODO: add our other inputs in like person_id, ....
        # TODO: start time, end time, etc... # retries...

        to_stream = kwargs.get("stream", False)

        if not to_stream:
            response = await self.__base_chat_completions.create(**kwargs)

            self._voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "openai-chat-completions",
                    "inputs": _rm_not_given(kwargs),
                    "output": response.model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

            return response

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

        manager = openai.lib.streaming.chat.AsyncChatCompletionStreamManager(
            api_request,
            response_format=kwargs.get("response_format", openai._types.omit),
            input_tools=kwargs.get("tools", openai._types.omit),
        )

        async def gen():
            async with manager as stream:
                async for chunk in stream:
                    c = getattr(chunk, "chunk", None)

                    if c is None or len(c.choices) == 0:
                        continue

                    yield c

                response = await stream.get_final_completion()

                self._voker_client.events.create(
                    event_name="llm",
                    properties={
                        "api": "openai-chat-completions",
                        "inputs": _rm_not_given(kwargs),
                        "output": response.model_dump(mode="json"),
                    },
                    agent=agent,
                    agent_version=agent_version,
                    person=person,
                    session=session,
                )

        return gen()


class AsyncResponses:
    """Send calls to the responses API to Voker as llm events."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_responses: openai_responses.responses.AsyncResponses,
    ) -> None:
        self._voker_client = voker_client
        self.__base_responses = base_responses

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

    async def create(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> (
        openai_types_responses.response.Response
        | openai_streaming.AsyncStream[
            openai_types_responses.response_stream_event.ResponseStreamEvent
        ]
    ):
        """Send a response request to Voker as an llm event.

        See `openai.OpenAI.responses.create`
        """
        # TODO: add our other inputs in like person_id, ....
        # TODO: start time, end time, etc... # retries...
        # TODO: instructions, conversation, previous_response_id, prompt, etc...
        to_stream = kwargs.get("stream", False)

        if not to_stream:
            response = await self.__base_responses.create(**kwargs)

            self._voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "openai-responses",
                    "inputs": _rm_not_given(kwargs),
                    "output": response.model_dump(mode="json"),
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

            return response

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

        from openai.lib.streaming.responses._responses import AsyncResponseStreamManager

        manager = AsyncResponseStreamManager(
            api_request,
            text_format=kwargs.get("text_format", openai._types.omit),
            input_tools=kwargs.get("tools", openai._types.omit),
            starting_after=kwargs.get("starting_after", None),
        )

        async def gen():
            async with manager as stream:
                async for chunk in stream:
                    yield chunk

                response = await stream.get_final_response()

                self._voker_client.events.create(
                    event_name="llm",
                    properties={
                        "api": "openai-responses",
                        "inputs": _rm_not_given(kwargs),
                        "output": response.model_dump(mode="json"),
                    },
                    agent=agent,
                    agent_version=agent_version,
                    person=person,
                    session=session,
                )

        return gen()
