import streamlit as st
import openai
import pyodbc
import pandas as pd

# -----------------------------
# 🔐 1. Configuration
# -----------------------------
openai.api_key = "YOUR_OPENAI_API_KEY"

# SQL Server connection string
server = 'dcdocstgdb01.dc68032.easn.morningstar.com'
database = 'DocumentAcquisition'
username = 'DocuUser'
password = 'StagTest@1'

conn_str = (
    f"DRIVER={{ODBC Driver 17 for SQL Server}};"
    f"SERVER={server};"
    f"DATABASE={database};"
    f"UID={username};"
    f"PWD={password}"
)

# -----------------------------
# 🤖 2. Natural Language → SQL
# -----------------------------
def generate_sql(nl_question):
    prompt = f"""
You are a helpful assistant that converts user questions into SQL Server queries.

Question: "{nl_question}"

SQL Query:
"""
    response = openai.ChatCompletion.create(
        model="gpt-4",  # or gpt-3.5-turbo
        messages=[{"role": "user", "content": prompt}],
        temperature=0
    )
    sql_query = response.choices[0].message["content"].strip()
    return sql_query

# -----------------------------
# 🧠 3. Run SQL and return results
# -----------------------------
def run_query(sql_query):
    try:
        conn = pyodbc.connect(conn_str)
        df = pd.read_sql(sql_query, conn)
        conn.close()
        return df
    except Exception as e:
        return str(e)

# -----------------------------
# 🎛️ 4. Streamlit UI
# -----------------------------
st.title("🧠 GenAI SQL Server Q&A")
st.markdown("Ask a question in natural language. I’ll generate SQL and return the results from your database.")

user_question = st.text_input("📌 Ask your question:")

if st.button("Submit") and user_question:
    with st.spinner("Thinking..."):
        sql_query = generate_sql(user_question)
        st.code(sql_query, language="sql")
        
        result = run_query(sql_query)
        
        if isinstance(result, pd.DataFrame):
            st.success("✅ Results:")
            st.dataframe(result)
        else:
            st.error(f"❌ Error:\n{result}")