import gradio as gr
import pandas as pd
import chromadb
from chromadb.config import Settings
from chromadb.utils import embedding_functions

# Initialize Chroma client and embedding function
EMBEDDING_MODEL = 'all-MiniLM-L6-v2'
embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=EMBEDDING_MODEL)
chroma_client = chromadb.Client(Settings(persist_directory="./chroma_db"))

# Global variables to store the current state
current_file_or_link = None
current_selected_column = None
current_collection = None
current_df = None

def load_data(file_or_link):
    try:
        if file_or_link.endswith(".csv"):
            df = pd.read_csv(file_or_link)
        elif "docs.google.com" in file_or_link:
            sheet_id = file_or_link.split('/d/')[1].split('/')[0]
            csv_url = f"https://docs.google.com/spreadsheets/d/{sheet_id}/export?format=csv"
            df = pd.read_csv(csv_url)
        else:
            raise ValueError("Unsupported file format or URL.")
        return df
    except Exception as e:
        print(f"Error in load_data: {str(e)}")
        return pd.DataFrame()

def embed_selected_column(df, selected_column, collection_name="dynamic_collection"):
    try:
        collection = chroma_client.get_or_create_collection(name=collection_name, embedding_function=embedding_function)

        # Check if collection is empty
        if collection.count() == 0:
            print("Embedding and storing data. This may take a while...")

            # Ensure the selected column exists in the DataFrame
            if selected_column not in df.columns:
                raise ValueError(f"Column '{selected_column}' does not exist in the data.")

            # Embed and store the selected column and all other columns as metadata
            collection.add(
                documents=df[selected_column].tolist(),  # Only embed the selected column
                metadatas=df.to_dict('records'),  # Store all columns as metadata
                ids=[str(i) for i in range(len(df))]
            )
            print("Data embedded and stored.")
        else:
            print("Data already embedded and stored.")

        return collection
    except Exception as e:
        print(f"Error in embed_selected_column: {str(e)}")
        return None

def search_similar_queries(query, collection, top_k=5, original_df=None):
    try:
        results = collection.query(
            query_texts=[query],
            n_results=top_k,
            include=["metadatas", "distances"]
        )

        # Create a DataFrame to hold the results
        formatted_results = []
        for i in range(len(results['ids'][0])):
            metadata = results['metadatas'][0][i]
            result = {key: metadata.get(key, 'N/A') for key in original_df.columns}  # Use original DataFrame columns order
            result['similarity'] = 1 - results['distances'][0][i]  # Convert distance to similarity
            formatted_results.append(result)

        # Convert results to DataFrame
        results_df = pd.DataFrame(formatted_results)

        return results_df
    except Exception as e:
        print(f"Error in search_similar_queries: {str(e)}")
        return pd.DataFrame()


def generate_embedding(file_or_link, selected_column):
    global current_file_or_link, current_selected_column, current_collection, current_df

    # Check if the data source or selected column has changed
    if file_or_link != current_file_or_link or selected_column != current_selected_column:
        df = load_data(file_or_link)
        if df.empty:
            return "Error: Failed to load data.", gr.DataFrame()

        # If data source or column changed, create a new collection
        collection_name = f"collection_{hash(file_or_link)}_{hash(selected_column)}"
        collection = embed_selected_column(df, selected_column, collection_name)

        if collection is None:
            return "Error: Failed to embed data.", gr.DataFrame()

        current_file_or_link = file_or_link
        current_selected_column = selected_column
        current_collection = collection
        current_df = df

        return "New embedding generated successfully.", gr.DataFrame(df.head())
    else:
        return "Using existing embedding.", gr.DataFrame(current_df.head() if current_df is not None else pd.DataFrame())

def search_queries(query, top_k):
    global current_file_or_link, current_selected_column, current_collection, current_df

    if current_collection is None:
        return "Error: Please generate embedding first.", gr.DataFrame()

    results_df = search_similar_queries(query, current_collection, top_k, original_df=current_df)

    if results_df.empty:
        return "No results found.", gr.DataFrame()

    return "Search completed.", results_df

def setup_interface():
    with gr.Blocks() as iface:
        gr.Markdown("# Semantic Search Interface")

        with gr.Row():
            file_or_link = gr.Textbox(label="CSV File Path or Google Sheets Link")
            selected_column = gr.Textbox(label="Enter Column Name to Embed")

        with gr.Row():
            generate_button = gr.Button("Generate Embedding")
            embed_status = gr.Textbox(label="Embedding Status", interactive=False)

        with gr.Row():
            query = gr.Textbox(label="Enter your query")
            top_k = gr.Slider(minimum=1, maximum=10, step=1, label="Top-K results", value=5)

        search_button = gr.Button("Search")
        search_status = gr.Textbox(label="Search Status", interactive=False)
        results = gr.DataFrame(label="Results", interactive=False)

        generate_button.click(
            fn=generate_embedding,
            inputs=[file_or_link, selected_column],
            outputs=[embed_status, results]
        )

        search_button.click(
            fn=search_queries,
            inputs=[query, top_k],
            outputs=[search_status, results]
        )

    iface.launch()

setup_interface()