import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
import io
from PIL import Image, ImageDraw, ImageFont
import traceback

def process_file(api_key, file, instructions):
    try:
        # Read uploaded file
        if file.name.endswith('.csv'):
            df = pd.read_csv(file.name)
        elif file.name.endswith('.xlsx'):
            df = pd.read_excel(file.name)
        else:
            raise ValueError("Unsupported file format")

        # Generate sample visualizations (replace with actual logic)
        fig1, ax1 = plt.subplots()
        df.plot(kind='bar', ax=ax1)
        ax1.set_title("Sample Bar Chart")
        
        fig2, ax2 = plt.subplots()
        df.plot(kind='line', ax=ax2)
        ax2.set_title("Sample Line Chart")
        
        fig3, ax3 = plt.subplots()
        df.plot(kind='hist', ax=ax3)
        ax3.set_title("Sample Histogram")

        # Convert plots to PIL Images
        def fig_to_image(fig):
            buf = io.BytesIO()
            fig.savefig(buf, format='png')
            buf.seek(0)
            return Image.open(buf)

        return [
            fig_to_image(fig1),
            fig_to_image(fig2),
            fig_to_image(fig3)
        ]

    except Exception as e:
        error_message = f"{str(e)}\n{traceback.format_exc()}"
        return [generate_error_image(error_message)] * 3

def generate_error_image(message):
    """Create error indication image with message"""
    try:
        img = Image.new('RGB', (800, 400), color=(255, 255, 255))
        draw = ImageDraw.Draw(img)
        font = ImageFont.load_default()
        
        # Wrap text
        lines = []
        for line in message.split('\n'):
            if len(line) > 80:
                lines.extend([line[i:i+80] for i in range(0, len(line), 80)])
            else:
                lines.append(line)

        y_text = 10
        for line in lines[:20]:  # Limit to 20 lines
            draw.text((10, y_text), line, font=font, fill=(255, 0, 0))
            y_text += 15

        return img
    except Exception as e:
        return Image.new('RGB', (800, 400), color=(255, 255, 255))

# Gradio interface
with gr.Blocks(theme=gr.themes.Default(spacing_size="lg")) as demo:
    gr.Markdown("# AutoData Visualizer")
    
    with gr.Row():
        api_key = gr.Textbox(label="Gemini API Key", type="password")
        file = gr.File(label="Upload Data File", file_types=[".csv", ".xlsx"])
    
    instructions = gr.Textbox(label="Visualization Instructions")
    submit = gr.Button("Generate Insights", variant="primary")
    
    with gr.Row():
        outputs = [gr.Image(label=f"Visualization {i+1}", width=600) for i in range(3)]

    submit.click(
        process_file,
        inputs=[api_key, file, instructions],
        outputs=outputs
    )

if __name__ == "__main__":
    demo.launch()