import os
import json
from google import genai
from google.genai import types
from dotenv import load_dotenv
from prompts import *

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

def format_search_results(results):
    """Format Search Results"""
    formatted_output = ""
    
    for i in range(1, 4):
        formatted_output += f"### {i}. {results[f'keyword{i}']}\n"
        formatted_output += f"{results[f'summary{i}']}"
        formatted_output += "\n"
    
    return formatted_output

def grounding_with_google_search(stt_data, content_type = "국민의힘"):
    """
    Extract Keywords and Perform Google Search
    """
    # Create a client for Google GenAI
    client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

    # Define a sample schema for search results
    sample_schema = {
        "keyword1": "Core Keyword 1",
        "summary1": "Summarize Search Result about Core Keyword 1",
        "keyword2": "Core Keyword 2",
        "summary2": "Summarize Search Result about Core Keyword 2",
        "keyword3": "Core Keyword 3",
        "summary3": "Summarize Search Result about Core Keyword 3"
    }

    # Select the appropriate system prompt based on content_type
    if content_type == "Agents for Amazon Bedrock":
        system_prompt = BEDROCK_SEARCH_PROMPT
    elif content_type == "Bundesliga Fan Experience":
        system_prompt = BUNDESLIGA_SEARCH_PROMPT
    elif content_type == "AWS_2024_recap":
        system_prompt = AWS_SEARCH_PROMPT

    # Format the system prompt with the sample schema
    system_prompt = system_prompt.format(sample_schema=sample_schema)

    # Prepare the human message with the input script
    human_message = f"""
    ## Input Script
    {stt_data}
    """

    # Generate content using the Google GenAI client
    response = client.models.generate_content(
        model="gemini-2.0-flash-001",
        contents=human_message,
        config=types.GenerateContentConfig(
            system_instruction=system_prompt,
            response_mime_type="application/json",
            tools=[
                types.Tool(
                    google_search=types.GoogleSearchRetrieval(
                        dynamic_retrieval_config = types.DynamicRetrievalConfig(
                            mode=types.DynamicRetrievalConfigMode.MODE_UNSPECIFIED,
                            dynamic_threshold=0.0
                        ) 
                    )
                )
            ]
        )
    )

    # Parse the response text and format the search results
    text = response.text
    results = json.loads(text)
    return format_search_results(results)