import os
import streamlit as st
from pandasai import SmartDataframe
from pandasai.callbacks import BaseCallback
from pandasai.llm import OpenAI
from pandasai.responses.response_parser import ResponseParser
from data import load_data

# Fungsi untuk mendapatkan daftar folder dalam direktori
def list_folders(directory):
    return [name for name in os.listdir(directory) if os.path.isdir(os.path.join(directory, name))]

class StreamlitCallback(BaseCallback):
    def __init__(self, container) -> None:
        """Initialize callback handler."""
        self.container = container

    def on_code(self, response: str):
        self.container.code(response)

class StreamlitResponse(ResponseParser):
    def __init__(self, context) -> None:
        super().__init__(context)

    def format_dataframe(self, result):
        st.dataframe(result["value"])
        return

    def format_plot(self, result):
        st.image(result["value"])
        return

    def format_other(self, result):
        st.write(result["value"])
        return

# Tampilkan judul aplikasi
st.write("# Chat with Tabular Dataframe 🦙")

# Menentukan direktori asal
root_dir = "./data"

# Mendapatkan daftar folder
folders = list_folders(root_dir)

# Menampilkan kotak pemilihan folder
selected_folder = st.selectbox("Pilih folder data", folders)

# Memuat data berdasarkan folder yang dipilih
if selected_folder:
    df = load_data(os.path.join(root_dir, selected_folder))

    # Menampilkan preview data
    with st.expander("🔎 Dataframe Preview"):
        st.write(df.tail(3))

    # Membuat input untuk query
    query = st.text_area("🗣️ Chat with Dataframe")
    container = st.container()

    # Eksekusi query jika ada input
    if query:
        llm = OpenAI(api_token=os.environ["BIT_OPENAI_API_KEY"])
        query_engine = SmartDataframe(
            df,
            config={
                "llm": llm,
                "response_parser": StreamlitResponse,
                # "callback": StreamlitCallback(container),
            },
        )

        answer = query_engine.chat(query)
else:
    st.write("Silakan pilih folder untuk memuat data.")