import streamlit as st
import pandas as pd
import os
from csv_agent import generate_code

st.set_page_config(page_title="CSV Assistant with Hugging Face", layout="wide")
st.title("🤖 CSV Assistant: Powered by FLAN-T5-Large")

# Check API token
if os.getenv("HUGGINGFACEHUB_API_TOKEN") is None:
    st.error("❌ API Token missing! Please set HUGGINGFACEHUB_API_TOKEN in your Space Secrets.")
    st.stop()

st.sidebar.title("Upload Your CSV")
uploaded_file = st.sidebar.file_uploader("Choose a CSV file", type=["csv"])

if uploaded_file:
    df = pd.read_csv(uploaded_file)
    st.write("### Preview of your data:", df.head())

    prompt = st.text_input("💬 What would you like to do with the data?")
    if prompt:
        with st.spinner("Generating Python code using AI..."):
            code = generate_code(prompt)
            st.code(code, language="python")
            
            try:
                exec_globals = {"df": df}
                exec(code, exec_globals)
                df = exec_globals["df"]
                st.success("✅ Data updated successfully!")
                st.write(df.head())

                csv = df.to_csv(index=False).encode('utf-8')
                st.download_button("⬇️ Download Updated CSV", csv, "updated_data.csv", "text/csv")
            except Exception as e:
                st.error(f"⚠️ Error while executing AI-generated code: {e}")