# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0

import json
from collections.abc import AsyncGenerator, Callable
from contextlib import asynccontextmanager
from logging import Logger, getLogger
from typing import TYPE_CHECKING, Any

from ......doc_utils import export_module
from ......import_utils import optional_import_block, require_optional_import
from ......llm_config import LLMConfig
from ...realtime_events import AudioDelta, FunctionCall, RealtimeEvent, SessionCreated
from ..realtime_client import RealtimeClientBase, Role, register_realtime_client

with optional_import_block():
    from websockets.asyncio.client import connect


if TYPE_CHECKING:
    from websockets.asyncio.client import ClientConnection

    from ..realtime_client import RealtimeClientProtocol

__all__ = ["GeminiRealtimeClient"]

global_logger = getLogger(__name__)


HOST = "generativelanguage.googleapis.com"
API_VERSION = "v1alpha"


@register_realtime_client()
@require_optional_import("websockets", "gemini", except_for=["get_factory", "__init__"])
@export_module("autogen.agentchat.realtime.experimental.clients")
class GeminiRealtimeClient(RealtimeClientBase):
    """(Experimental) Client for Gemini Realtime API."""

    def __init__(
        self,
        *,
        llm_config: LLMConfig | dict[str, Any],
        logger: Logger | None = None,
    ) -> None:
        """(Experimental) Client for Gemini Realtime API.

        Args:
            llm_config: The config for the client.
            logger: The logger for the client.
        """
        super().__init__()
        self._llm_config = llm_config
        self._logger = logger

        self._connection: ClientConnection | None = None
        config = llm_config["config_list"][0]

        self._model: str = config["model"]
        self._voice = config.get("voice", "charon")
        self._temperature: float = config.get("temperature", 0.8)  # type: ignore[union-attr]

        self._response_modality = "AUDIO"

        self._api_key = config.get("api_key", None)
        # todo: add test with base_url just to make sure it works
        self._base_url: str = config.get(
            "base_url",
            f"wss://{HOST}/ws/google.ai.generativelanguage.{API_VERSION}.GenerativeService.BidiGenerateContent?key={self._api_key}",
        )
        self._final_config: dict[str, Any] = {}
        self._pending_session_updates: dict[str, Any] = {}
        self._is_reading_events = False

    @property
    def logger(self) -> Logger:
        """Get the logger for the Gemini Realtime API."""
        return self._logger or global_logger

    @property
    def connection(self) -> "ClientConnection":
        """Get the Gemini WebSocket connection."""
        if self._connection is None:
            raise RuntimeError("Gemini WebSocket is not initialized")
        return self._connection

    async def send_function_result(self, call_id: str, result: str) -> None:
        """Send the result of a function call to the Gemini Realtime API.

        Args:
            call_id (str): The ID of the function call.
            result (str): The result of the function call.
        """
        msg = {
            "tool_response": {"function_responses": [{"id": call_id, "response": {"result": {"string_value": result}}}]}
        }
        if self._is_reading_events:
            await self.connection.send(json.dumps(msg))

    async def send_text(self, *, role: Role, text: str, turn_complete: bool = True) -> None:
        """Send a text message to the Gemini Realtime API.

        Args:
            role: The role of the message.
            text: The text of the message.
            turn_complete: A flag indicating if the turn is complete.
        """
        msg = {
            "client_content": {
                "turn_complete": turn_complete,
                "turns": [{"role": role, "parts": [{"text": text}]}],
            }
        }
        if self._is_reading_events:
            await self.connection.send(json.dumps(msg))

    async def send_audio(self, audio: str) -> None:
        """Send audio to the Gemini Realtime API.

        Args:
            audio (str): The audio to send.
        """
        msg = {
            "realtime_input": {
                "media_chunks": [
                    {
                        "data": audio,
                        "mime_type": "audio/pcm",
                    }
                ]
            }
        }
        await self.queue_input_audio_buffer_delta(audio)
        if self._is_reading_events:
            await self.connection.send(json.dumps(msg))

    async def truncate_audio(self, audio_end_ms: int, content_index: int, item_id: str) -> None:
        self.logger.info("This is not natively supported by Gemini Realtime API.")
        pass

    async def _initialize_session(self) -> None:
        """Initialize the session with the Gemini Realtime API."""
        session_config = {
            "setup": {
                "system_instruction": {
                    "role": "system",
                    "parts": [{"text": self._pending_session_updates.get("instructions", "")}],
                },
                "model": f"models/{self._model}",
                "tools": [
                    {
                        "function_declarations": [
                            {
                                "name": tool_schema["name"],
                                "description": tool_schema["description"],
                                "parameters": tool_schema["parameters"],
                            }
                            for tool_schema in self._pending_session_updates.get("tools", [])
                        ]
                    },
                ],
                "generation_config": {
                    "response_modalities": [self._response_modality],
                    "speech_config": {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": self._voice}}},
                    "temperature": self._temperature,
                },
            }
        }

        self.logger.info(f"Sending session update: {session_config}")
        await self.connection.send(json.dumps(session_config))

    async def session_update(self, session_options: dict[str, Any]) -> None:
        """Record session updates to be applied when the connection is established.

        Args:
            session_options (dict[str, Any]): The session options to update.
        """
        if self._is_reading_events:
            self.logger.warning("Is reading events. Session update will be ignored.")
        else:
            self._pending_session_updates.update(session_options)

    @asynccontextmanager
    async def connect(self) -> AsyncGenerator[None, None]:
        """Connect to the Gemini Realtime API."""
        try:
            async with connect(
                self._base_url, additional_headers={"Content-Type": "application/json"}
            ) as self._connection:
                yield
        finally:
            self._connection = None

    async def read_events(self) -> AsyncGenerator[RealtimeEvent, None]:
        """Read Events from the Gemini Realtime Client"""
        if self._connection is None:
            raise RuntimeError("Client is not connected, call connect() first.")
        await self._initialize_session()

        self._is_reading_events = True

        async for event in self._read_events():
            yield event

    async def _read_from_connection(self) -> AsyncGenerator[RealtimeEvent, None]:
        """Read messages from the Gemini Realtime connection."""
        async for raw_message in self.connection:
            message = raw_message.decode("ascii") if isinstance(raw_message, bytes) else raw_message
            events = self._parse_message(json.loads(message))
            for event in events:
                yield event

    def _parse_message(self, response: dict[str, Any]) -> list[RealtimeEvent]:
        """Parse a message from the Gemini Realtime API.

        Args:
            response (dict[str, Any]): The response to parse.

        Returns:
            list[RealtimeEvent]: The parsed events.
        """
        if "serverContent" in response and "modelTurn" in response["serverContent"]:
            try:
                b64data = response["serverContent"]["modelTurn"]["parts"][0]["inlineData"].pop("data")
                return [
                    AudioDelta(
                        delta=b64data,
                        item_id=None,
                        raw_message=response,
                    )
                ]
            except KeyError:
                return []
        elif "toolCall" in response:
            return [
                FunctionCall(
                    raw_message=response,
                    call_id=call["id"],
                    name=call["name"],
                    arguments=call["args"],
                )
                for call in response["toolCall"]["functionCalls"]
            ]
        elif "setupComplete" in response:
            return [
                SessionCreated(raw_message=response),
            ]
        else:
            return [RealtimeEvent(raw_message=response)]

    @classmethod
    def get_factory(
        cls, llm_config: LLMConfig | dict[str, Any], logger: Logger, **kwargs: Any
    ) -> Callable[[], "RealtimeClientProtocol"] | None:
        """Create a Realtime API client.

        Args:
            llm_config: The LLM config for the client.
            logger: The logger for the client.
            **kwargs: Additional arguments.

        Returns:
            RealtimeClientProtocol: The Realtime API client is returned if the model matches the pattern
        """
        if llm_config["config_list"][0].get("api_type") == "google" and list(kwargs.keys()) == []:
            return lambda: GeminiRealtimeClient(llm_config=llm_config, logger=logger, **kwargs)
        return None


# needed for mypy to check if GeminiRealtimeClient implements RealtimeClientProtocol
if TYPE_CHECKING:
    _client: RealtimeClientProtocol = GeminiRealtimeClient(llm_config={})
