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")

# ---------------- Session State Defaults ----------------
defaults = {
    "api_key_input": "",
    "api_key": None,
    "api_validated": False,
    "client": None,
    "models": [],
    "selected_model": None,
    "messages": []
}

for k, v in defaults.items():
    if k not in st.session_state:
        st.session_state[k] = v

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

    if not st.session_state.api_validated:
        with st.form("api_form", clear_on_submit=False):
            api_key_input = st.text_input(
                "Enter Google API Key",
                type="password",
                key="api_key_input"
            )
            submitted = st.form_submit_button("Validate API Key")

            if submitted:
                if not api_key_input:
                    st.warning("Please enter an API key.")
                else:
                    try:
                        client = genai.Client(api_key=api_key_input)

                        try:
                            models = [
                                m.name for m in client.models.list()
                                if "gemini" in m.name.lower()
                            ]
                        except Exception:
                            models = []

                        st.session_state.api_key = api_key_input
                        st.session_state.client = client
                        st.session_state.models = models or ["gemini-2.5-flash"]
                        st.session_state.selected_model = st.session_state.models[0]
                        st.session_state.api_validated = True

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

                    except Exception as e:
                        error_msg = str(e).lower()
                        if "401" in error_msg or "403" in error_msg or "unauthorized" in error_msg:
                            st.error("❌ Invalid API key")
                        else:
                            st.warning("⚠️ API key accepted, but model access may be restricted")
                            st.caption(str(e))

    else:
        st.success("🔒 API Key validated (hidden)")

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

        if st.button("🔄 Change API Key"):
            for k in defaults:
                st.session_state[k] = defaults[k]
# ---------------- 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 an 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=st.session_state.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}")