from pandasai import SmartDataframe
import pandas as pd
from pandasai.llm.google_gemini import GoogleGemini

# df = pd.DataFrame({
#     "country": [
#         "United States",
#         "United Kingdom",
#         "France",
#         "Germany",
#         "Italy",
#         "Spain",
#         "Canada",
#         "Australia",
#         "Japan",
#         "China",
#     ],
#     "gdp": [
#         19294482071552,
#         2891615567872,
#         2411255037952,
#         3435817336832,
#         1745433788416,
#         1181205135360,
#         1607402389504,
#         1490967855104,
#         4380756541440,
#         14631844184064,
#     ],
#     "happiness_index": [6.94, 7.16, 6.66, 7.07, 6.38, 6.4, 7.23, 7.22, 5.87, 5.12],
# })
llm = GoogleGemini(api_key="AIzaSyCW-TP3IlbbdQmp_nDMEEaip0uVcK9lbgA")

def pandas_ai_res(file, input_text, llm=llm):
    # check if the file is csv or xlsx
    if file.name.endswith(".csv"):
        df = pd.read_csv(file.name)
    elif file.name.endswith(".xlsx"):
        df = pd.read_excel(file.name)

    sdf = SmartDataframe(df, config={"llm": llm})

    return sdf.chat(input_text)


# sdf = SmartDataframe(df, config={"llm": llm})
# sdf.chat("Return the top 5 countries by GDP")