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

def generate_text(user_input):
    try:
        # Initialize client with API key
        genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
        
        model = "gemini-2.5-pro-exp-03-25"
        contents = [
            types.Content(
                role="user",
                parts=[types.Part.from_text(user_input)],
            ),
        ]
        config = types.GenerationConfig(
            temperature=2,
            response_mime_type="text/plain",
        )

        # Generate response
        response = []
        for chunk in genai.generate_content_stream(
            model=model,
            contents=contents,
            generation_config=config,
        ):
            response.append(chunk.text)
            
        return ''.join(response)
    
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(lines=4, placeholder="Enter your prompt...", label="Input"),
    outputs=gr.Textbox(label="Output"),
    title="Gemini 2.5 Pro Demo",
    description="Generate text with Gemini 2.5 Pro (streaming enabled)",
    allow_flagging="never",
)

if __name__ == "__main__":
    iface.launch()