# Copyright 2024 Klira AI
#
# 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.

"""Google Gemini adapter.

Patches google.generativeai GenerativeModel.generate_content to create
klira.llm.gemini spans with standard gen_ai attributes.
"""

from __future__ import annotations

import logging
from typing import Any

from klira.sdk.adapters.base_llm import BaseLLMAdapter

logger = logging.getLogger(__name__)

PROVIDER = "gemini"


class GeminiAdapter(BaseLLMAdapter):
    """Adapter for Google Gemini API."""

    def patch(self, client_class: type) -> None:
        """Patch the Gemini GenerativeModel's generate_content method."""
        if self._is_already_patched(client_class):
            logger.debug("Gemini GenerativeModel already patched, skipping")
            return

        adapter = self

        original_generate = client_class.generate_content  # type: ignore[attr-defined]

        def patched_generate(self_model: Any, *args: Any, **kwargs: Any) -> Any:
            model_name = getattr(self_model, "model_name", "unknown")

            with adapter.create_llm_span(PROVIDER) as span:
                adapter.set_request_attributes(span, model_name)

                response = original_generate(self_model, *args, **kwargs)

                # Extract token counts from usage metadata
                usage = getattr(response, "usage_metadata", None)
                input_tokens = (
                    getattr(usage, "prompt_token_count", None) if usage else None
                )
                output_tokens = (
                    getattr(usage, "candidates_token_count", None) if usage else None
                )

                # Extract finish reason
                candidates = getattr(response, "candidates", [])
                finish_reasons = []
                for candidate in candidates:
                    reason = getattr(candidate, "finish_reason", None)
                    if reason is not None:
                        finish_reasons.append(str(reason))

                adapter.set_response_attributes(
                    span,
                    model=model_name,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    finish_reasons=finish_reasons,
                )

                # Capture output
                output_text = getattr(response, "text", None) or ""
                if output_text:
                    if len(output_text) > adapter.OUTPUT_TRUNCATION_LIMIT:
                        output_text = output_text[: adapter.OUTPUT_TRUNCATION_LIMIT]
                    span.set_attribute("klira.output", output_text)

                return response

        self._mark_function(patched_generate)
        client_class.generate_content = patched_generate  # type: ignore[attr-defined]
        self._mark_as_patched(client_class)
        self.verify_patch(client_class)

    def verify_patch(self, client_class: type) -> bool:
        """Verify the patch took effect by inspecting the actual method."""
        method = getattr(client_class, "generate_content", None)
        if method is None or not getattr(method, "_klira_patched", False):
            raise RuntimeError(
                f"Gemini adapter: patch not applied to {client_class}"
            )
        return True
