import streamlit as st
import pandas as pd
import requests
from pandasql import sqldf

# 🛠 Page setup
st.set_page_config(page_title="🧠 Excel SQL Assistant", layout="centered")
st.title("📊 Excel to SQL with Together AI")

# 🔐 Load API Key securely
TOGETHER_API_KEY = st.secrets.get("TOGETHER_API_KEY", None)

if not TOGETHER_API_KEY:
    st.error("❌ Together API key not found. Please add it under Secrets on Hugging Face.")
    st.stop()

# 📂 File upload
uploaded_file = st.file_uploader("Upload your Excel file", type=["xlsx"])

if uploaded_file is not None:
    try:
        df = pd.read_excel(uploaded_file)
        st.success("✅ File loaded successfully")
        st.subheader("📄 Preview of Uploaded Data")
        st.dataframe(df.head(10))
    except Exception as e:
        st.error(f"❌ Error reading Excel file: {e}")
        st.stop()
else:
    st.info("📂 Please upload an Excel file to begin.")

# 💬 User input
user_query = st.text_input("💬 Ask a question about your dataset")

# 🚀 On click
if st.button("Generate & Run SQL"):
    if uploaded_file is not None and user_query.strip():
        try:
            # 🧠 Prompt construction
            preview = df.head(5).to_string(index=False)
            prompt = f"""
You are a SQL expert assistant. Generate an SQL query for a pandas DataFrame named 'df' based on the user's question.

Data preview:
{preview}

User question:
{user_query}

IMPORTANT:
- Only return a valid SQL query (no Python, no comments, no markdown, no explanations).
- Use SQL compatible with SQLite.
- Table name is 'df'.
- If column names contain spaces or special characters, wrap them in double quotes.
- DO NOT return unmatched quotes or symbols like *, ', ", ; — unless used correctly in SQL syntax.
"""

            # 🔗 Together API request
            with st.spinner("🧠 Thinking..."):
                headers = {
                    "Authorization": f"Bearer {TOGETHER_API_KEY}",
                    "Content-Type": "application/json"
                }
                payload = {
                    "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
                    "max_tokens": 256,
                    "temperature": 0.3,
                    "messages": [
                        {"role": "system", "content": "You are a SQL assistant. You return only clean SQL code."},
                        {"role": "user", "content": prompt}
                    ]
                }

                response = requests.post(
                    "https://api.together.xyz/v1/chat/completions",
                    headers=headers,
                    json=payload
                )

                if response.status_code != 200:
                    raise ValueError(f"API Error: {response.status_code} - {response.text}")

                # 🧼 Clean the response
                sql_query = response.json()['choices'][0]['message']['content'].strip()
                sql_query = sql_query.replace("```sql", "").replace("```", "").strip()

                # Auto-correct common malformed syntax
                sql_query = sql_query.replace("*\"", "*").replace("\" FROM", " FROM").replace(";*", ";")
                if sql_query.startswith('"SELECT'):
                    sql_query = sql_query.replace('"SELECT', 'SELECT')
                if sql_query.startswith('SELECT "'):
                    sql_query = sql_query.replace('SELECT "', 'SELECT ')

                st.code(sql_query, language='sql')

                # 🧪 Execute SQL
                try:
                    result_df = sqldf(sql_query, {"df": df})
                    st.success("✅ SQL query executed successfully!")
                    st.dataframe(result_df)
                except Exception as e:
                    st.error(f"⚠️ Failed to run SQL query:\n{e}")

        except Exception as e:
            st.error(f"❌ Something went wrong: {e}")
    else:
        st.warning("📢 Please upload a file and enter your question.")