import pandas as pd
import gradio as gr
import numpy as np
from pandasai import SmartDataframe
from pandasai.llm import OpenAI
from PIL import Image
import base64
from io import BytesIO

def ask(query, file, api_token):
    data = pd.read_csv(file, index_col=0)
    llm = OpenAI(api_token=api_token)
    df = SmartDataframe(data, config={"llm": llm})
    result = df.chat(query)
    if isinstance(result, str) and result.endswith('.png'):
        # Open the image file
        img = Image.open('exports/charts/temp_chart.png')
        
        # Convert the PIL Image to a base64 encoded string
        buffered = BytesIO()
        img.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        
        # Construct the HTML string to display the image
        html = f"<img src='data:image/png;base64,{img_str}' alt='Generated Image' style='width:100%;'>"
        return html
    elif isinstance(result, pd.DataFrame):
        html_table = result.to_html()
        return html_table
    else:
        return result

def on_file_upload(file, api_token):
    data = pd.read_csv(file, index_col=0)
    llm = OpenAI(api_token=api_token)
    df = SmartDataframe(data, config={"llm": llm})
    result = df.chat('What are the name of my columns')
    return result

with gr.Blocks() as demo:
    headers = gr.Markdown("<div style='text-align: center;'><h1>    This is a project from Chasers using Pandas AI</h1></div>")
    
    with gr.Row():
        # First column for the first image
        with gr.Column():
            gr.HTML("""
                <img src="https://d1b66s7evqera.cloudfront.net/f8d9b474-d1d2-4931-b38d-c5f918b8dc82/images/Chasers320x132.png?61554cbec993ea37a9d52399a0a0d3cd" width="320">
            """)
        
        # Second column for the clickable image
        with gr.Column():
            gr.HTML("""
                <a href='https://pandas-ai.com/' target='_blank'>
                    <img src="https://framerusercontent.com/images/tDg1U6HZYYK3azrDbdVXLhPkOlk.png" width="150" style="opacity: 0.3";>
                </a>
            """)
    
    api_token_input = gr.Textbox(lines=1, label="Enter your API token")
    file_input = gr.File() 
    upload_button = gr.Button("Process Uploaded File")
    upload_message = gr.Label("") 
    upload_button.click(on_file_upload, inputs=[file_input, api_token_input], outputs=[upload_message])
    query_input = gr.Textbox(lines=2, label="Use exactly the name of the columns in your querys")
    output_html = gr.HTML()



    # Button to trigger the ask function
    submit_button = gr.Button("Submit")
    submit_button.click(ask, inputs=[query_input, file_input, api_token_input], outputs=[output_html])

demo.launch()