import os
import pandas as pd
import matplotlib.pyplot as plt
from langchain_ollama import OllamaLLM
import gradio as gr
from fpdf import FPDF

# Initialize Homelander AI
ollama = OllamaLLM(base_url='http://localhost:11434', model="Homelander")

# Directory for uploads
UPLOAD_FOLDER = './uploads'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# Function to upload and read the file
def upload_file(file):
    """Reads and returns a Pandas DataFrame from a file."""
    try:
        # Check if the input is a file-like object or string
        if isinstance(file, str):  # If it's already a file path
            file_path = file
        else:  # Save the file-like object to the server
            file_path = os.path.join(UPLOAD_FOLDER, file.name)
            with open(file_path, "wb") as f:
                f.write(file.read())

        # Load the file into a Pandas DataFrame
        if file_path.endswith('.csv'):
            return pd.read_csv(file_path), file_path
        elif file_path.endswith('.xlsx'):
            return pd.read_excel(file_path), file_path
        else:
            raise ValueError("Unsupported file format. Please upload a CSV or Excel file.")
    except Exception as e:
        raise ValueError(f"Error loading file: {e}")

def generate_plots(df):
    """
    Generates meaningfull plots for all numeric and categorical columns in the dataset.
    - Numeric columns: Histograms.
    - Categorical columns: Bar charts (if unique values < 20).
    """

    try:
        num_cols = df.select_dtypes(include=['number']).columns.tolist()
        cat_cols = df.select_dtypes(include=['object', 'category']).columns.tolist()

        # Define plot size based on the number of subplots
        total_plots = len(num_cols) + len(cat_cols)
        if total_plots == 0:
            return "No suitable columns to plot."

        plt.figure(figsize=(12, 6 * total_plots))

        # Plot numeric columns
        for i, col in enumerate(num_cols):
            plt.subplot(total_plots, 1, i + 1)
            df[col].plot(kind='hist', bins=30, color='lightblue', edgecolor='black')
            plt.title(f"Histogram of '{col}'")
            plt.xlabel(col)
            plt.ylabel("Frequency")
            plt.grid(True)

        # Plot categorical columns (only if unique values < 20 for readability)
        for j, col in enumerate(cat_cols):
            if df[col].nunique() < 20:
                plt.subplot(total_plots, 1, len(num_cols) + j + 1)
                df[col].value_counts().plot(kind='bar', color='lightcoral')
                plt.title(f"Bar Plot of '{col}'")
                plt.xlabel(col)
                plt.ylabel("Count")
                plt.xticks(rotation=45)
                plt.grid(axis='y')

        plt.tight_layout()
        plot_path = os.path.join(UPLOAD_FOLDER, "full_dataset_plots.png")
        plt.savefig(plot_path, dpi=300)
        plt.close()
        
        return plot_path

    except Exception as e:
        return None


# Function to query Homelander AI for insights
def query_homelander(file_path, question):
    """
    Query Homelander to generate insights based on the user's question.
    """
    prompt = f"""
    Analyze the file '{file_path}' and provide detailed insights for the following question:
    {question}

    Ensure the response contains insights only, without Python code.
    """
    try:
        response = ollama(prompt)
        return response.encode('utf-8').decode('utf-8')  # Ensure proper encoding
    except Exception as e:
        raise ValueError(f"Error querying Homelander: {e}")


# Main analysis function
def analyze_file_and_plot(file, question):
    """
    Main function to handle file upload, analysis, and graph generation.
    """
    try:
        # Load the uploaded file
        df, file_path = upload_file(file)
        
        # Query Homelander for insights
        insights = query_homelander(file_path, question)

        # Generate and display a basic plot (customize as needed)
        plt.figure(figsize=(10, 6))
        df.iloc[:, 0].value_counts().plot(kind='bar', color='skyblue')
        plt.title("Sample Bar Plot of the First Column")
        plt.xlabel("Categories")
        plt.ylabel("Counts")
        plot_path = os.path.join(UPLOAD_FOLDER, "plot.png")
        plt.savefig(plot_path)
        plt.close()

        return insights, plot_path
    except Exception as e:
        return str(e), None


def chat_with_homelander(file=None, question=""):
    try:
        if not question.strip():
            return "Please provide a question.", None, None  # Adjusted for 3 outputs

        plot_path = None
        pdf_path = None

        if file is not None:
            df, file_path = upload_file(file)
            plot_path = generate_plots(df)
            insights = query_homelander(file_path, question)
            pdf_path = save_insights_as_pdf(insights)  # Generate PDF from insights
        else:
            insights = ollama.invoke(question).strip()  # Correct way to query the AI

        # Handle potential Unicode issues
        insights = insights.encode('utf-8', 'ignore').decode('utf-8')

        return insights, plot_path, pdf_path

    except Exception as e:
        return f"Error: {str(e)}", None, None




def save_insights_as_pdf(insights_text):
    """Generate a PDF from the insights and return the file path."""
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    pdf.multi_cell(0, 10, insights_text)

    pdf_path = os.path.join(UPLOAD_FOLDER, "insights_report.pdf")
    pdf.output(pdf_path)
    return pdf_path

interface = gr.Interface(
    fn=chat_with_homelander,
    inputs=[
        gr.File(label="Upload CSV or Excel File (Optional)"),
        gr.Textbox(lines=2, placeholder="Enter your question", label="Question"),
    ],
    outputs=[
        gr.Textbox(label="Insights"),
        gr.Image(type="filepath", label="Generated Plot"),
        gr.File(label="Download Insights PDF"),
    ],
    title="🤖Homelander AI Assistant",
    description="Ask Homelander any question or upload a file to analyze and gain insights. "
                "You can download the insights as a PDF.",
)


# Launch the interface
interface.launch()