run_code.py

# run.py
import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import io
from scanDataFrames import read_local_csv, suggest_prompts
from promptToCode import prompt_to_code
from ask_gemini import ask_gemini
import uuid

def run_code(df: pd.DataFrame, code: str):
    local_ns = {
        "df": df,
        "pd": pd,
        "np": np,
        "plt": plt,
    }

    SAFE_BUILTINS = {
        "len": len, "range": range, "min": min, "max": max, "abs": abs, "sum": sum,
        "print": print, "str": str, "int": int, "float": float, "bool": bool,
        "list": list, "dict": dict, "df": df,
        "pd": pd,
        "np": np,
        "plt": plt, "isinstance": isinstance, "type": type,
        "object": object, "__import__": __import__,
    }

    old_stdout = sys.stdout
    stdout_buf = io.StringIO()
    sys.stdout = stdout_buf

    try:
        exec(code, {"__builtins__": SAFE_BUILTINS}, local_ns)

        # 1️⃣ Gemini text-based result (TOP PRIORITY)
        result_val = local_ns.get("result")
        if isinstance(result_val, str):
            return {"type": "text", "output": result_val}

        # 2️⃣ DataFrame return
        if isinstance(result_val, pd.DataFrame):
            return {"type": "dataframe", "df": result_val}

        # 3️⃣ Plot created
        if plt.get_fignums():
            save_dir = os.path.join(os.getcwd(), "plots")
            os.makedirs(save_dir, exist_ok=True)
            fname = f"plot_{uuid.uuid4().hex[:8]}.png"
            fpath = os.path.join(save_dir, fname)
            plt.savefig(fpath, bbox_inches="tight", dpi=150)
            plt.close("all")
            return {"type": "image", "path": fpath}

        # 4️⃣ Any printed output
        out = stdout_buf.getvalue().strip()
        if out:
            return {"type": "text", "output": out}

        # 5️⃣ Nothing happened
        return {"type": "text", "output": "⚠️ No output produced. (Code ran but returned nothing)"}

    except NameError as e:
        return {"type": "text", "output": f"❌ NameError: {e}\n💡 Check variable/column names."}
    except KeyError as e:
        return {"type": "text", "output": f"❌ KeyError: column {e} not found in dataset."}
    except Exception as e:
        return {"type": "text", "output": f"❌ ERROR: {e}"}
    finally:
        sys.stdout = old_stdout
     
        
        
        
        
        
# Example usage
if __name__ == "__main__":
    df = read_local_csv("salary.csv")
    prompts = suggest_prompts(df)

    print("\n=== Suggested Prompts ===")
    for i, p in enumerate(prompts, 1):
        print(f"{i}. {p}")

    while True:
        user_input = input("\n👉 Enter prompt number or text (or 'exit'): ").strip()

        if user_input.lower() == "exit":
            break

        # 🔥 If user entered a number, map it
        if user_input.isdigit():
            idx = int(user_input) - 1
            if 0 <= idx < len(prompts):
                prompt = prompts[idx]
                print(f"➡️ Using prompt: {prompt}")
            else:
                print("❌ Invalid number. Try again.")
                continue
        else:
            # They typed a sentence instead of a number
            prompt = user_input

        # 🎯 Send to translator
        code = prompt_to_code(prompt, df)

        if not code:
            print("💬 Unrecognized prompt. (Next step: GPT fallback here)")
            clean_prompt = prompt.replace('"', "").replace("'", "").strip()

            code = ask_gemini(clean_prompt, df)

        # print("\n=== Generated Code ===\n", code)

        # ⚙️ Execute
        result = run_code(df, code)

        # 🧾 Show results
        if result["type"] == "image":
            print(f"🖼️ Image saved at: {result['path']}")
        elif result["type"] == "dataframe":
            print("📊 DataFrame Output:")
            print(result["df"].head())
        else:
            print("📝 Output:")
            print(result["output"])