import os
import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
from groq import Groq
import io

# Load Groq API key from Hugging Face secrets
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

if not GROQ_API_KEY:
    raise ValueError("❌ GROQ_API_KEY not found. Please add it in Hugging Face Secrets.")

client = Groq(api_key=GROQ_API_KEY)

# Function to query Groq LLM
def ask_groq(prompt):
    response = client.chat.completions.create(
        model="mixtral-8x7b-32768",  # good for structured reasoning
        messages=[
            {"role": "system", "content": "You are a data visualization assistant. Convert user requests into Python Pandas + Matplotlib code."},
            {"role": "user", "content": prompt}
        ],
        temperature=0.3,
        max_tokens=500
    )
    return response.choices[0].message["content"]

# Function to generate chart
def generate_dashboard(file, query):
    try:
        df = pd.read_csv(file.name)

        # Build prompt for Groq
        prompt = f"""
        I have the following dataset columns: {', '.join(df.columns)}.
        User wants: {query}.
        Write Python code using Pandas + Matplotlib to create the chart. 
        The dataframe variable is `df`. Do not load the CSV again.
        Only return the code inside ```python ... ``` fences.
        """

        code_response = ask_groq(prompt)

        # Extract python code block
        if "```python" in code_response:
            code = code_response.split("```python")[1].split("```")[0].strip()
        else:
            code = code_response.strip()

        # Run the code safely
        local_vars = {"df": df, "plt": plt}
        exec(code, {}, local_vars)

        # Save plot to buffer
        buf = io.BytesIO()
        plt.savefig(buf, format="png")
        buf.seek(0)
        plt.close()

        return buf, code

    except Exception as e:
        return None, f"❌ Error: {str(e)}"

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 📊 AI-Powered Dashboard Generator (Groq + Matplotlib)")
    
    file_input = gr.File(label="Upload CSV", file_types=[".csv"])
    query_input = gr.Textbox(label="What do you want to see?", placeholder="e.g. Show a bar chart of subscribers per subject")
    output_plot = gr.Image(label="Generated Chart")
    output_code = gr.Code(label="Generated Python Code")

    run_btn = gr.Button("Generate Chart")
    run_btn.click(fn=generate_dashboard, inputs=[file_input, query_input], outputs=[output_plot, output_code])

demo.launch()