import asyncio
from typing import Optional, Type
from google import genai
from pydantic import BaseModel
from google.genai.types import GenerateContentConfig
import bilbystats as bs


async def gemini_api_async(
    content: str,
    instructions: str = "",
    model_name: str = "gemini-2.0-flash",
    response_format: Optional[Type[BaseModel]] = None,
    temperature: float = 1.0,
    max_tokens: int = 1000
) -> str:
    """
    Async generate content using the Gemini API.

    This function sends an async request to the Gemini API to generate content based on
    the provided input. If the default model name "gemini" is used, it is replaced
    with "gemini-2.0-flash" to ensure compatibility. The function returns the
    generated text content.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    - content : str
        The input prompt or content to be sent to the Gemini API for generation.

    - model_name : str, optional (default="gemini-2.0-flash")
        The name of the Gemini model to use for content generation. If "gemini"
        is specified, it defaults to "gemini-2.0-flash". Other options are e.g.
        gemini-2.5-flash-preview-04-17 and gemini-2.0-flash-lite and
        gemini-2.5-pro-preview-05-06. See
        https://ai.google.dev/gemini-api/docs/models for a list of the different
        possible models.

    ---------------------------------------------------------------------------
    OUTPUT:
    - output : str
        The generated text returned by the Gemini API.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport and Stephen Enwright-Ward
    ---------------------------------------------------------------------------
    """
    if model_name == "gemini":
        model_name = "gemini-2.0-flash"

    # Note: The genai client doesn't have native async support
    # We'll use asyncio.run_in_executor to make it async
    if not response_format:
        def _generate_sync():
            client = genai.Client(api_key=bs.read_api_key("gemini"))
            response = client.models.generate_content(
                model=model_name,
                contents=instructions + content,
                config=GenerateContentConfig(
                    temperature=temperature,
                    max_output_tokens=max_tokens,
                ),
            )
            return response.text
    else:
        def _generate_sync():
            client = genai.Client(
                api_key=bs.read_api_key("gemini"),
                )

            response = client.models.generate_content(
                model=model_name,
                contents=content,
                config={
                    "response_mime_type": "application/json",
                    "system_instruction": instructions,
                    "response_schema": response_format,
                    "temperature": temperature,
                    "max_output_tokens": max_tokens,
                    },
                )

            # return response.text
            return response.parsed
    

    loop = asyncio.get_event_loop()
    output = await loop.run_in_executor(None, _generate_sync)
    return output