import base64
import gradio as gr
import os
import json
from google import genai
from google.genai import types
from gradio_client import Client

def clean_json_string(json_str):
    """
    Removes any comments or prefixes before the actual JSON content.
    """
    # Find the first occurrence of '{'
    json_start = json_str.find('{')
    if json_start == -1:
        # If no '{' is found, try with '[' for arrays
        json_start = json_str.find('[')
        if json_start == -1:
            return json_str  # Return original if no JSON markers found
    
    # Extract everything from the first JSON marker
    cleaned_str = json_str[json_start:]
    return cleaned_str
    # Verify it's valid JSON
    try:
        json.loads(cleaned_str)
        return cleaned_str
    except json.JSONDecodeError:
        return json_str  # Return original if cleaning results in invalid JSON

def generate(input_text):
    try:
        client = genai.Client(
            api_key=os.environ.get("GEMINI_API_KEY"),
        )
    except Exception as e:
        return f"Error initializing client: {e}.  Make sure GEMINI_API_KEY is set."

    model = "gemini-2.0-flash"
    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(text=input_text),
            ],
        ),
    ]
    tools = [
        types.Tool(google_search=types.GoogleSearch()),
    ]
    generate_content_config = types.GenerateContentConfig(
        temperature=0.4,
        thinking_config = types.ThinkingConfig(
            thinking_budget=0,
        ),
        tools=tools,
        response_mime_type="text/plain",
    )


    response_text = ""
    try:
      for chunk in client.models.generate_content_stream(
          model=model,
          contents=contents,
          config=generate_content_config,
      ):
          response_text += chunk.text
    except Exception as e:
        return f"Error during generation: {e}"
    data = clean_json_string(response_text)
    data = data[:-1]
    return response_text, ""



def generate1(input_text):
    client = genai.Client(
        api_key=os.environ.get("GEMINI_API_KEY"),
    )

    model = "gemini-1.5-pro"
    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(text=f"return json object with keys name and email. name = {input_text}  search the web for the email value. return json object only, no additional text or comments"),
            ],
        ),
    ]
    tools = [
        types.Tool(googleSearchRetrieval=types.DynamicRetrievalConfig(dynamicThreshold=0.3, mode=types.DynamicRetrievalConfigMode.MODE_DYNAMIC)),
    ]
    generate_content_config = types.GenerateContentConfig(
        temperature=0.45,
        tools=tools,
        response_mime_type="text/plain",
    )

    for chunk in client.models.generate_content_stream(
        model=model,
        contents=contents,
        config=generate_content_config,
    ):
        print(chunk.text, end="")
    

if __name__ == '__main__':

    with gr.Blocks() as demo:
        title=gr.Markdown("# Gemini 2.0 Flash + Websearch")
        output_textbox = gr.Markdown()
        input_textbox = gr.Textbox(lines=3, label="", placeholder="Enter event details here...")
        submit_button = gr.Button("send")
        submit_button.click(fn=generate,inputs=input_textbox,outputs=[output_textbox, input_textbox])
    demo.launch(show_error=True)