import os
from typing import List, Optional, Union

import gradio as gr
import pandas as pd
from dotenv import load_dotenv
from pandas import DataFrame
from smolagents import CodeAgent, LiteLLMModel, tool

# Load environment variables
load_dotenv()

def create_agent():
    """Create a CodeAgent instance with GPT-4 backend."""
    model = LiteLLMModel(model_id="gpt-4o-mini")
    
    @tool
    def read_csv(filepath: str) -> DataFrame:
        """
        Read a CSV file and return a pandas DataFrame.
        
        Args:
            filepath: Path to the CSV file
        """
        return pd.read_csv(filepath)
    
    @tool
    def read_excel(filepath: str) -> DataFrame:
        """
        Read an Excel file and return a pandas DataFrame.
        
        Args:
            filepath: Path to the Excel file
        """
        return pd.read_excel(filepath)

    agent = CodeAgent(
        tools=[read_csv, read_excel],
        model=model,
        additional_authorized_imports=[
            "pandas",
            "numpy",
            "matplotlib",
            "seaborn",
            "plotly",
            "sklearn",
            "scipy",
        ],
        max_steps=5,
        verbosity_level=1
    )
    return agent

def process_request(
    files: Union[str, List[str]],
    user_query: str,
    api_key: str = "",
    temperature: float = 0.7,
    history: Optional[List[tuple]] = None
) -> tuple:
    """
    Process user request with uploaded files and query.
    
    Args:
        files: Path or list of paths to uploaded files
        user_query: Natural language query from user
        api_key: Optional API key for GPT-4
        temperature: Model temperature
        history: Chat history
        
    Returns:
        Tuple of (output, error, new_history)
    """
    if api_key:
        os.environ["OPENAI_API_KEY"] = api_key
        
    try:
        # Create agent instance
        agent = create_agent()
        
        # Build context from files
        file_context = ""
        if isinstance(files, str):
            files = [files]
            
        for file in files:
            filename = os.path.basename(file)
            file_context += f"File uploaded: {filename}\n"
            
        # Build complete prompt
        prompt = f"""
        {file_context}
        
        User request: {user_query}
        
        Please analyze the data and provide:
        1. Code to perform the analysis
        2. Explanation of approach
        3. Visualizations if relevant
        4. Key insights and findings
        """
        
        # Execute agent
        result = agent.run(prompt)
        
        # Update history
        new_history = history or []
        new_history.append((user_query, result))
        
        return result, None, new_history
        
    except Exception as e:
        return None, str(e), history

# Create Gradio interface
def create_interface():
    """Create Gradio interface for the AI coding assistant."""
    
    with gr.Blocks(title="AI Coding Assistant") as interface:
        gr.Markdown("""
        # AI Coding Assistant
        Upload data files and ask questions in natural language to get code, analysis and visualizations.
        """)
        
        with gr.Row():
            with gr.Column():
                files = gr.File(
                    label="Upload Data Files",
                    file_types=[".csv", ".xlsx", ".xls"],
                    multiple=True
                )
                query = gr.Textbox(
                    label="What would you like to analyze?",
                    placeholder="e.g., Create a scatter plot comparing column A vs B"
                )
                api_key = gr.Textbox(
                    label="API Key (Optional)",
                    placeholder="Your OpenAI API key",
                    type="password"
                )
                temperature = gr.Slider(
                    label="Temperature",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.7,
                    step=0.1
                )
                submit = gr.Button("Analyze")
                
            with gr.Column():
                output = gr.Markdown(label="Output")
                error = gr.Markdown(label="Errors")
        
        # Hidden state for chat history
        history = gr.State([])
        
        # Handle submissions
        submit.click(
            process_request,
            inputs=[files, query, api_key, temperature, history],
            outputs=[output, error, history]
        )
        
        # Add examples
        gr.Examples(
            examples=[
                [
                    None,
                    "Create a scatter plot showing the relationship between column A and B, with a trend line",
                ],
                [
                    None,
                    "Calculate summary statistics and identify any outliers in the numerical columns",
                ],
                [
                    None,
                    "Perform clustering analysis on the data and visualize the clusters",
                ],
            ],
            inputs=[files, query],
        )
        
    return interface

if __name__ == "__main__":
    interface = create_interface()
    interface.launch()