import streamlit as st
import os
from google import genai
from io import StringIO
import prompts # Your refactored prompts.py

st.set_page_config(page_title="Gemini SQL Pro", layout="wide")

# 1. Sidebar Context
st.sidebar.title("🛠️ Settings")
dialect = st.sidebar.selectbox("Dialect", ["PostgreSQL", "MySQL", "SQLite", "Snowflake"])
db_name = st.sidebar.text_input("DB Name", "prod_database")

# Schema Upload Logic
st.sidebar.subheader("📄 Schema")
uploaded_file = st.sidebar.file_uploader("Upload DDL (.sql)", type=["sql", "txt"])
default_schema = uploaded_file.getvalue().decode("utf-8") if uploaded_file else ""

schema = st.sidebar.text_area("Schema Context", value=default_schema, height=250)
show_explain = st.sidebar.checkbox("Explain Logic", value=False)

# 2. Secret Handling (The Fix)
api_key = os.environ.get("GOOGLE_API_KEY") or st.secrets.get("GOOGLE_API_KEY")
if not api_key:
    st.error("API Key missing!")
    st.stop()

client = genai.Client(api_key=api_key)

# 3. Chat & Generation
if "messages" not in st.session_state: st.session_state.messages = []
if "last_sql" not in st.session_state: st.session_state.last_sql = ""

for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

if prompt := st.chat_input("Ask for a query..."):
    st.chat_message("user").write(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})

    explain_req = "Explain your steps." if show_explain else "Code only."
    sys_msg = prompts.SYSTEM_INSTRUCTION.format(
        dialect=dialect, db_name=db_name, schema=schema, explain=explain_req
    )

    with st.spinner("Gemini is thinking..."):
        response = client.models.generate_content(
            model="gemini-2.0-flash",
            config={'system_instruction': sys_msg},
            contents=prompt
        )
        st.session_state.last_sql = response.text
        st.chat_message("assistant").code(response.text, language="sql")
        st.session_state.messages.append({"role": "assistant", "content": response.text})

# 4. Download
if st.session_state.last_sql:
    st.download_button("💾 Download SQL", st.session_state.last_sql, "query.sql")