import chainlit as cl
import pandas as pd
import io
import matplotlib.pyplot as plt
import base64
from io import BytesIO
from pandasai import SmartDataframe
import pandas as pd
from pandasai.llm import OpenAI
from io import StringIO
import matplotlib.pyplot as plt
import csv
from collections import defaultdict
import os


got_csv = False


def find_most_valuable_feature(csv_file):
    print("find_most_valuable_feature")
    print(csv_file)

    smart_llm = OpenAI(api_token=os.environ["OPENAI_API_KEY"])

    


    # Initialize a defaultdict to store column data
    columns = defaultdict(list)
  
    # Read the CSV file and populate the defaultdict
    with open("upload.csv") as f:
        reader = csv.reader(f)
        headers = next(reader)

        for row in reader:
            for header, value in zip(headers, row):
                columns[header].append(value)

    # Manually create a DataFrame from the defaultdict
    smart_df = pd.DataFrame({
        "ID": columns["ID"],
        "Date and Time": columns["Date and Time"],
        "Business Unit": columns["Business Unit"],
        "Usage Change": columns["Usage Change"],
        "Wolftech Improvement": columns["Wolftech Improvement"],
        "Likelihood to Recommend": columns["Likelihood to Recommend"],
        "Effective Training": columns["Effective Training"],
        "Most Valuable Feature": columns["Most Valuable Feature"]
    })

    smart_df = SmartDataframe(smart_df, config={"llm": smart_llm})
    out = smart_df.chat('Summarize the top three "Most Valuable Feature" for people where Usage Changed was Increased?')

    print(out)


    df = out

    # Plotting
    plt.figure(figsize=(10, 6))
    plt.bar(df["Most Valuable Feature"], df["Count"], color='blue')
    plt.xlabel('Most Valuable Feature')
    plt.ylabel('Count')
    plt.title('Count of Most Valuable Features')
    plt.xticks(rotation=45, ha="right")  # Rotate labels for better readability
    plt.tight_layout()  # Adjust layout for better fit

    # Save the plot to a BytesIO object
    image_buffer = BytesIO()
    plt.savefig(image_buffer, format='png')
    image_buffer.seek(0)

    return image_buffer





def process_and_analyze_data(csv_file):
    # Read CSV file
    csv_data = pd.read_csv(csv_file)

    # Logging to check data loading
    print(f"CSV Data Loaded: {csv_data.head()}")

    # Count of responses in each category of 'Business Unit'
    business_unit_counts = csv_data['Business Unit'].value_counts()

    # Plotting the count of responses in each 'Business Unit' category
    plt.figure(figsize=(10, 6))
    business_unit_counts.plot(kind='bar')
    plt.title('Count of Responses by Business Unit')
    plt.xlabel('Business Unit')
    plt.ylabel('Count')
    plt.xticks(rotation=45)
    plt.tight_layout()

   # Save the plot to a BytesIO object
    image_buffer = BytesIO()
    plt.savefig(image_buffer, format='png')
    image_buffer.seek(0)

    return image_buffer


# Function to handle message events

@cl.on_message
async def handle_message(message: cl.Message):
    global got_csv
    # Retrieve the CSV file from the message
    csv_file = next(
        (
            io.BytesIO(file.content)
            for file in message.elements or []
            if file.mime and "csv" in file.mime
        ),
        None,
    )

    # Logging to check file retrieval
    print(f"CSV File: {csv_file}")

    if csv_file:
        got_csv = True
        try:
            
            image_buffer = find_most_valuable_feature(csv_file)

            # Get bytes data from BytesIO object and send the image data
            image_data = image_buffer.getvalue()
            name = "chart"
            cl.user_session.set(name, image_data)
            cl.user_session.set("generated_image", name)
           
            await cl.Message(content="Based on the people who increased usage, here are the most valuable features...").send()

            generated_image = cl.user_session.get(name)

            elements = []
            actions = []

            elements = [
                cl.Image(
                    content=generated_image,
                    name=name,
                    display="inline",
                    size="large"
                )
            ]

            await cl.Message(content=name, elements=elements, actions=actions).send()


        except Exception as e:
            await cl.Message(content=f"An error occurred: {str(e)}").send()
    else:
        if not got_csv:
            await cl.Message(content="Please upload a CSV file.").send()


# Run the ChainLit app
if __name__ == "__main__":
    cl.run()