import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
import io
import google.generativeai as genai
from PIL import Image
import ast
import re

def process_file(api_key, file, instructions):
    genai.configure(api_key=api_key)
    model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')

    try:
        df = pd.read_csv(file.name) if file.name.endswith('.csv') else pd.read_excel(file.name)
    except Exception as e:
        print(f"File Error: {str(e)}")
        return [None]*3

    # Enhanced prompt with strict syntax rules
    prompt = f"""Generate 3 Python code blocks for matplotlib visualizations with:
    1. Perfect Python syntax
    2. No markdown or incomplete lines
    3. Each block must start with:
       plt.figure(figsize=(16,9), dpi=120)
       plt.style.use('ggplot')
    4. Use ONLY df and plt variables
    5. End with plt.tight_layout()
    
    Columns: {list(df.columns)}
    Sample: {df.head(3).to_dict()}
    Instructions: {instructions or 'None'}
    
    Format EXACTLY as:
    # Visualization 1
    [code]
    
    # Visualization 2
    [code]
    
    # Visualization 3
    [code]
    """

    try:
        response = model.generate_content(prompt)
        code_blocks = re.split(r'# Visualization \d+', response.text)[1:4]
    except Exception as e:
        print(f"Gemini Error: {str(e)}")
        return [None]*3

    visualizations = []
    for i, block in enumerate(code_blocks, 1):
        buf = io.BytesIO()
        try:
            # Advanced code cleaning
            cleaned_code = '\n'.join(
                line.strip().replace("'", "").replace('"', '')  # Remove stray quotes
                for line in block.split('\n') 
                if line.strip() and 
                not any(c in line for c in ['`', '```', 'Annotation']) and
                re.match(r'^[a-zA-Z0-9_().=, <>:]+$', line)  # Basic syntax validation
            )
            
            # Add missing parentheses check
            cleaned_code = re.sub(r'plt.style.use\([\'"]ggplot$', 
                                'plt.style.use("ggplot")', cleaned_code)
            
            # Syntax validation
            ast.parse(cleaned_code)
            
            # Execute code
            exec_env = {'df': df, 'plt': plt}
            plt.figure(figsize=(16, 9), dpi=120)
            exec(cleaned_code, exec_env)
            
            plt.savefig(buf, format='png', bbox_inches='tight')
            plt.close()
            visualizations.append(Image.open(buf))
        except Exception as e:
            print(f"Visualization {i} Error: {str(e)}")
            print(f"Cleaned Code:\n{cleaned_code}")
            visualizations.append(None)

    return visualizations + [None]*(3-len(visualizations))

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Professional Data 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", 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
    )

demo.launch()