import os
import io
import sys
import gradio as gr
import pandas as pd
import matplotlib
# Use a non-interactive backend for Matplotlib to prevent cloud server crashes
matplotlib.use('Agg') 
import matplotlib.pyplot as plt
from google import genai

# 1. Initialize the Gemini client securely from the environment
api_key = os.environ.get('GEMINI_API_KEY')
client = genai.Client(api_key=api_key)

# 2. Define the core data and plotting logic
def analyze_and_plot(uploaded_file, user_question):
    if uploaded_file is None:
        # Return error message, hide chart, disable download button
        return "⚠️ Please upload a CSV file first.", None, gr.update(visible=False)
        
    try:
        # Clear any previous figures and delete old reports to keep clean states
        plt.close('all')
        if os.path.exists("analysis_report.txt"):
            os.remove("analysis_report.txt")
        if os.path.exists("output_chart.png"):
            os.remove("output_chart.png")
        
        # Read the uploaded CSV file using Pandas
        df = pd.read_csv(uploaded_file.name)
        
        # Gather basic details about the dataset to feed the AI context
        row_count = len(df)
        col_count = len(df.columns)
        column_names = ", ".join(df.columns.tolist())
        data_preview = df.head(3).to_string()
        data_types = df.dtypes.to_string()
        
        # System instructions forcing Gemini to return structured code and metrics
        system_context = (
            f"You are an expert Python Data Analyst. The user has uploaded a dataset.\n"
            f"Dataset Statistics:\n"
            f"- Total Rows: {row_count}\n"
            f"- Total Columns: {col_count}\n"
            f"- Column Names: {column_names}\n\n"
            f"Data Types:\n{data_types}\n\n"
            f"Preview:\n{data_preview}\n\n"
            f"INSTRUCTIONS:\n"
            f"1. Answer the user's question clearly.\n"
            f"2. If performing a regression or statistical test, you MUST explicitly state the key numerical metrics "
            f"(such as R-squared value, Slope/Coefficients, Intercept, or p-values) in your final written text response.\n"
            f"3. If a chart is requested, include a clean Python code block using 'matplotlib.pyplot' and 'sklearn'. "
            f"Assume the dataframe is loaded as 'df'.\n"
            f"4. Crucial: The LAST line of your chart code block must save the image exactly as: plt.savefig('output_chart.png')\n"
            f"5. Do not include plt.show(). Only save it to 'output_chart.png'."
        )
        
        # Send the metadata context and user query to Gemini
        response = client.models.generate_content(
            model='gemini-2.5-flash',
            contents=[system_context, f"User Question: {user_question}"]
        )
        
        text_response = response.text
        chart_path = None
        
        # Extract and execute the python code if Gemini generated a graph script
        if "```python" in text_response:
            try:
                # Isolate the code inside the markdown blocks
                code_block = text_response.split("```python")[1].split("```")[0]
                
                # Setup a string buffer to catch anything the executed script tries to print out
                old_stdout = sys.stdout
                redirected_output = sys.stdout = io.StringIO()
                
                # Create a safe execution environment containing the user's data
                local_vars = {'df': df, 'plt': plt}
                exec(code_block, {}, local_vars)
                
                # Restore normal system output
                sys.stdout = old_stdout
                printed_results = redirected_output.getvalue()
                
                # If the script printed extra statistical text, add it to the final display
                if printed_results:
                    text_response += f"\n\n### 🖥️ Calculated Execution Output:\n```text\n{printed_results}\n```"
                
                # Check if the chart image file was successfully written
                if os.path.exists("output_chart.png"):
                    chart_path = "output_chart.png"
            except Exception as e:
                text_response += f"\n\n*(Note: There was a background error rendering the chart: {str(e)})*"
        
        # Save the finalized text analysis into a downloadable file on the server
        with open("analysis_report.txt", "w", encoding="utf-8") as f:
            f.write(text_response)
            
        # Return text, chart path, and make the download button visible with the report file attached
        return text_response, chart_path, gr.update(value="analysis_report.txt", visible=True)
        
    except Exception as e:
        return f"❌ An error occurred processing the file: {str(e)}", None, gr.update(visible=False)

# 3. Design the custom layout using Gradio Blocks
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📊 Visual Cloud AI Data Analyst")
    gr.Markdown("Upload a CSV file, ask a question, and download your final text reports easily.")
    
    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(label="Step 1: Upload CSV Dataset", file_types=[".csv"])
            question_input = gr.Textbox(
                label="Step 2: Ask a question or request a model", 
                placeholder="e.g., 'Run a linear regression between ColumnA and ColumnB'..."
            )
            submit_btn = gr.Button("Analyze & Generate Chart", variant="primary")
            
        with gr.Column(scale=1):
            output_text = gr.Markdown(label="AI Written Insights")
            
            # The download button is hidden until an analysis runs successfully
            download_btn = gr.DownloadButton("📥 Download Text Report", visible=False)
            
            output_image = gr.Image(label="Generated Data Visualization", type="filepath")
            
    # Connect inputs to the execution function and distribute the 3 outputs
    submit_btn.click(
        fn=analyze_and_plot, 
        inputs=[file_input, question_input], 
        outputs=[output_text, output_image, download_btn]
    )

# 4. Launch the public cloud server
if __name__ == "__main__":
    demo.launch()