app.py

import os
import streamlit as st
import pandas as pd
import pickle
import base64
from io import BytesIO, StringIO
import sys
import operator
from typing import Literal, Sequence, TypedDict, Annotated, List, Dict, Tuple
import tempfile
import shutil
import plotly.io as pio
import io
import re
import json
import openai
# from fpdf import FPDF
import base64
from datetime import datetime
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from PIL import Image as PILImage

# Import LangChain and LangGraph components
from langchain_core.messages import AIMessage, ToolMessage, HumanMessage, BaseMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_experimental.utilities import PythonREPL
from langgraph.prebuilt import ToolInvocation, ToolExecutor
from langchain_core.tools import tool
from langgraph.prebuilt import InjectedState
from langgraph.graph import StateGraph, END
from reportlab.platypus import PageBreak


# Initialize session state for AI provider settings
if 'ai_provider' not in st.session_state:
    st.session_state.ai_provider = "openai"
    
if 'api_key' not in st.session_state:
    st.session_state.api_key = ""
    
if 'selected_model' not in st.session_state:
    st.session_state.selected_model = "gpt-4"
    
# Define model options for each provider
OPENAI_MODELS = ["gpt-4", "gpt-4-turbo", "gpt-4-mini", "gpt-3.5-turbo"]
GROQ_MODELS = ["llama3.3-70b-versatile", "gemma2-9b-it", "llama-3-8b-8192"]

# Create temporary directory for file storage
if 'temp_dir' not in st.session_state:
    st.session_state.temp_dir = tempfile.mkdtemp()
    st.session_state.images_dir = os.path.join(st.session_state.temp_dir, "images/plotly_figures/pickle")
    os.makedirs(st.session_state.images_dir, exist_ok=True)
    print(f"Created temporary directory: {st.session_state.temp_dir}")
    print(f"Created images directory: {st.session_state.images_dir}")

# Define the system prompt
SYSTEM_PROMPT = """## Role
You are a professional data scientist helping a non-technical user understand, analyze, and visualize their data.

## Capabilities
1. **Execute python code** using the `complete_python_task` tool.

## Goals
1. Understand the user's objectives clearly.
2. Take the user on a data analysis journey, iterating to find the best way to visualize or analyse their data to solve their problems.
3. Investigate if the goal is achievable by running Python code via the `python_code` field.
4. Gain input from the user at every step to ensure the analysis is on the right track and to understand business nuances.

## Code Guidelines
- **ALL INPUT DATA IS LOADED ALREADY**, so use the provided variable names to access the data.
- **VARIABLES PERSIST BETWEEN RUNS**, so reuse previously defined variables if needed.
- **TO SEE CODE OUTPUT**, use `print()` statements. You won't be able to see outputs of `pd.head()`, `pd.describe()` etc. otherwise.
- **ONLY USE THE FOLLOWING LIBRARIES**:
  - `pandas`
  - `sklearn` (including all major ML models)
  - `plotly`
  - `numpy`
  
All these libraries are already imported for you.

## Machine Learning Guidelines
- For regression tasks:
  - Linear Regression: `LinearRegression`
  - Logistic Regression: `LogisticRegression`
  - Ridge Regression: `Ridge`
  - Lasso Regression: `Lasso`
  - Random Forest Regression: `RandomForestRegressor`
  
- For classification tasks:
  - Logistic Regression: `LogisticRegression`
  - Decision Trees: `DecisionTreeClassifier`
  - Random Forests: `RandomForestClassifier`
  - Support Vector Machines: `SVC`
  - K-Nearest Neighbors: `KNeighborsClassifier`
  - Naive Bayes: `GaussianNB`
  
- For clustering:
  - K-Means: `KMeans`
  - DBSCAN: `DBSCAN`
  
- For dimensionality reduction:
  - PCA: `PCA`
  
- Always preprocess data appropriately:
  - Scale numerical features with `StandardScaler` or `MinMaxScaler`
  - Encode categorical variables with `OneHotEncoder` when needed
  - Handle missing values with `SimpleImputer`
  
- Always split data into training and testing sets using `train_test_split`
- Evaluate models using appropriate metrics:
  - For regression: `mean_squared_error`, `mean_absolute_error`, `r2_score`
  - For classification: `accuracy_score`, `confusion_matrix`, `classification_report`
  - For clustering: `silhouette_score`
  
- Consider using `cross_val_score` for more robust evaluation
- Visualize ML results with plotly when possible

## Plotting Guidelines
- Always use the `plotly` library for plotting.
- Store all plotly figures inside a `plotly_figures` list, they will be saved automatically.
- Do not try and show the plots inline with `fig.show()`.
"""

# Define the State class
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    input_data: Annotated[List[Dict], operator.add]
    intermediate_outputs: Annotated[List[dict], operator.add]
    current_variables: dict
    output_image_paths: Annotated[List[str], operator.add]

# Initialize session state variables
if 'in_memory_datasets' not in st.session_state:
    st.session_state.in_memory_datasets = {}

if 'persistent_vars' not in st.session_state:
    st.session_state.persistent_vars = {}

if 'dataset_metadata_list' not in st.session_state:
    st.session_state.dataset_metadata_list = []

if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

if 'dashboard_plots' not in st.session_state:
    st.session_state.dashboard_plots = [None, None, None, None]

if 'columns' not in st.session_state:
    st.session_state.columns = ["No columns available"]

if 'custom_plots_to_save' not in st.session_state:
    st.session_state.custom_plots_to_save = {}

# Set up the tools
repl = PythonREPL()
plotly_saving_code = """import pickle

import uuid
import os
for figure in plotly_figures:
    pickle_filename = f"{images_dir}/{uuid.uuid4()}.pickle"
    with open(pickle_filename, 'wb') as f:
        pickle.dump(figure, f)
"""

@tool
def complete_python_task(
    graph_state: Annotated[dict, InjectedState],
    thought: str,
    python_code: str
) -> Tuple[str, dict]:
    """Execute Python code for data analysis and visualization."""

    current_variables = graph_state.get("current_variables", {})

    # Load datasets from in-memory storage
    for input_dataset in graph_state.get("input_data", []):
        var_name = input_dataset.get("variable_name")
        if var_name and var_name not in current_variables and var_name in st.session_state.in_memory_datasets:
            print(f"Loading {var_name} from in-memory storage")
            current_variables[var_name] = st.session_state.in_memory_datasets[var_name]
    current_image_pickle_files = os.listdir(st.session_state.images_dir)

    try:
        # Capture stdout
        old_stdout = sys.stdout
        sys.stdout = StringIO()
       
        # Execute the code and capture the result
        exec_globals = globals().copy()
        exec_globals.update(st.session_state.persistent_vars)
        exec_globals.update(current_variables)
        
        # Add scikit-learn modules to execution environment
        import sklearn
        import numpy as np
        
        # Import scikit-learn components
        from sklearn.linear_model import LinearRegression, LogisticRegression, Ridge, Lasso # type: ignore
        from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor, GradientBoostingClassifier # type: ignore
        from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
        from sklearn.svm import SVC, SVR
        from sklearn.naive_bayes import GaussianNB
        from sklearn.decomposition import PCA
        from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
        from sklearn.cluster import KMeans, DBSCAN
        from sklearn.preprocessing import StandardScaler, MinMaxScaler, OneHotEncoder
        from sklearn.model_selection import train_test_split, cross_val_score, GridSearchCV
        from sklearn.metrics import (
            accuracy_score, confusion_matrix, classification_report,
            mean_squared_error, r2_score, mean_absolute_error, silhouette_score
        )
        from sklearn.pipeline import Pipeline
        from sklearn.impute import SimpleImputer
        
        # Update execution globals with all ML components
        exec_globals.update({
            "plotly_figures": [], 
            "images_dir": st.session_state.images_dir,
            "np": np,
            # Linear models
            "LinearRegression": LinearRegression,
            "LogisticRegression": LogisticRegression,
            "Ridge": Ridge,
            "Lasso": Lasso,
            # Tree-based models
            "DecisionTreeClassifier": DecisionTreeClassifier,
            "DecisionTreeRegressor": DecisionTreeRegressor,
            "RandomForestClassifier": RandomForestClassifier,
            "RandomForestRegressor": RandomForestRegressor,
            "GradientBoostingClassifier": GradientBoostingClassifier,
            # SVM models
            "SVC": SVC,
            "SVR": SVR,
            # Other models
            "GaussianNB": GaussianNB,
            "PCA": PCA,
            "KNeighborsClassifier": KNeighborsClassifier,
            "KNeighborsRegressor": KNeighborsRegressor,
            "KMeans": KMeans,
            "DBSCAN": DBSCAN,
            # Preprocessing
            "StandardScaler": StandardScaler,
            "MinMaxScaler": MinMaxScaler,
            "OneHotEncoder": OneHotEncoder,
            "SimpleImputer": SimpleImputer,
            # Model selection and evaluation
            "train_test_split": train_test_split,
            "cross_val_score": cross_val_score,
            "GridSearchCV": GridSearchCV,
            "accuracy_score": accuracy_score,
            "confusion_matrix": confusion_matrix,
            "classification_report": classification_report,
            "mean_squared_error": mean_squared_error,
            "r2_score": r2_score,
            "mean_absolute_error": mean_absolute_error,
            "silhouette_score": silhouette_score,
            # Pipeline
            "Pipeline": Pipeline
        })
        
        exec(python_code, exec_globals)

        st.session_state.persistent_vars.update({k: v for k, v in exec_globals.items() if k not in globals()})

        # Get the captured stdout
        output = sys.stdout.getvalue()

        # Restore stdout
        sys.stdout = old_stdout

        updated_state = {
            "intermediate_outputs": [{"thought": thought, "code": python_code, "output": output}],
            "current_variables": st.session_state.persistent_vars
        }

        if 'plotly_figures' in exec_globals and exec_globals['plotly_figures']:
            exec(plotly_saving_code, exec_globals)
           
            # Check if any images were created
            new_image_folder_contents = os.listdir(st.session_state.images_dir)
            new_image_files = [file for file in new_image_folder_contents if file not in current_image_pickle_files]
           
            if new_image_files:
                updated_state["output_image_paths"] = new_image_files
            st.session_state.persistent_vars["plotly_figures"] = []
        return output, updated_state

    except Exception as e:
        sys.stdout = old_stdout  # Restore stdout in case of error
        print(f"Error in complete_python_task: {str(e)}")
        return str(e), {"intermediate_outputs": [{"thought": thought, "code": python_code, "output": str(e)}]}

# Function to initialize the LLM based on selected provider and model
def initialize_llm():
    api_key = st.session_state.api_key
    model = st.session_state.selected_model
    
    if not api_key:
        return None
    
    try:
        if st.session_state.ai_provider == "openai":
            os.environ["OPENAI_API_KEY"] = api_key
            return ChatOpenAI(model=model, temperature=0)
        elif st.session_state.ai_provider == "groq":
            os.environ["GROQ_API_KEY"] = api_key
            # For Groq, set the base URL and use the model
            from langchain_groq import ChatGroq
            return ChatGroq(model=model, temperature=0)
    except Exception as e:
        print(f"Error initializing LLM: {str(e)}")
        return None

# Set up the tools
tools = [complete_python_task]
tool_executor = ToolExecutor(tools)

# Load the prompt template
chat_template = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    ("placeholder", "{messages}"),
])

def create_data_summary(state: AgentState) -> str:
    summary = ""
    variables = []
   
    # Add sample data for each dataset
    for d in state.get("input_data", []):
        var_name = d.get("variable_name")
        if var_name:
           
            variables.append(var_name)
            summary += f"\n\nVariable: {var_name}\n"
            summary += f"Description: {d.get('data_description', 'No description')}\n"

            # Add sample data if available
            if var_name in st.session_state.in_memory_datasets:
                df = st.session_state.in_memory_datasets[var_name]
                summary += "\nSample Data (first 5 rows):\n"
                summary += df.head(5).to_string()

    if "current_variables" in state:
        remaining_variables = [v for v in state["current_variables"] if v not in variables and not v.startswith("_")]
       
        for v in remaining_variables:
           
            var_value = state["current_variables"].get(v)

            if isinstance(var_value, pd.DataFrame):
                summary += f"\n\nVariable: {v} (DataFrame with shape {var_value.shape})"
            else:
                summary += f"\n\nVariable: {v}"
    return summary

def route_to_tools(state: AgentState) -> Literal["tools", "__end__"]:
    """Determine if we should route to tools or end the chain"""
    if messages := state.get("messages", []):
        ai_message = messages[-1]
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")

    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        return "tools"
   
    return "__end__"

def call_model(state: AgentState):
    """Call the LLM to get a response"""
    current_data_template = """The following data is available:\n{data_summary}"""
    current_data_message = HumanMessage(
        content=current_data_template.format(data_summary=create_data_summary(state))
    )
    messages = [current_data_message] + state["messages"]
    
    # Get the initialized LLM
    llm = initialize_llm()
    if llm is None:
        return {"messages": [AIMessage(content="Please configure a valid API key and model in the settings tab.")]}
    
    # Create the model with bound tools
    model = llm.bind_tools(tools)
    model = chat_template | model
    
    llm_outputs = model.invoke({"messages": messages})
    return {"messages": [llm_outputs], "intermediate_outputs": [current_data_message.content]}

def call_tools(state: AgentState):
    """Execute tools called by the LLM"""
    last_message = state["messages"][-1]
    tool_invocations = []

    if isinstance(last_message, AIMessage) and hasattr(last_message, 'tool_calls'):
        tool_invocations = [
            ToolInvocation(
                tool=tool_call["name"],
                tool_input={**tool_call["args"], "graph_state": state}
            ) for tool_call in last_message.tool_calls
        ]
    responses = tool_executor.batch(tool_invocations, return_exceptions=True)

    tool_messages = []
    state_updates = {}

    for tc, response in zip(last_message.tool_calls, responses):
        if isinstance(response, Exception):
            print(f"Exception in tool execution: {str(response)}")
            tool_messages.append(ToolMessage(                
                content=f"Error: {str(response)}",
                name=tc["name"],
                tool_call_id=tc["id"]
            ))
            continue

        message, updates = response
        tool_messages.append(ToolMessage(
            content=str(message),
            name=tc["name"],
            tool_call_id=tc["id"]
        ))

        # Merge updates instead of overwriting
        for key, value in updates.items():
            if key in state_updates:
                if isinstance(value, list) and isinstance(state_updates[key], list):
                    state_updates[key].extend(value)
                elif isinstance(value, dict) and isinstance(state_updates[key], dict):
                    state_updates[key].update(value)
                else:
                    state_updates[key] = value
            else:
                state_updates[key] = value

    if 'messages' not in state_updates:
        state_updates["messages"] = []

    state_updates["messages"] = tool_messages
    return state_updates

# Set up the graph
workflow = StateGraph(AgentState)
workflow.add_node("agent", call_model)
workflow.add_node("tools", call_tools)
workflow.add_conditional_edges(
    "agent",
    route_to_tools,
    {
        "tools": "tools",
        "__end__": END
    }
)
workflow.add_edge("tools", "agent")
workflow.set_entry_point("agent")

chain = workflow.compile()

def process_file_upload(files):
    """Process uploaded files and return dataframe previews and column names"""
    st.session_state.in_memory_datasets = {}  # Clear previous datasets
    st.session_state.dataset_metadata_list = []  # Clear previous metadata
    st.session_state.persistent_vars.clear()  # Clear persistent variables for new session

    if not files:
        return "No files uploaded.", [], ["No columns available"]

    results = []
    all_columns = []  # Track all columns from all datasets

    for file in files:
        try:
            # Use file object directly
            if file.name.endswith('.csv'):
                df = pd.read_csv(file)
            elif file.name.endswith(('.xls', '.xlsx')):
                df = pd.read_excel(file)
            else:
                results.append(f"Unsupported file format: {file.name}. Please upload CSV or Excel files.")
                continue

            var_name = file.name.split('.')[0].replace('-', '_').replace(' ', '_').lower()
            st.session_state.in_memory_datasets[var_name] = df

            # Collect all columns
            all_columns.extend(df.columns.tolist())

            # Create dataset metadata
            dataset_metadata = {
                "variable_name": var_name,
                "data_path": "in_memory",
                "data_description": f"Dataset containing {df.shape[0]} rows and {df.shape[1]} columns. Columns: {', '.join(df.columns.tolist())}",
                "original_filename": file.name
            }

            st.session_state.dataset_metadata_list.append(dataset_metadata)

            # Return preview of the dataset
            preview = f"### Dataset: {file.name}\nVariable name: `{var_name}`\n\n"
            preview += df.head(10).to_markdown()
            results.append(preview)
            print(f"Successfully processed {file.name}")

        except Exception as e:
            print(f"Error processing {file.name}: {str(e)}")
            results.append(f"Error processing {file.name}: {str(e)}")

    # Get unique columns
    unique_columns = []
    seen = set()

    for col in all_columns:
        if col not in seen:
            seen.add(col)
            unique_columns.append(col)

    if not unique_columns:
        unique_columns = ["No columns available"]

    print(f"Found {len(unique_columns)} unique columns across datasets")
    return "\n\n".join(results), st.session_state.dataset_metadata_list, unique_columns

def get_columns():
    """Directly gets columns from in-memory datasets"""
    all_columns = []

    for var_name, df in st.session_state.in_memory_datasets.items():
        if isinstance(df, pd.DataFrame):
            all_columns.extend(df.columns.tolist())

    # Remove duplicates while preserving order
    unique_columns = []
    seen = set()
   
    for col in all_columns:
        if col not in seen:
            seen.add(col)
            unique_columns.append(col)

    if not unique_columns:
        unique_columns = ["No columns available"]

    print(f"Populating dropdowns with {len(unique_columns)} columns")
    return unique_columns

# === FUNCTIONS ===
import openai
import pandas as pd
import json
import re

def standard_clean(df):
    df.columns = [re.sub(r'\W+', '_', col).strip().lower() for col in df.columns]
    df.drop_duplicates(inplace=True)
    df.dropna(axis=1, how='all', inplace=True)
    df.dropna(axis=0, how='all', inplace=True)
    for col in df.select_dtypes(include='object').columns:
        df[col] = df[col].astype(str).str.strip()
    return df

def query_openai(prompt):
    try:
        # Use the configured API key and model from session state
        api_key = st.session_state.api_key
        model = st.session_state.selected_model
        
        if st.session_state.ai_provider == "openai":
            client = openai.OpenAI(api_key=api_key)
            response = client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7
            )
            return response.choices[0].message.content
        elif st.session_state.ai_provider == "groq":
            from groq import Groq
            client = Groq(api_key=api_key)
            response = client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7
            )
            return response.choices[0].message.content
    except Exception as e:
        print(f"API Error: {e}")
        return "{}"

def llm_suggest_cleaning(df):
    sample = df.head(10).to_csv(index=False)
    prompt = f"""
You are a professional data wrangler. Below is a sample of a messy dataset.

Return a Python dictionary with the following keys:

1. rename_columns – fix unclear or inconsistent column names
2. convert_types – correct datatypes: int, float, str, or date
3. fill_missing – use 'mean', 'median', 'mode', or a constant like 'Unknown' or 0
4. value_map – map inconsistent values (e.g., yes/Yes/Y → Yes)

Do not drop any rows or columns. Your output must be a valid Python dict.

Example:
{{
  "rename_columns": {{"dob": "date_of_birth"}},
  "convert_types": {{"age": "int", "salary": "float", "signup_date": "date"}},
  "fill_missing": {{"gender": "mode", "salary": -1}},
  "value_map": {{
    "gender": {{"M": "Male", "F": "Female"}},
    "subscribed": {{"Y": "Yes", "N": "No"}}
  }}
}}
Apart from these mentioned steps, study the data and also do whatever things are good and needed for that particular dataset and do the cleaning.
Sample data:
{sample}
"""
    raw_response = query_openai(prompt)
    try:
        suggestions = eval(raw_response)
        return suggestions
    except:
        print("Could not parse suggestions.")
        return {
            "rename_columns": {},
            "convert_types": {},
            "fill_missing": {},
            "value_map": {}
        }

def apply_suggestions(df, suggestions):
    df.rename(columns=suggestions.get("rename_columns", {}), inplace=True)

    for col, dtype in suggestions.get("convert_types", {}).items():
        if col not in df.columns:
            continue
        try:
            if dtype == "int":
                df[col] = pd.to_numeric(df[col], errors='coerce').astype("Int64")
            elif dtype == "float":
                df[col] = pd.to_numeric(df[col], errors='coerce')
            elif dtype == "str":
                df[col] = df[col].astype(str)
            elif dtype == "date":
                df[col] = pd.to_datetime(df[col], errors='coerce')
        except:
            print(f"Failed to convert {col} to {dtype}")

    for col, method in suggestions.get("fill_missing", {}).items():
        if col not in df.columns:
            continue
        try:
            if method == "mean":
                df[col].fillna(df[col].mean(), inplace=True)
            elif method == "median":
                df[col].fillna(df[col].median(), inplace=True)
            elif method == "mode":
                df[col].fillna(df[col].mode().iloc[0], inplace=True)
            elif isinstance(method, str):
                df[col].fillna(method, inplace=True)
        except:
            print(f"Could not fill missing values for {col}")

    for col, mapping in suggestions.get("value_map", {}).items():
        if col in df.columns:
            try:
                df[col] = df[col].replace(mapping)
            except:
                print(f"Could not map values in {col}")

    return df

def capture_dashboard_screenshot():
    """Capture the entire dashboard as a single image"""
    try:
        # Create a figure that combines all dashboard plots
        import plotly.graph_objects as go
        from plotly.subplots import make_subplots
       
        # Create a 2x2 subplot
        fig = make_subplots(rows=2, cols=2,
                           subplot_titles=["Visualization 1", "Visualization 2",
                                          "Visualization 3", "Visualization 4"])

        # Add each plot from the dashboard to the combined figure
        for i, plot in enumerate(st.session_state.dashboard_plots):
            if plot is not None:
                row = (i // 2) + 1
                col = (i % 2) + 1

                # Extract traces from the original figure and add to our subplot
                for trace in plot.data:
                    fig.add_trace(trace, row=row, col=col)
               
                # Copy layout properties for each subplot
                for axis_type in ['xaxis', 'yaxis']:
                    axis_name = f"{axis_type}{i+1 if i > 0 else ''}"
                    subplot_name = f"{axis_type}{row}{col}"

                    # Copy axis properties if they exist
                    if hasattr(plot.layout, axis_name):
                        axis_props = getattr(plot.layout, axis_name)
                        fig.update_layout({subplot_name: axis_props})
       
        # Update layout for better appearance
        fig.update_layout(
            height=800,
            width=1000,
            title_text="Dashboard Overview",
            showlegend=False,
        )

        # Save to a temporary file
        dashboard_path = f"{st.session_state.temp_dir}/dashboard_combined.png"
        fig.write_image(dashboard_path, scale=2)  # Higher scale for better resolution
        return dashboard_path

    except Exception as e:
        import traceback
        print(f"Error capturing dashboard: {str(e)}")
        print(traceback.format_exc())
        return None
    
def generate_enhanced_pdf_report():
    """Generate an enhanced PDF report with proper handling of base64 image data"""
    try:
        # Create a buffer for the PDF
        buffer = io.BytesIO()
        
        # Create the PDF document
        doc = SimpleDocTemplate(buffer, pagesize=letter,
                               leftMargin=36, rightMargin=36,
                               topMargin=36, bottomMargin=36)

        # Create custom styles with better formatting
        styles = getSampleStyleSheet()
        
        # Add custom styles with improved formatting
        styles.add(ParagraphStyle(
            name='ReportTitle',
            parent=styles['Heading1'],
            fontSize=24,
            alignment=1,  # Centered
            spaceAfter=20,
            textColor='#2C3E50'  # Dark blue color
        ))
        
        styles.add(ParagraphStyle(
            name='SectionHeader',
            parent=styles['Heading2'],
            fontSize=16,
            spaceBefore=15,
            spaceAfter=10,
            textColor='#2C3E50',
            borderWidth=1,
            borderColor='#95A5A6',
            borderPadding=5,
            borderRadius=5
        ))
        
        styles.add(ParagraphStyle(
            name='SubHeader',
            parent=styles['Heading3'],
            fontSize=14,
            spaceBefore=10,
            spaceAfter=8,
            textColor='#34495E',
            fontWeight='bold'
        ))
        styles.add(ParagraphStyle(
                name='UserMessage',
                parent=styles['Normal'],
                fontSize=11,
                leftIndent=10,
                spaceBefore=8,
                spaceAfter=4
            ))
            
        styles.add(ParagraphStyle(
            name='AssistantMessage',
            parent=styles['Normal'],
            fontSize=11,
            leftIndent=10,
            spaceBefore=4,
            spaceAfter=12,
            textColor='#2980B9'
        ))
        
        styles.add(ParagraphStyle(
            name='Timestamp',
            parent=styles['Italic'],
            fontSize=10,
            textColor='#7F8C8D',
            alignment=2  # Right aligned
        ))

        # Create the document content
        elements = []
        
        # Add title
        elements.append(Paragraph('Data Analysis Report', styles['ReportTitle']))
        
        # Add timestamp
        elements.append(Paragraph(f'Generated on: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}',
                              styles['Timestamp']))
        elements.append(Spacer(1, 0.5*inch))
        
        # Add conversation history with better formatting
        elements.append(Paragraph('Analysis Conversation History', styles['SectionHeader']))
        
        if st.session_state.chat_history:
            for i, (user_msg, assistant_msg) in enumerate(st.session_state.chat_history):
                # Format user message with proper styling
                elements.append(Paragraph(f'<b>You:</b>', styles['SubHeader']))
                user_msg_formatted = user_msg.replace('\n', '<br/>')
                elements.append(Paragraph(user_msg_formatted, styles['UserMessage']))
                
                # Process assistant message to handle visualization
                # Look for markdown image syntax with base64 data
                base64_pattern = r'!\[Visualization\]\(data:image\/png;base64,([^\)]+)\)'
                
                # Check if the message contains visualizations
                if '### Visualizations' in assistant_msg or re.search(base64_pattern, assistant_msg):
                    # Split the message at the Visualizations header if it exists
                    if '### Visualizations' in assistant_msg:
                        parts = assistant_msg.split('### Visualizations', 1)
                        text_part = parts[0]
                        viz_part = "### Visualizations" + parts[1] if len(parts) > 1 else ""
                    else:
                        # If no header but still has visualization
                        match = re.search(base64_pattern, assistant_msg)
                        text_part = assistant_msg[:match.start()]
                        viz_part = assistant_msg[match.start():]
                    
                    # Format the text part
                    elements.append(Paragraph(f'<b>Assistant:</b>', styles['SubHeader']))
                    text_part = text_part.replace('\n', '<br/>')
                    elements.append(Paragraph(text_part, styles['AssistantMessage']))
                    
                    # Process visualizations
                    matches = re.findall(base64_pattern, viz_part)
                    for j, base64_data in enumerate(matches):
                        try:
                            # Decode the base64 image
                            image_data = base64.b64decode(base64_data)
                            
                            # Create a temporary file for the image
                            temp_img_path = f"{st.session_state.temp_dir}/chat_viz_{i}_{j}.png"
                            
                            with open(temp_img_path, 'wb') as f:
                                f.write(image_data)
                            
                            # Add the image to the PDF
                            elements.append(Paragraph(f'<b>Visualization:</b>', styles['SubHeader']))
                            elements.append(Spacer(1, 0.1*inch))
                            img = Image(temp_img_path, width=6*inch, height=4*inch)
                            elements.append(img)
                            elements.append(Spacer(1, 0.2*inch))
                        except Exception as e:
                            print(f"Error processing base64 image: {str(e)}")
                            elements.append(Paragraph(f"[Error displaying visualization: {str(e)}]", 
                                                    styles['Normal']))
                else:
                    # No visualizations, just format the text
                    elements.append(Paragraph(f'<b>Assistant:</b>', styles['SubHeader']))
                    assistant_msg_formatted = assistant_msg.replace('\n', '<br/>')
                    if len(assistant_msg_formatted) > 1500:
                        assistant_msg_formatted = assistant_msg_formatted[:1500] + '...'
                    elements.append(Paragraph(assistant_msg_formatted, styles['AssistantMessage']))
                
                elements.append(Spacer(1, 0.2*inch))
        else:
            elements.append(Paragraph('No conversation history available.', styles['Normal']))
        
        # Force a page break before the dashboard
        elements.append(PageBreak())
        
        # Add dashboard section header
        elements.append(Paragraph('Dashboard Overview', styles['SectionHeader']))
        elements.append(Spacer(1, 0.2*inch))
        
        # Capture the dashboard as a single image
        dashboard_img_path = capture_dashboard_screenshot()
        
        if dashboard_img_path:
            # Calculate available width
            available_width = doc.width
            
            # Create PIL image to get dimensions
            pil_img = PILImage.open(dashboard_img_path)
            img_width, img_height = pil_img.size
            
            # Calculate scaling factor to fit within page width
            scale_factor = available_width / img_width
            
            # Calculate new height based on aspect ratio
            new_height = img_height * scale_factor
            
            # Add the image with scaled dimensions
            img = Image(dashboard_img_path, width=available_width, height=new_height)
            elements.append(img)
        else:
            # Fallback: Add individual plots if combined dashboard fails
            plot_count = 0
            for i, plot in enumerate(st.session_state.dashboard_plots):
                if plot is not None:
                    plot_count += 1
                    
                    # Convert plotly figure to image
                    img_bytes = io.BytesIO()
                    plot.write_image(img_bytes, format='png', width=500, height=300)
                    img_bytes.seek(0)
                    
                    # Create a temporary file for the image
                    temp_img_path = f"{st.session_state.temp_dir}/plot_{i}.png"
                    
                    with open(temp_img_path, 'wb') as f:
                        f.write(img_bytes.getvalue())
                    
                    # Add to PDF with appropriate caption and formatting
                    elements.append(Paragraph(f'Dashboard Visualization {i+1}', styles['SubHeader']))
                    elements.append(Spacer(1, 0.1*inch))
                    
                    # Add the image with proper scaling
                    img = Image(temp_img_path, width=6.5*inch, height=4*inch)
                    elements.append(img)
                    elements.append(Spacer(1, 0.3*inch))
            
            if plot_count == 0:
                elements.append(Paragraph('No visualizations have been added to the dashboard.',
                                        styles['Normal']))
        
        # Build the PDF with improved formatting
        doc.build(elements)
        
        # Get the value of the buffer
        pdf_value = buffer.getvalue()
        buffer.close()
        
        return pdf_value
        
    except Exception as e:
        import traceback
        print(f"Error generating enhanced PDF report: {str(e)}")
        print(traceback.format_exc())
        return None

def chat_with_workflow(message, history, dataset_info):
    """Send user query to the workflow and get response"""
    
    if not dataset_info:
        return "Please upload at least one dataset before asking questions."

    # Check if we have a valid API key and model
    if not st.session_state.api_key:
        return "Please set up your API key and model in the Settings tab before chatting."

    print(f"Chat with workflow called with {len(dataset_info)} datasets")

    try:
        # Extract chat history for context (last 3 exchanges)
        max_history = 3
        previous_messages = []
        
        if history:
            start_idx = max(0, len(history) - max_history)
            recent_history = history[start_idx:]
            
            for exchange in recent_history:
                if exchange[0]:  # User message
                    previous_messages.append(HumanMessage(content=exchange[0]))
                if exchange[1]:  # AI response
                    previous_messages.append(AIMessage(content=exchange[1]))

        # Initialize the workflow state
        state = AgentState(
            messages=previous_messages + [HumanMessage(content=message)],
            input_data=dataset_info,
            intermediate_outputs=[],
            current_variables=st.session_state.persistent_vars,
            output_image_paths=[]
        )

        # Execute the workflow
        print("Executing workflow...")
        result = chain.invoke(state)
        print("Workflow execution completed")

        # Extract messages from the result
        messages = result["messages"]

        # Format the response - only get the latest response
        response = ""
        if messages:
            latest_message = messages[-1]  # Get only the last message
            if hasattr(latest_message, "content"):
                content = latest_message.content
                
                # Clean up the response
                # Remove any instances where the user's message is repeated
                if message in content:
                    content = content.split(message)[-1].strip()
                
                # Remove any chat history markers
                content_lines = content.split('\n')
                filtered_lines = [line for line in content_lines 
                                if not line.strip().startswith(("You:", "User:", "Human:", "Assistant:"))]
                content = '\n'.join(filtered_lines)
                
                response = content.strip() + "\n\n"

        # Handle visualizations
        if "output_image_paths" in result and result["output_image_paths"]:
            response += "### Visualizations\n\n"
            for img_path in result["output_image_paths"]:
                try:
                    full_path = os.path.join(st.session_state.images_dir, img_path)
                    with open(full_path, 'rb') as f:
                        fig = pickle.load(f)

                    # Convert plotly figure to image
                    img_bytes = BytesIO()
                    fig.update_layout(width=800, height=500)
                    pio.write_image(fig, img_bytes, format='png')
                    img_bytes.seek(0)

                    # Convert to base64 for markdown image
                    b64_img = base64.b64encode(img_bytes.read()).decode()
                    response += f"![Visualization](data:image/png;base64,{b64_img})\n\n"
                except Exception as e:                    
                    response += f"Error loading visualization: {str(e)}\n\n"
                    
        return response

    except Exception as e:
        import traceback
        print(f"Error in chat_with_workflow: {str(e)}")
        print(traceback.format_exc())
        return f"Error executing workflow: {str(e)}"

def auto_generate_dashboard(dataset_info):
    """Generate an automatic dashboard with four plots"""

    if not dataset_info:
        return "Please upload a dataset first.", [None, None, None, None]

    prompt = """
    You are a data visualization expert. Given a dataset, identify the top 4 most insightful plots using statistical reasoning or patterns (correlation, distribution, trends).

    Use plotly and store the plots in a list named plotly_figures.

    Include multivariate plots using color/size/facets when helpful.
    """

    state = AgentState(
        messages=[HumanMessage(content=prompt)],
        input_data=dataset_info,
        intermediate_outputs=[],
        current_variables=st.session_state.persistent_vars,
        output_image_paths=[]
    )

    result = chain.invoke(state)
    figures = []

    if "output_image_paths" in result:
        for img_path in result["output_image_paths"][:4]:
            try:
                full_path = os.path.join(st.session_state.images_dir, img_path)
                with open(full_path, 'rb') as f:
                    fig = pickle.load(f)
                    figures.append(fig)
            except Exception as e:
                print(f"Error loading figure: {e}")

    while len(figures) < 4:
        figures.append(None)

    st.session_state.dashboard_plots = figures
    return "Dashboard generated!", figures

def generate_custom_plots_with_llm(dataset_info, x_col, y_col, facet_col):
    """Generate custom plots based on user-selected columns"""

    if not dataset_info or not x_col or not y_col:
        return [None, None, None]

    prompt = f"""
    You are a data visualization expert.

    Create 3 insightful visualizations using Plotly based on:

    - X-axis: {x_col}
    - Y-axis: {y_col}
    - Facet (optional): {facet_col if facet_col != 'None' else 'None'}

    Try to find interesting relationships, trends, or clusters using appropriate chart types.

    Use `plotly_figures` list and avoid using fig.show().
    """

    state = AgentState(
        messages=[HumanMessage(content=prompt)],
        input_data=dataset_info,
        intermediate_outputs=[],
        current_variables=st.session_state.persistent_vars,
        output_image_paths=[]
    )

    result = chain.invoke(state)
    figures = []

    if "output_image_paths" in result:
        for img_path in result["output_image_paths"][:3]:
            try:
                full_path = os.path.join(st.session_state.images_dir, img_path)
                with open(full_path, 'rb') as f:
                    fig = pickle.load(f)
                    figures.append(fig)
            except Exception as e:
                print(f"Error loading figure: {e}")
               
    while len(figures) < 3:
        figures.append(None)
    return figures

def remove_plot(index):
    """Remove a plot from the dashboard"""
    if 0 <= index < len(st.session_state.dashboard_plots):
        st.session_state.dashboard_plots[index] = None

def respond(message):
    """Handle chat message response"""
    if not st.session_state.dataset_metadata_list:
        bot_message = "Please upload at least one dataset before asking questions."
    else:
        bot_message = chat_with_workflow(message, st.session_state.chat_history, st.session_state.dataset_metadata_list)

    st.session_state.chat_history.append((message, bot_message))
    st.rerun()

def save_plot_to_dashboard(plot_index):
    """Callback for the Add Plot button"""
    for i in range(len(st.session_state.dashboard_plots)):
        if st.session_state.dashboard_plots[i] is None:
            # Found an empty slot
            st.session_state.dashboard_plots[i] = st.session_state.custom_plots_to_save[plot_index]
            return

# New function to check if settings are valid
def is_settings_valid():
    """Check if API key and model are configured"""
    return st.session_state.api_key != ""

# Streamlit UI with left panel settings
st.set_page_config(page_title="QueryMind", layout="wide")

# Create side panel for settings
with st.sidebar:
    st.header("Settings")
    st.info("⚠️ Configure your API settings before uploading data")
    
    # AI Provider selection
    provider = st.radio("Select AI Provider", 
                         options=["OpenAI", "Groq"],
                         index=0 if st.session_state.ai_provider == "openai" else 1,
                         horizontal=True)
    
    # Update session state based on selection
    st.session_state.ai_provider = provider.lower()
    
    # API Key input
    api_key = st.text_input("Enter API Key", 
                           value=st.session_state.api_key, 
                           type="password",
                           help="Your API key for the selected provider")
    
    # Display different model options based on provider
    if st.session_state.ai_provider == "openai":
        model_options = OPENAI_MODELS
        model_help = "GPT-4 provides the best results but is slower. GPT-3.5-Turbo is faster but less capable."
    else:  # groq
        model_options = GROQ_MODELS
        model_help = "Llama 3.3 70B is most capable. Gemma 2 9B offers good balance. Llama 3 8B is fastest."
    
    # Model selection
    selected_model = st.selectbox("Select Model", 
                                 options=model_options,
                                 index=model_options.index(st.session_state.selected_model) if st.session_state.selected_model in model_options else 0,
                                 help=model_help)
    
    # Save button
    if st.button("Save Settings"):
        st.session_state.api_key = api_key
        st.session_state.selected_model = selected_model
        
        # Test the API key and model
        try:
            # Initialize LLM using the provided settings
            test_llm = initialize_llm()
            if test_llm:
                st.success(f"✅ Successfully configured {provider} with model: {selected_model}")
            else:
                st.error("Failed to initialize the AI provider. Please check your API key and model selection.")
        except Exception as e:
            st.error(f"Error testing settings: {str(e)}")
    
    # Display current settings
    st.subheader("Current Settings")
    settings_info = f"""
    - **Provider**: {st.session_state.ai_provider.upper()}
    - **Model**: {st.session_state.selected_model}
    - **API Key**: {'✅ Set' if st.session_state.api_key else '❌ Not Set'}
    """
    st.markdown(settings_info)
    
    # Provider-specific information
    if st.session_state.ai_provider == "openai":
        with st.expander("OpenAI Models Info"):
            st.info("""
            - **GPT-4**: Most powerful model, best for complex analysis and detailed explanations
            - **GPT-4-Turbo**: Faster than GPT-4 with similar capabilities
            - **GPT-4-Mini**: Economical option with good performance for standard tasks
            - **GPT-3.5-Turbo**: Fastest option, suitable for basic analysis and visualization
            """)
    else:
        with st.expander("Groq Models Info"):
            st.info("""
            - **llama3.3-70b-versatile**: Most powerful model for comprehensive analysis
            - **gemma2-9b-it**: Good balance of speed and capabilities
            - **llama-3-8b-8192**: Fastest option for basic analysis tasks
            """)
    
    # Integration instructions
    with st.expander("How to get API Keys"):
        if st.session_state.ai_provider == "openai":
            st.markdown("""
            ### Getting an OpenAI API Key
            
            1. Go to [OpenAI's platform](https://platform.openai.com)
            2. Sign up or log in to your account
            3. Navigate to the API section
            4. Create a new API key
            5. Copy the key and paste it above
            
            Note: OpenAI API usage incurs charges based on tokens used.
            """)
        else:
            st.markdown("""
            ### Getting a Groq API Key
            
            1. Go to [Groq's website](https://console.groq.com/keys)
            2. Sign up or log in to your account
            3. Navigate to API Keys section
            4. Create a new API key
            5. Copy the key and paste it above
            
            Note: Check Groq's pricing page for current rates.
            """)

# Main content
st.title("QureyMind 🧠 - Data Assistant")
st.markdown("Upload your datasets, ask questions, and generate visualizations to gain insights.")

# Check if settings are valid before showing tabs
if not is_settings_valid():
    st.warning("⚠️ Please configure your API key in the sidebar settings panel before proceeding.")
    st.stop()

# Create main tabs - only if settings are valid
tab1, tab2, tab3, tab4, tab5 = st.tabs(["Upload Datasets", "Data Cleaning", "Chat with AI Assistant", "Auto Dashboard Generator", "Generate Report"])

with tab1:
    st.header("Upload Datasets")
    uploaded_files = st.file_uploader("Upload CSV or Excel Files",
                                    accept_multiple_files=True,
                                    type=['csv', 'xlsx', 'xls'])

    if uploaded_files and st.button("Process Uploaded Files"):
        with st.spinner("Processing files..."):
            preview, metadata_list, columns = process_file_upload(uploaded_files)
            st.session_state.columns = columns
           
            # Display basic information about processed files
            st.success(f"✅ Successfully processed {len(uploaded_files)} file(s)")
           
            # Show detailed preview for each dataset
            st.subheader("Dataset Previews")
           
            for dataset_name, df in st.session_state.in_memory_datasets.items():
                with st.expander(f"Preview: {dataset_name}"):
                    # Display dataset info
                    st.write(f"**Rows:** {df.shape[0]} | **Columns:** {df.shape[1]}")
                   
                    # Display column information
                    col_info = pd.DataFrame({
                        'Column Name': df.columns,
                        'Data Type': df.dtypes.astype(str),
                        'Non-Null Count': df.count().values,
                        'Sample Values': [', '.join(df[col].dropna().astype(str).head(3).tolist()) for col in df.columns]
                    })
                   
                    # Show column information in a compact table
                    st.write("**Column Information:**")
                    st.dataframe(col_info, use_container_width=True)
                   
                    # Show actual data preview
                    st.write("**Data Preview (First 10 rows):**")
                    st.dataframe(df.head(10), use_container_width=True)
           
            # Provide hint for the next steps
            st.info("👆 Click on the dataset names above to see detailed previews. Then proceed to the Data Cleaning tab to clean your data or Chat with AI Assistant to analyze it.")

with tab2:
    st.header("Data Cleaning")

    if 'cleaning_done' not in st.session_state:
        st.session_state.cleaning_done = False

    if 'cleaned_datasets' not in st.session_state:
        st.session_state.cleaned_datasets = {}

    if 'cleaning_summaries' not in st.session_state:
        st.session_state.cleaning_summaries = {}

    if st.session_state.get("in_memory_datasets"):
        if not st.session_state.cleaning_done:
            if st.button("Run Data Cleaning"):
                with st.spinner("Running LLM-assisted cleaning..."):
                    for name, df in st.session_state.in_memory_datasets.items():
                        raw_df = df.copy()
                        df_std = standard_clean(raw_df.copy())
                        suggestions = llm_suggest_cleaning(df_std.copy())
                        df_clean = apply_suggestions(df_std.copy(), suggestions)
                        st.session_state.cleaned_datasets[name] = df_clean
                        st.session_state.cleaning_summaries[name] = suggestions
                    st.session_state.cleaning_done = True
                    st.rerun()
            else:
                st.info("Click Run Data Cleaning to clean your datasets using the LLM.")
        else:
            for name, df_clean in st.session_state.cleaned_datasets.items():
                raw_df = st.session_state.in_memory_datasets[name]

                st.subheader(f"Dataset: {name}")
                col1, col2 = st.columns(2)

                with col1:
                    st.markdown("Original Data (First 5 Rows)")
                    st.dataframe(raw_df.head())

                with col2:
                    st.markdown("Cleaned Data (First 5 Rows)")
                    st.dataframe(df_clean.head())

                st.markdown("Summary of Cleaning Actions")
                suggestions = st.session_state.cleaning_summaries[name]
                summary_text = ""

                if suggestions:
                    for key, value in suggestions.items():
                        summary_text += f"**{key}**: {json.dumps(value, indent=2)}\n\n"
                    st.markdown(summary_text)

                st.markdown("Refine the Cleaning (Natural Language Instructions)")
                user_input = st.text_input("Example: Convert 'dob' to datetime and fill missing with '2000-01-01'",
                                           key=f"user_input_{name}")

                if f'corrections_{name}' not in st.session_state:
                    st.session_state[f'corrections_{name}'] = []

                if st.button("Apply Correction", key=f'apply_correction_{name}'):
                    if user_input.strip():
                        correction_prompt = f"""
You are a data cleaning expert. Below is a previously cleaned dataset with these actions:

{summary_text}

The user now wants the following additional instruction:
\"{user_input.strip()}\"

Write only the Python code that modifies the pandas DataFrame `df` accordingly.
Do not include explanations or markdown.
"""
                        correction_code = query_openai(correction_prompt)

                        try:
                            df = st.session_state.cleaned_datasets[name].copy()
                            local_vars = {"df": df}
                            exec(correction_code, {}, local_vars)
                            df_updated = local_vars["df"]

                            st.session_state.cleaned_datasets[name] = df_updated
                            st.session_state[f'corrections_{name}'].append((user_input, correction_code))
                            st.success("Correction applied.")
                            st.rerun()

                        except Exception as e:
                            st.error(f"Failed to apply correction: {str(e)}")

                if st.session_state[f'corrections_{name}']:
                    st.markdown("Applied Corrections")
                    for i, (msg, code) in enumerate(st.session_state[f'corrections_{name}']):
                        st.markdown(f"**Instruction:** {msg}")
                        st.code(code, language='python')

            col1, col2 = st.columns([1, 2])
            with col1:
                if st.button("Reset Cleaning and Re-run"):
                    st.session_state.cleaning_done = False
                    st.rerun()

            with col2:
                if st.button("Finalize and Proceed to Visualizations"):
                    st.session_state.cleaning_finalized = True
                    st.rerun()
    else:
        st.info("Please upload and process datasets first.")

with tab3:
    st.header("Chat with AI Assistant")
   
    st.markdown("""
    ## Example Questions
    - "What analysis can you perform on this dataset?"
    - "Show me basic statistics for all columns"
    - "Create a correlation heatmap"
    - "Plot the distribution of a specific column"
    - "What is the relationship between two columns?"
    """)

    # Display chat history
    for exchange in st.session_state.chat_history:
        with st.chat_message("user"):
            st.write(exchange[0])
        with st.chat_message("assistant"):
            st.write(exchange[1])

    # Chat input
    if prompt := st.chat_input("Your question"):
        with st.spinner("Thinking..."):
            respond(prompt)

with tab4:
    st.header("Auto Dashboard Generator")

    # Dashboard controls
    dashboard_title = st.text_input("Dashboard Title", placeholder="Enter your dashboard title")
 
    col1, col2 = st.columns(2)

    with col1:
        if st.button("Generate Suggested Dashboard (Auto)"):
            with st.spinner("Generating dashboard..."):
                message, figures = auto_generate_dashboard(st.session_state.dataset_metadata_list)
                st.success(message)

    with col2:
        if st.button("Refresh Column Options"):
            st.session_state.columns = get_columns()
            st.rerun()

    # Dashboard display
    st.subheader("Dashboard")

    # Row 1
    col1, col2 = st.columns(2)

    with col1:
        if st.session_state.dashboard_plots[0]:
            st.plotly_chart(st.session_state.dashboard_plots[0], use_container_width=True)
            if st.button("Remove Plot 1"):
                remove_plot(0)
                st.rerun()

    with col2:
        if st.session_state.dashboard_plots[1]:
            st.plotly_chart(st.session_state.dashboard_plots[1], use_container_width=True)
            if st.button("Remove Plot 2"):
                remove_plot(1)
                st.rerun()

    # Row 2
    col3, col4 = st.columns(2)

    with col3:
        if st.session_state.dashboard_plots[2]:
            st.plotly_chart(st.session_state.dashboard_plots[2], use_container_width=True)
            if st.button("Remove Plot 3"):
                remove_plot(2)
                st.rerun()

    with col4:
        if st.session_state.dashboard_plots[3]:
            st.plotly_chart(st.session_state.dashboard_plots[3], use_container_width=True)
            if st.button("Remove Plot 4"):
                remove_plot(3)
                st.rerun()

    # Custom plot generator
    st.subheader("Custom Plot Generator")

    # Column selection
    col1, col2, col3 = st.columns(3)

    with col1:
        x_axis = st.selectbox("X-axis Column", options=st.session_state.columns)

    with col2:
        y_axis = st.selectbox("Y-axis Column", options=st.session_state.columns)

    with col3:
        facet = st.selectbox("Facet (optional)", options=["None"] + st.session_state.columns)
 
    if st.button("Generate Custom Visualizations"):
        with st.spinner("Generating custom visualizations..."):
            custom_plots = generate_custom_plots_with_llm(st.session_state.dataset_metadata_list, x_axis, y_axis, facet)
            # Store plots in session state
            for i, plot in enumerate(custom_plots):
                if plot:
                    st.session_state.custom_plots_to_save[i] = plot

            # Display custom plots with add buttons
            for i, plot in enumerate(custom_plots):
                if plot:
                    st.plotly_chart(plot, use_container_width=True)
                    st.button(
                        f"Add Plot {i+1} to Dashboard",
                        key=f"add_plot_{i}",
                        on_click=save_plot_to_dashboard,
                        args=(i,)
                    )

with tab5:
    st.header("Generate Analysis Report")

    st.markdown("""
    Generate a PDF report containing:
    - Dashboard visualizations
    - Chat conversation history
    """)

    report_title = st.text_input("Report Title (Optional)", "Data Analysis Report")
   
    if st.button("Generate PDF Report"):
        with st.spinner("Generating report..."):
            pdf_data = generate_enhanced_pdf_report()
            if pdf_data:
                # Create download button for PDF
                b64_pdf = base64.b64encode(pdf_data).decode('utf-8')
                # Create download link
                pdf_download_link = f'<a href="data:application/pdf;base64,{b64_pdf}" download="data_analysis_report.pdf">Download PDF Report</a>'
                st.markdown("### Your report is ready!")
                st.markdown(pdf_download_link, unsafe_allow_html=True)
                # Preview option (simplified)
                with st.expander("Preview Report"):
                    st.warning("PDF preview is not available in Streamlit, please download the report to view it.")
            else:
                st.error("Failed to generate the report. Please try again.")

# Cleanup on app exit
def cleanup():
    try:
        shutil.rmtree(st.session_state.temp_dir)
        print(f"Cleaned up temporary directory: {st.session_state.temp_dir}")
    except Exception as e:
        print(f"Error cleaning up: {e}")

import atexit
atexit.register(cleanup)