import streamlit as st
from google import genai
from google.genai import types
import json
from PIL import Image

# Page config
st.set_page_config(page_title="Gemini API Physics Questions", layout="wide")

# Title
st.title("AP Physics C Practice Question Generator")

# Initialize the API key input
api_key = st.text_input("Enter your Google API Key", type="password")

if api_key:
    # Initialize the client
    client = genai.Client(
        api_key=api_key,
        http_options={'api_version': 'v1alpha'},
    )
    
    model_name = "gemini-2.0-flash-thinking-exp-01-21"
    
    if st.button("Generate New Question"):
        # Create containers for thinking and answer
        thinking_container = st.container()
        answer_container = st.container()
        
        # Generate content
        response = client.models.generate_content_stream(
            model=model_name,
            config={'thinking_config': {'include_thoughts': True}},
            contents="Give me a practice question I can use for the AP Physics C exam?"
        )
        
        mode = 'starting'
        
        # Process the response
        for chunk in response:
            for part in chunk.candidates[0].content.parts:
                if part.thought:
                    if mode != "thinking":
                        with thinking_container:
                            st.subheader("Thinking Process")
                        mode = "thinking"
                    with thinking_container:
                        st.write(part.text)
                else:
                    if mode != "answering":
                        with answer_container:
                            st.subheader("Generated Question")
                        mode = "answering"
                    with answer_container:
                        st.write(part.text)
else:
    st.warning("Please enter your Google API key to start generating questions.")