import streamlit as st
import pandas as pd
import os
from groq import Groq

# Set page configuration
st.set_page_config(page_title="📊 Excel SQL Assistant", layout="centered")
st.title("🧠 Excel to SQL with GROQ + Streamlit")

# Load Groq API key from Streamlit secrets
GROQ_API_KEY = st.secrets.get("GROQ_API_KEY", None)

# Fallback warning if API key not found
if not GROQ_API_KEY:
    st.error("❌ GROQ API key not found. Please set it in your Hugging Face `secrets` tab.")
    st.stop()

# Initialize Groq client
client = Groq(api_key=GROQ_API_KEY)

# Upload Excel file
uploaded_file = st.file_uploader("📂 Upload your Excel file", type=["xlsx"])
st.write("📁 Uploaded file object:", uploaded_file)  # Debugging aid

# Text input for user query
user_query = st.text_input("💬 Ask a question about your data")

# Generate button
if st.button("🚀 Generate & Run SQL"):
    if uploaded_file is not None and user_query.strip() != "":
        try:
            df = pd.read_excel(uploaded_file)
            st.success("✅ File successfully loaded!")

            # Create schema-like string from the DataFrame
            preview = df.head(5).to_string(index=False)
            schema = f"Here are the first few rows of the dataset:\n\n{preview}"

            # Prompt to Groq
            full_prompt = f"""You are a data expert. Write a Pandas code snippet to answer the following question based on the data:\n\n{schema}\n\nQuestion: {user_query}"""

            with st.spinner("⏳ Generating SQL-like response using Groq..."):
                response = client.chat.completions.create(
                    model="mixtral-8x7b-32768",
                    messages=[
                        {"role": "system", "content": "You are an expert data assistant that writes pandas code based on user questions."},
                        {"role": "user", "content": full_prompt}
                    ]
                )

            answer = response.choices[0].message.content
            st.code(answer, language="python")

            # Optionally: try to run it
            try:
                local_vars = {"df": df.copy()}
                exec(answer, {}, local_vars)
                if 'df' in local_vars:
                    st.dataframe(local_vars['df'])
            except Exception as e:
                st.warning(f"⚠️ Could not execute the code:\n\n{e}")

        except Exception as e:
            st.error(f"❌ Error processing file: {e}")
    else:
        st.warning("⚠️ Please upload a file and enter a question.")