"""Google Gemini API integration for Akern-Genai project.

This module provides functionality to generate content using Google's Gemini model
with Vertex AI RAG (Retrieval-Augmented Generation) support.
"""

from google import genai
from google.genai import types
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Vertex AI RAG Corpus resource path
CORPUS: str = "projects/520464122471/locations/europe-west3/ragCorpora/2305843009213693952"

# Gemini model name
GEMINI_MODEL: str = "gemini-3-pro-preview"


def generate(prompt: str):
    """Generate content using Gemini model with RAG retrieval.

    This function creates a streaming response from the Gemini model,
    augmented with content from the configured RAG corpus.

    Args:
        prompt: The user's input prompt to generate content for.

    Yields:
        str: Text chunks from the generated response.
    """
    client = genai.Client(
        vertexai=True,
    )

    contents = [
        types.Content(role="user", parts=[types.Part.from_text(text=prompt)]),
    ]
    tools = [
        types.Tool(
            retrieval=types.Retrieval(
                vertex_rag_store=types.VertexRagStore(
                    rag_resources=[
                        types.VertexRagStoreRagResource(rag_corpus=CORPUS)],
                )
            )
        )
    ]

    generate_content_config = types.GenerateContentConfig(
        temperature=1,
        top_p=0.95,
        max_output_tokens=65535,
        safety_settings=[
            types.SafetySetting(
                category="HARM_CATEGORY_HATE_SPEECH", threshold="OFF"),
            types.SafetySetting(
                category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="OFF"
            ),
            types.SafetySetting(
                category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="OFF"
            ),
            types.SafetySetting(
                category="HARM_CATEGORY_HARASSMENT", threshold="OFF"),
        ],
        tools=tools,
        thinking_config=types.ThinkingConfig(
            thinking_level="HIGH",
        ),
    )

    for chunk in client.models.generate_content_stream(
        model=GEMINI_MODEL,
        contents=contents,
        config=generate_content_config,
    ):
        if (
            not chunk.candidates
            or not chunk.candidates[0].content
            or not chunk.candidates[0].content.parts
        ):
            continue

        yield chunk.text


if __name__ == "__main__":
    for chunk in generate("Come si calcola il rapporto sodio potassio?"):
        print(chunk, end="")