import gradio as gr
import pandas as pd
import json
from agents import analyze_data_with_agent 
import io
import asyncio
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def process_data_and_prompt(file, prompt):
    """Process uploaded file and prompt using the data analysis agent."""
    try:
        if not file:
            return "Please upload a data file.", None, None
            
        if not prompt or prompt.strip() == "":
            return "Please enter an analysis prompt.", None, None

        # Read the uploaded file
        if file.name.endswith('.csv'):
            df = pd.read_csv(file.name)
        elif file.name.endswith(('.xlsx', '.xls')):
            df = pd.read_excel(file.name)
        elif file.name.endswith('.json'):
            df = pd.read_json(file.name)
        else:
            return "Error: Unsupported file format. Please upload CSV, Excel, or JSON files.", None, None

        # Clean column names
        df.columns = [str(col).strip().lower().replace(' ', '_').replace('-', '_') for col in df.columns]
        
        # Show data preview
        # data_preview = f"""
        # <div class="data-section">
        #     <h3>Data Preview</h3>
        #     <p><strong>Shape:</strong> {df.shape[0]} rows × {df.shape[1]} columns</p>
        #     <p><strong>Columns:</strong> {', '.join(df.columns.tolist())}</p>
        #     {df.head().to_html(classes='table data-table', table_id='data-preview')}
        # </div>
        # """
        data_preview = f"""
<div></div>"""

        # Process with agent
        logger.info(f"Processing prompt: {prompt}")
        result = await analyze_data_with_agent(prompt, df)
        logger.info(f"Agent result type: {result.get('type')}")

        # Handle different result types
        if result["type"] == "error":
            error_html = f"""
            <div class="error-box">
                <h3>Error</h3>
                <p><strong>Message:</strong> {result['message']}</p>
                {f"<p><strong>Suggestions:</strong></p><ul>{''.join([f'<li>{s}</li>' for s in result.get('suggestions', [])])}</ul>" if result.get('suggestions') else ""}
            </div>
            """
            return data_preview + error_html, None, None
        
        elif result["type"] == "visualization":
            # Display the chart
            image_base64 = result.get("image")
            if image_base64:
                chart_html = f"""
                <div class="analysis-result">
                    <h3>Visualization Result</h3>
                    <p><strong>Chart Type:</strong> {result.get('chart_type', 'Unknown').title()}</p>
                    <div class="chart-container">
                        <img src="data:image/png;base64,{image_base64}" class="chart-image">
                    </div>
                    <p><em>{result.get('message', 'Visualization created successfully')}</em></p>
                </div>
                """
                return data_preview + chart_html, None, None
            else:
                return data_preview + "<p>Error: Could not generate visualization</p>", None, None
        
        elif result["type"] == "statistical":
            # Format statistical results
            stat_html = f"""
            <div class="analysis-result">
                <h3>Statistical Analysis Results</h3>
                <div class="stat-output-box">
                    {result.get('data', 'No statistical results available')}
                </div>
                <p><em>{result.get('message', 'Statistical analysis completed')}</em></p>
            </div>
            """
            return data_preview + stat_html, None, None
        
        elif result["type"] == "transformation":
            # Return transformed data
            transformed_df = result.get("dataframe")
            if transformed_df is not None:
                # Create CSV for download
                csv_buffer = io.StringIO()
                transformed_df.to_csv(csv_buffer, index=False)
                csv_data = csv_buffer.getvalue()
                
                # Create temporary file for download (Gradio handles temporary files for downloads)
                temp_file_name = "transformed_data.csv"
                with open(temp_file_name, 'w', encoding='utf-8') as f:
                    f.write(csv_data)
                
                transform_html = f"""
                <div class="analysis-result">
                    <h3>Data Transformation Results</h3>
                    <p><strong>Original Shape:</strong> {df.shape[0]} rows × {df.shape[1]} columns</p>
                    <p><strong>New Shape:</strong> {result.get('shape', 'Unknown')}</p>
                    <p><strong>New Columns:</strong> {', '.join(result.get('columns', []))}</p>
                    <div class="transformed-data-preview">
                        <h4>Preview of Transformed Data:</h4>
                        {result.get('preview', 'No preview available')}
                    </div>
                    <p><em>{result.get('message', 'Data transformation completed')}</em></p>
                    <p><strong>Download the transformed data using the button below.</strong></p>
                </div>
                """
                return data_preview + transform_html, temp_file_name, None
            else:
                return data_preview + "<p>Error: Could not retrieve transformed data</p>", None, None
        
        else:
            return data_preview + f"<p>Unknown result type: {result.get('type')}</p>", None, None

    except Exception as e:
        logger.error(f"Error processing data: {str(e)}")
        error_html = f"""
        <div class="error-box">
            <h3>Processing Error</h3>
            <p><strong>Error:</strong> {str(e)}</p>
            <p><strong>Please check:</strong></p>
            <ul>
                <li>File format is supported (CSV, Excel)</li>
                <li>File is not corrupted</li>
                <li>Prompt is clear and specific</li>
                <li>Ollama server is running</li>
            </ul>
        </div>
        """
        return error_html, None, None

def process_sync(file, prompt):
    """Synchronous wrapper for the async processing function."""
    try:
        # Check if an event loop is already running
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
        return loop.run_until_complete(process_data_and_prompt(file, prompt))
    except Exception as e:
        logger.error(f"Error in sync wrapper: {str(e)}")
        return f"Error: {str(e)}", None, None

def generate_preview(file):
    """Generate a preview of the uploaded file."""
    try:
        if not file:
            return "Please upload a data file to see preview."
            
        # Read the uploaded file
        if file.name.endswith('.csv'):
            df = pd.read_csv(file.name)
        elif file.name.endswith(('.xlsx', '.xls')):
            df = pd.read_excel(file.name)
        elif file.name.endswith('.json'):
            df = pd.read_json(file.name)
        else:
            return "Error: Unsupported file format. Please upload CSV, Excel, or JSON files."

        # Clean column names
        df.columns = [str(col).strip().lower().replace(' ', '_').replace('-', '_') for col in df.columns]
        
        # Show data preview
        data_preview = f"""
        <div class="data-section">
            <h3>📊 Data Preview</h3>
            <div class="data-stats">
                <span class="stat-badge">📏 {df.shape[0]} rows</span>
                <span class="stat-badge">📋 {df.shape[1]} columns</span>
            </div>
            <div class="columns-info">
                <strong>Columns:</strong> {', '.join(df.columns.tolist())}
            </div>
            <div class="table-container">
                {df.head(4).to_html(classes='table data-table', table_id='data-preview')}
            </div>
        </div>
        """
        return data_preview
    except Exception as e:
        logger.error(f"Error generating preview: {str(e)}")
        return f"<div class='error-box'>Error generating preview: {str(e)}</div>"

# Sample prompts for different analysis types
sample_prompts = {
    "Data Transformation": [
        "Filter data where [column] > 1000 ",
        "Group by [column] and calculate average [values]",
        "Create new columns based on existing ones",
        "Remove duplicates and sort by date",
    ],
    "Visualization": [
        "Create a bar chart showing the distribution of [categories]",
        "Generate a line plot of sales over time",
        "Make a scatter plot of [column1] vs [column2]",
        "Show a histogram of [column2]",
        "Create a pie chart of market share by region"
    ],
    "Statistical Analysis": [
        "Calculate correlation matrix for all numeric columns",
        "Perform descriptive statistics analysis",
    ]
}

# Create the Gradio interface
with gr.Blocks(
    title="Data Analysis Agent",
    theme=gr.themes.Soft(),
    css="""
    /* Main container */
    .gradio-container {
        max-width: 900px;
        margin: auto;
        padding: 20px;
    }

    /* Header styling */
    .main-header {
        text-align: center;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 30px;
        border-radius: 15px;
        margin-bottom: 30px;
        box-shadow: 0 8px 32px rgba(0,0,0,0.1);
    }
    
    .main-header h1 {
        margin: 0;
        font-size: 2.5em;
        font-weight: 600;
    }
    
    .main-header p {
        margin: 10px 0 0 0;
        font-size: 1.1em;
        opacity: 0.9;
    }

    /* Accordion styling */
    .gr-accordion {
        margin-bottom: 20px !important;
        border-radius: 12px !important;
        border: 1px solid var(--border-color-primary) !important;
        box-shadow: 0 2px 8px rgba(0,0,0,0.05) !important;
        overflow: hidden !important;
    }
    
    .gr-accordion-header {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
        color: white !important;
        padding: 15px 20px !important;
        font-weight: 600 !important;
        font-size: 1.1em !important;
        border: none !important;
        cursor: pointer !important;
        transition: all 0.3s ease !important;
    }
    
    .gr-accordion-header:hover {
        background: linear-gradient(135deg, #5a6fd8 0%, #6a4190 100%) !important;
        transform: translateY(-1px) !important;
    }
    
    .gr-accordion-content {
        background: var(--background-fill-secondary) !important;
        padding: 25px !important;
        border-top: 1px solid var(--border-color-primary) !important;
    }
    
    /* Special styling for example prompt accordions */
    .gr-accordion .gr-accordion {
        margin-bottom: 15px !important;
        border-radius: 8px !important;
        box-shadow: 0 1px 4px rgba(0,0,0,0.1) !important;
    }
    
    .gr-accordion .gr-accordion .gr-accordion-header {
        background: var(--color-accent-soft) !important;
        color: var(--text-color-body) !important;
        padding: 12px 16px !important;
        font-size: 1em !important;
        font-weight: 500 !important;
    }
    
    .gr-accordion .gr-accordion .gr-accordion-header:hover {
        background: var(--color-accent) !important;
        color: white !important;
        transform: none !important;
    }
    
    .gr-accordion .gr-accordion .gr-accordion-content {
        background: var(--background-fill-primary) !important;
        padding: 15px !important;
    }

    /* Section styling (keeping for compatibility) */
    .section {
        background: var(--background-fill-secondary);
        border-radius: 12px;
        padding: 25px;
        margin-bottom: 25px;
        border: 1px solid var(--border-color-primary);
        box-shadow: 0 2px 8px rgba(0,0,0,0.05);
    }
    
    .section h2 {
        margin: 0 0 20px 0;
        color: var(--text-color-body);
        font-size: 1.4em;
        font-weight: 600;
        display: flex;
        align-items: center;
        gap: 10px;
    }

    /* File upload styling */
    .upload-area {
        border: 2px dashed var(--border-color-accent);
        border-radius: 10px;
        padding: 20px;
        text-align: center;
        background: var(--background-fill-primary);
        transition: all 0.3s ease;
    }
    
    .upload-area:hover {
        border-color: var(--color-accent);
        background: var(--background-fill-hover);
    }

    /* Data preview styling */
    .data-section {
        background: var(--background-fill-primary);
        border-radius: 10px;
        padding: 20px;
        border: 1px solid var(--border-color-primary);
        margin: 15px 0;
    }
    
    .data-section h3 {
        margin: 0 0 15px 0;
        color: var(--text-color-body);
        font-size: 1.2em;
    }
    
    .data-stats {
        display: flex;
        gap: 10px;
        margin-bottom: 15px;
        flex-wrap: wrap;
    }
    
    .stat-badge {
        background: var(--color-accent-soft);
        color: var(--text-color-body);
        padding: 6px 12px;
        border-radius: 20px;
        font-size: 0.9em;
        font-weight: 500;
    }
    
    .columns-info {
        margin-bottom: 15px;
        padding: 10px;
        background: var(--background-fill-secondary);
        border-radius: 8px;
        font-size: 0.9em;
    }
    
    .table-container {
        overflow-x: auto;
        border-radius: 8px;
    }

    /* Table styling */
    .table {
        width: 100%;
        border-collapse: collapse;
        font-size: 0.85em;
        background: var(--background-fill-primary);
    }
    
    .table th {
        background: var(--background-fill-secondary);
        color: var(--text-color-body);
        font-weight: 600;
        padding: 12px 8px;
        border: 1px solid var(--border-color-primary);
        text-align: left;
    }
    
    .table td {
        padding: 10px 8px;
        border: 1px solid var(--border-color-primary);
        color: var(--text-color-body);
    }
    
    .table tr:nth-child(even) {
        background: var(--background-fill-hover);
    }

    /* Prompt examples styling */
    .prompt-examples {
        display: grid;
        gap: 15px;
        margin-top: 15px;
    }
    
    .prompt-category {
        background: var(--background-fill-primary);
        border-radius: 8px;
        padding: 15px;
        border: 1px solid var(--border-color-primary);
    }
    
    .prompt-category h4 {
        margin: 0 0 10px 0;
        color: var(--text-color-body);
        font-size: 1em;
    }
    
    .prompt-buttons {
        display: flex;
        flex-wrap: wrap;
        gap: 8px;
    }
    
    .prompt-btn {
        font-size: 0.8em !important;
        padding: 6px 12px !important;
        border-radius: 15px !important;
        background: var(--color-accent-soft) !important;
        color: var(--text-color-body) !important;
        border: 1px solid var(--border-color-accent) !important;
        cursor: pointer;
        transition: all 0.2s ease;
    }
    
    .prompt-btn:hover {
        background: var(--color-accent) !important;
        color: white !important;
    }

    /* Analysis results styling */
    .analysis-result {
        background: var(--background-fill-primary);
        border-radius: 10px;
        padding: 20px;
        margin: 15px 0;
        border: 1px solid var(--border-color-primary);
    }
    
    .analysis-result h3 {
        margin: 0 0 15px 0;
        color: var(--text-color-body);
    }

    /* Chart styling */
    .chart-container {
        text-align: center;
        margin: 20px 0;
        background: var(--background-fill-primary);
        padding: 15px;
        border-radius: 8px;
        border: 1px solid var(--border-color-primary);
    }
    
    .chart-image {
        max-width: 100%;
        height: auto;
        border-radius: 8px;
        box-shadow: 0 4px 12px rgba(0,0,0,0.1);
    }

    /* Error styling */
    .error-box {
        background: #fee;
        border: 1px solid #fcc;
        color: #c33;
        padding: 15px;
        border-radius: 8px;
        margin: 15px 0;
    }
    
    .error-box h3 {
        margin: 0 0 10px 0;
        color: #c33;
    }

    /* Button styling */
    .analyze-btn {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
        color: white !important;
        border: none !important;
        border-radius: 25px !important;
        padding: 15px 30px !important;
        font-size: 1.1em !important;
        font-weight: 600 !important;
        box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4) !important;
        transition: all 0.3s ease !important;
    }
    
    .analyze-btn:hover {
        transform: translateY(-2px) !important;
        box-shadow: 0 6px 20px rgba(102, 126, 234, 0.6) !important;
    }

    /* Responsive design */
    @media (max-width: 768px) {
        .gradio-container {
            padding: 10px;
        }
        
        .main-header h1 {
            font-size: 2em;
        }
        
        .section {
            padding: 15px;
        }
        
        .data-stats {
            flex-direction: column;
        }
        
        .prompt-buttons {
            flex-direction: column;
        }
    }
    """
) as demo:
    
    # Header
    gr.Markdown("""
    # 🤖 Data Analysis Agent
    
    Upload your data file and describe what analysis you want to perform. The AI agent will:
    - 📊 Create visualizations (charts, plots, graphs)
    - 🔢 Perform statistical analysis (correlations, tests, summaries)
    - 🔧 Transform your data (filter, aggregate, compute new columns)
    
    **Supported formats:** CSV, Excel (.xlsx, .xls)
    """)
    
    # Step 1: File Upload
    with gr.Accordion("📁 Step 1: Upload Your Data", open=True):
        file_input = gr.File(
            label="Choose your data file (CSV, Excel)",
            file_types=[".csv", ".xlsx", ".xls"],
            type="filepath"
        )
    
    # Step 2: Data Preview
    with gr.Accordion("👀 Step 2: Data Preview", open=True):
        preview_output = gr.HTML(value="<p style='text-align: center; color: #888; padding: 40px;'>Upload a file to see data preview</p>")
    
    # Step 3: Analysis Prompt
    with gr.Accordion("💬 Step 3: Describe Your Analysis", open=True):
        prompt_input = gr.Textbox(
            label="What would you like to analyze?",
            placeholder="e.g., 'Create a bar chart showing sales by category' or 'Calculate correlation between price and quantity'",
            lines=3
        )
        
        # Example prompts in separate collapsible sections
        gr.HTML('<h4 style="margin: 20px 0 10px 0;">💡 Need inspiration? Try these examples:</h4>')
        

        with gr.Accordion("🔧 Data Transformation Examples", open=False):
            for prompt in sample_prompts["Data Transformation"]:
                gr.Button(prompt, size="sm", elem_classes=["prompt-btn"]).click(
                    lambda p=prompt: p, inputs=[], outputs=prompt_input, queue=False
                )

        with gr.Accordion("📊 Visualization Examples", open=False):
            for prompt in sample_prompts["Visualization"]:
                gr.Button(prompt, size="sm", elem_classes=["prompt-btn"]).click(
                    lambda p=prompt: p, inputs=[], outputs=prompt_input, queue=False
                )
        
        with gr.Accordion("📈 Statistical Analysis Examples", open=False):
            for prompt in sample_prompts["Statistical Analysis"]:
                gr.Button(prompt, size="sm", elem_classes=["prompt-btn"]).click(
                    lambda p=prompt: p, inputs=[], outputs=prompt_input, queue=False
                )
        

    
    # Step 4: Analysis Button
    with gr.Accordion("🚀 Step 4: Run Analysis", open=True):
        submit_btn = gr.Button("🚀 Analyze Data", variant="primary", size="lg", elem_classes=["analyze-btn"])
    
    # Step 5: Results
    with gr.Accordion("📊 Step 5: Analysis Results", open=True):
        output = gr.HTML(value="<p style='text-align: center; color: #888; padding: 40px;'>Click 'Analyze Data' to see results here</p>")
    
    # Step 6: Downloads
    with gr.Accordion("📥 Step 6: Downloads", open=True):
        download_output = gr.File(label="Transformed Data (if applicable)", visible=True)
        gr.HTML("<p style='color: #666; font-size: 0.9em;'>Download will appear here for data transformation results</p>")
    
    # Event handlers
    file_input.change(
        fn=generate_preview,
        inputs=[file_input],
        outputs=[preview_output]
    )
    
    submit_btn.click(
        fn=process_sync,
        inputs=[file_input, prompt_input],
        outputs=[output, download_output],
        show_progress=True
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )