import os
import gradio as gr
from google import genai
from google.genai import types

# Client initialization
client_genai = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))

def generate_with_history(input_text, history):
    """
    - input_text: The new message from JS
    - history: A list of previous messages [user, model, user, model...]
    """
    if history is None:
        history = []

    # Build the contents list including history
    # Gemini API expects a list of Content objects
    contents = []
    for msg in history:
        contents.append(types.Content(
            role=msg["role"],
            parts=[types.Part.from_text(text=msg["content"])]
        ))
    
    # Add the current user prompt
    contents.append(types.Content(
        role="user",
        parts=[types.Part.from_text(text=input_text)]
    ))

    tools = [types.Tool(google_search=types.GoogleSearch())]
    config = types.GenerateContentConfig(tools=tools)

    try:
        response = client_genai.models.generate_content(
            model="gemini-2.5-flash-lite",
            contents=contents,
            config=config,
        )
        
        response_text = response.text
        
        # Update history: add the current exchange
        history.append({"role": "user", "content": input_text})
        history.append({"role": "model", "content": response_text})
        
        # Return the response for JS and the updated history for the State
        return response_text, history

    except Exception as e:
        return f"Error: {str(e)}", history

# Building the Blocks interface
with gr.Blocks() as demo:
    # State component (always invisible)
    history_state = gr.State([])
    
    # Change visible=False to True (or remove it)
    input_box = gr.Textbox(label="Input", visible=True) 
    output_box = gr.Markdown(label="Response", visible=True)
    
    btn = gr.Button("Submit", visible=True)
    
    btn.click(
        fn=generate_with_history,
        inputs=[input_box, history_state],
        outputs=[output_box, history_state],
        api_name="generate"
    )
if __name__ == "__main__":
    demo.launch()