from __future__ import annotations

from functools import lru_cache
from typing import Any, Mapping

from google import genai
from google.genai import types


def _build_generation_config(generation_config: Any) -> Any:
    if generation_config is None:
        return None
    if isinstance(generation_config, types.GenerateContentConfig):
        return generation_config
    if isinstance(generation_config, Mapping):
        try:
            return types.GenerateContentConfig(**dict(generation_config))
        except TypeError:
            return dict(generation_config)
    return generation_config


class GeminiModelAdapter:
    """Compat adapter that keeps the old generate_content(...) shape."""

    def __init__(self, api_key: str, model_name: str) -> None:
        self._client = genai.Client(api_key=api_key)
        self._model_name = model_name

    def generate_content(self, contents: Any, generation_config: Any = None):
        return self._client.models.generate_content(
            model=self._model_name,
            contents=contents,
            config=_build_generation_config(generation_config),
        )


@lru_cache(maxsize=16)
def get_gemini_model(api_key: str, model_name: str) -> GeminiModelAdapter:
    return GeminiModelAdapter(api_key=api_key, model_name=model_name)