import gradio as gr
import pandas as pd
import re

session_df = None

# ---------- Load file ----------
def load_file(file):
    global session_df

    if file is None:
        return "Please upload a file first."

    try:
        file_path = file  # Gradio provides file path as string

        if file_path.lower().endswith(".csv"):
            session_df = pd.read_csv(file_path)
        else:
            session_df = pd.read_excel(file_path)

        return f"File loaded successfully ({session_df.shape[0]} rows, {session_df.shape[1]} columns)"

    except Exception as e:
        return f"Error loading file: {e}"

# ---------- Ask question ----------
def ask_question(query):
    global session_df

    if session_df is None:
        return None, "Please upload a file first."

    q = query.lower().strip()
    df = session_df.copy()

    try:
        if q.startswith("sum:"):
            col, grp = q.replace("sum:", "").split(" by ")
            result = df.groupby(grp.strip())[col.strip()].sum().reset_index()
            return result, "Done"

        if q.startswith("avg:"):
            col, grp = q.replace("avg:", "").split(" by ")
            result = df.groupby(grp.strip())[col.strip()].mean().reset_index()
            return result, "Done"

        if q.startswith("top"):
            n = int(q.split()[1])
            rest = q.split(":", 1)[1]
            col, grp = rest.split(" by ")
            result = (
                df.groupby(grp.strip())[col.strip()]
                .sum()
                .sort_values(ascending=False)
                .head(n)
                .reset_index()
            )
            return result, "Done"

        if q.startswith("count:"):
            col = q.replace("count:", "").strip()
            result = df[col].value_counts().reset_index()
            result.columns = [col, "count"]
            return result, "Done"

        return None, (
            "Examples:\n"
            "sum: Sales by Region\n"
            "avg: Salary by Department\n"
            "top 5: Revenue by Customer\n"
            "count: Country"
        )

    except Exception as e:
        return None, f"Error: {e}"

# ---------- Simple UI ----------
with gr.Blocks() as demo:
    gr.Markdown("## 📊 Simple Chat with CSV (No AI, No API key)")

    file_input = gr.File(label="Upload CSV or Excel")
    file_status = gr.Textbox(label="File status")

    query = gr.Textbox(
        label="Ask a question",
        placeholder="Example: sum: Sales by Region"
    )

    answer = gr.Dataframe(label="Answer")
    message = gr.Textbox(label="Message")

    file_input.change(load_file, file_input, file_status)
    query.submit(ask_question, query, [answer, message])

demo.launch(share=True)