import streamlit as st
from google import genai

# Page Config
st.set_page_config(page_title="Gemini API Tester", page_icon="⚡")

# --- Title ---
st.title("⚡ Gemini API Tester")
st.markdown("Enter your Google API key to see which Gemini models are available.")

# --- Sidebar ---
with st.sidebar:
    st.header("🔐 API Configuration")

    api_key = st.text_input(
        "Enter your Google API Key",
        type="password"
    )

    models = []
    client = None

    if api_key:
        try:
            # Initialize client
            client = genai.Client(api_key=api_key)

            # Fetch available models
            models = [
                model.name
                for model in client.models.list()
                if "gemini" in model.name.lower()
            ]

            if models:
                st.success("✅ API Key is valid")
            else:
                st.warning("⚠️ API key is valid, but no Gemini models found")

        except Exception as e:
            st.error("❌ Invalid API Key")
            st.caption(str(e))

    # Model selection (only show if models exist)
    selected_model = None
    if models:
        selected_model = st.selectbox(
            "Available Models",
            models
        )

# --- Chat State ---
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat history
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# --- Chat Input ---
if prompt := st.chat_input("Ask something..."):

    if not client or not selected_model:
        st.error("❌ Please enter a valid API key and select a model first.")
        st.stop()

    # User message
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Assistant response
    with st.chat_message("assistant"):
        placeholder = st.empty()
        placeholder.markdown("Thinking...")

        try:
            response = client.models.generate_content(
                model=selected_model,
                contents=prompt
            )

            reply = response.text
            placeholder.markdown(reply)

            st.session_state.messages.append(
                {"role": "assistant", "content": reply}
            )

        except Exception as e:
            placeholder.error(f"❌ Error: {e}")