import streamlit as st
from google import genai

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

st.title("⚡ Gemini API Tester")
st.markdown("Enter your Google API key to load available Gemini models.")

# --- Initialize session state ---
if "api_key" not in st.session_state:
    st.session_state.api_key = ""

if "api_validated" not in st.session_state:
    st.session_state.api_validated = False

if "models" not in st.session_state:
    st.session_state.models = []

if "client" not in st.session_state:
    st.session_state.client = None

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

    # 🔹 API Key Input (only show if not validated)
    if not st.session_state.api_validated:
        api_key_input = st.text_input(
            "Enter your Google API Key",
            type="password"
        )

        if st.button("Validate API Key"):
            if not api_key_input:
                st.warning("⚠️ Please enter an API key")
            else:
                try:
                    client = genai.Client(api_key=api_key_input)

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

                    if not models:
                        st.error("❌ API key valid, but no Gemini models found")
                    else:
                        # Save to session state
                        st.session_state.api_key = api_key_input
                        st.session_state.client = client
                        st.session_state.models = models
                        st.session_state.api_validated = True

                        st.success("✅ API Key validated successfully")

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

    # 🔹 After validation (API key hidden)
    else:
        st.success("🔐 API Key validated")

        selected_model = st.selectbox(
            "Available Models",
            st.session_state.models
        )

        if st.button("🔄 Change API Key"):
            st.session_state.api_key = ""
            st.session_state.api_validated = False
            st.session_state.models = []
            st.session_state.client = None
            st.experimental_rerun()

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

# Show 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 st.session_state.api_validated:
        st.error("❌ Please validate your API key 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 = st.session_state.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}")