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

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

from __future__ import annotations

from collections.abc import AsyncIterator, Iterator
from typing import TYPE_CHECKING, Any, Optional, overload

from voker.client import VokerClient

if TYPE_CHECKING:
    from google.genai import _interactions as genai_interactions
    from google.genai import models as genai_models
    from google.genai import types as genai_types
    from google.genai._interactions import resources as genai_interactions_resources
    from google.genai._interactions import types as genai_interactions_types

try:
    import google
    from google import genai
except ImportError:
    raise ModuleNotFoundError


class Client(genai.Client):
    """Gemini wrapper that automatically creates llm events in Voker."""

    @overload
    def __init__(
        self,
        voker_client: Optional[VokerClient] = None,
        base_genai_client: Optional[genai.Client] = None,
        *,
        vertexai: Optional[bool] = None,
        api_key: Optional[str] = None,
        credentials: Optional[google.auth.credentials.Credentials] = None,
        project: Optional[str] = None,
        location: Optional[str] = None,
        debug_config: Optional[genai.client.DebugConfig] = None,
        http_options: Optional[
            genai_types.HttpOptions | genai_types.HttpOptionsDict
        ] = None,
    ) -> None: ...

    def __init__(
        self,
        voker_client: Optional[VokerClient] = None,
        base_genai_client: Optional[genai.Client] = None,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        if voker_client is not None:
            self.__voker_client = voker_client
        else:
            self.__voker_client = VokerClient()

        if base_genai_client is None:
            base_genai_client = self
            self.__base_genai_client = base_genai_client

        if (base_models := getattr(base_genai_client, "_models", None)) is not None:
            self._models = Models(self.__voker_client, base_models)

        if (
            base_interactions := getattr(base_genai_client, "interactions", None)
        ) is not None:
            self._nextgen_client_instance

            self._nextgen_client.interactions = Interactions(
                self.__voker_client,
                base_interactions,
            )

    @property
    def aio(self) -> AsyncClient:
        """Pass through to the underlying AsyncClient instance."""
        return AsyncClient(
            self.__voker_client,
            self.__base_genai_client._aio,
        )


class Models:
    """Send calls to the models endpoints to Voker."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_models: genai_models.Models,
    ) -> None:
        self.__voker_client = voker_client
        self.__base_models = base_models

    def __getattr__(self, name: str) -> Any:
        """Pass through to the underlying Models instance."""
        return getattr(self.__base_models, name)

    def generate_content(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> genai_types.GenerateContentResponse:
        """Send a content generation request to Voker as an llm event.

        See `google.genai.Client.models.generate_content`.
        """
        response = self.__base_models.generate_content(**kwargs)

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

        return response

    def generate_content_stream(
        self,
        **kwargs,
    ) -> Iterator[genai_types.GenerateContentResponse]:
        """Send a streaming content generation request to Voker as an llm event.

        See `google.genai.Client.models.generate_content_stream`.
        """
        for chunk in self.__base_models.generate_content_stream(**kwargs):
            yield chunk

        # TODO: accumulate


class Interactions:
    """Send calls to the interactions endpoints to Voker."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_interactions: genai_interactions_resources.InteractionsResource,
    ) -> None:
        self.__voker_client = voker_client
        self.__base_interactions = base_interactions

    def __getattr__(self, name: str) -> Any:
        """Pass through to the underlying Interactions instance."""
        return getattr(self.__base_interactions, name)

    def create(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> (
        genai_interactions_types.Interaction
        | genai_interactions.Stream[genai_interactions_types.InteractionSSEEvent]
    ):
        """Send a create interaction request to Voker as an llm event.

        See: `google.genai.Client.interactions.create`
        """
        from google.genai._interactions.types import Interaction

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

        interaction = self.__base_interactions.create(**kwargs)

        if not stream:
            self.__voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "gemini-interactions",
                    "inputs": kwargs,
                    "output": interaction,
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

            return interaction

        # TODO: stream and accumulate...

        return interaction


class AsyncClient:
    """Async Gemini client wrapper that automatically creates llm events in Voker."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_genai_async_client: genai.client.AsyncClient,
    ) -> None:
        self.__voker_client = voker_client
        self.__base_genai_async_client = base_genai_async_client

        if (
            base_models := getattr(self.__base_genai_async_client, "models", None)
        ) is not None:
            self._models = AsyncModels(voker_client, base_models)

        if (
            base_interactions := getattr(base_genai_async_client, "interactions", None)
        ) is not None:
            self._nextgen_client.interactions = AsyncInteractions(
                self.__voker_client,
                base_interactions,
            )

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


class AsyncModels:
    """Send async calls to the models endpoints to Voker as llm events."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_models: genai_models.AsyncModels,
    ) -> None:
        self.__voker_client = voker_client
        self.__base_models = base_models

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

    async def generate_content(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> genai_types.GenerateContentResponse:
        """Send an async content generation request to Voker as an llm event.

        See: `google.genai.client.AsyncClient.models.generate_content`
        """
        response = await self.__base_models.generate_content(**kwargs)

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

        return response

    async def generate_content_stream(
        self,
        **kwargs,
    ) -> AsyncIterator[genai_types.GenerateContentResponse]:
        """Send an async streaming content generation request to Voker as an llm event.

        See: `google.genai.client.AsyncClient.models.generate_content_stream`
        """
        stream = await self.__base_models.generate_content_stream(**kwargs)

        async for chunk in stream:
            yield chunk

        # TODO: accumulate


class AsyncInteractions:
    """Send async calls to the interactions endpoints to Voker."""

    def __init__(
        self,
        voker_client: VokerClient,
        base_interactions: genai_interactions_resources.AsyncInteractionsResource,
    ) -> None:
        self.__voker_client = voker_client
        self.__base_interactions = base_interactions

    def __getattr__(self, name: str) -> Any:
        """Pass through to the underlying Interactions instance."""
        return getattr(self.__base_interactions, name)

    async def create(
        self,
        agent: Optional[str] = None,
        agent_version: Optional[str] = None,
        person: Optional[str] = None,
        session: Optional[str] = None,
        **kwargs,
    ) -> (
        genai_interactions_types.Interaction
        | genai_interactions.AsyncStream[genai_interactions_types.InteractionSSEEvent]
    ):
        """Send a, async create interaction request to Voker as an llm event.

        See: `google.genai.client.AsyncClient.interactions.create`
        """
        from google.genai._interactions.types import Interaction

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

        interaction = await self.__base_interactions.create(**kwargs)

        if not stream:
            self.__voker_client.events.create(
                event_name="llm",
                properties={
                    "api": "gemini-interactions",
                    "inputs": kwargs,
                    "output": interaction,
                },
                agent=agent,
                agent_version=agent_version,
                person=person,
                session=session,
            )

            return interaction

        # TODO: stream and accumulate...

        return interaction
