import json
import logging
from copy import deepcopy

from openai import OpenAI

from rhodecode.apps.ai_agents.ai_settings import AISettings

from rhodecode.apps.ai_agents.models.base import (
    PingRequest,
    CodeReviewRequest,
    CustomFunctions,
    Response,
    AIService,
)


class GeminiService(AIService):
    """
    uses google compatibility option with OpenAI library: https://ai.google.dev/gemini-api/docs/openai
    """

    def __init__(self, model_settings: AISettings):
        super().__init__(model_settings)
        self._client = OpenAI(
            api_key=self.model_settings.api_key,
            base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
        )
        self.log = logging.getLogger(GeminiService.__name__)

    def list_model_api_and_display_names(self) -> list[tuple[str, str]]:
        res = super().list_model_api_and_display_names()
        return self._remove_models_from_api_name(res)

    def _remove_models_from_api_name(self, api_model_names: list[tuple[str, str]]) -> list[tuple[str, str]]:
        res = []
        for api_name, display_name in api_model_names:
            res.append((api_name.replace("models/", ""), display_name))
        return res

    def _list_non_text_models(self) -> list[str]:
        return [
            "audio",
            "lyria",
            "aqa",
            "veo",
            "learnlm",
            "imagen",
            "image",
            "embedding",
            "computer-use",
            "robotics",
        ]

    def _get_response(self, request: PingRequest | CodeReviewRequest):
        _input = self._get_response_input(request)

        function = self._adapt_function(self._get_function(request.function_name))

        return self._client.chat.completions.create(
            model=self.get_api_model_name(),
            messages=_input,
            tools=[function],
            tool_choice={"type": "function", "function": {"name": request.function_name}},
            parallel_tool_calls=False,
        )

    def _transform(self, resp):
        for item in resp.choices:
            self.log.debug("response: %s", item)
            if item.finish_reason == "tool_calls":
                for tool in item.message.tool_calls:
                    if tool.function.name in [CustomFunctions.Code_review, CustomFunctions.Ping]:
                        return Response(model=resp.model, message=json.loads(tool.function.arguments))

        raise ValueError("No function call found in response")

    def _adapt_function(self, function_dict: dict):
        original_copy = deepcopy(function_dict)

        del original_copy["type"]
        del original_copy["strict"]

        return {
            "type": "function",
            "function": original_copy,
        }
