import streamlit as st
import openai
import pandas as pd
import subprocess
import os
import io
import contextlib

# Setup OpenAI
openai.api_key = os.getenv("OPENAI_API_KEY")

st.set_page_config(page_title="🧠 Build Any Model", layout="centered")
st.title("🤖 Dynamic Code Generator Chatbot")

prompt = st.text_input("🧠 Tell me what you want to build:")

uploaded_file = st.file_uploader("📂 Upload your dataset (CSV)", type=["csv"])

if st.button("🚀 Generate and Run"):
    if prompt and uploaded_file:
        df = pd.read_csv(uploaded_file)
        st.success("✅ File loaded successfully.")
        
        # Call OpenAI to generate code
        system_prompt = "You are a Python code generator. Only output clean, executable Python code inside one code block. No explanation."
        client = openai.OpenAI()
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,
        )
        code = response.choices[0].message.content
        if "```python" in code:
            code = code.split("```python")[-1].split("```")[0].strip()
        
        st.code(code, language="python")
        
        # Execute the generated code
        output_buffer = io.StringIO()
        try:
            with contextlib.redirect_stdout(output_buffer):
                exec(code, {"df": df})
            st.success("✅ Code executed successfully.")
            st.text(output_buffer.getvalue())
        except Exception as e:
            st.error(f"❌ Error during execution: {e}")
    else:
        st.warning("⚠️ Please upload a dataset and enter a prompt.")