import os
import json
import pandas as pd
from dotenv import load_dotenv
from os.path import join, dirname
from llm import embedding_llm
from langchain_community.vectorstores import SupabaseVectorStore
from langchain_core.documents import Document
from supabase.client import Client, create_client
from langchain.tools.retriever import create_retriever_tool


dotenv_path = join(dirname(__file__), '.env')
load_dotenv(dotenv_path)

def get_supabase_client():
    supabase: Client = create_client(
        os.environ.get("SUPABASE_URL"), 
        os.environ.get("SUPABASE_SERVICE_KEY")
    )
    return supabase


def __upload_gaia_data(supabase_client):
    supabase_client = get_supabase_client()
    df = pd.read_csv('supabase_docs.csv')[['content', 'metadata']]
    df['metadata'] = df['metadata'].apply(lambda x: x.replace("'", '"')).apply(json.loads)
    docs_dict = df.to_dict(orient='records')
    docs = [Document(page_content=d['content'], metadata=d['metadata']) for d in docs_dict]
    vector_store = SupabaseVectorStore.from_documents(
        docs,
        embedding_llm,
        client=supabase_client,
        table_name="documents",
        query_name="match_documents_langchain",
    )
    return vector_store

def get_gaia_data_vector_store(supabase_client):
    return SupabaseVectorStore(
        client=supabase_client,
        embedding=embedding_llm,
        table_name="documents",
        query_name="match_documents_langchain"
    )

def get_gaia_data_retriever_tool(supabase_client):
    vector_store = get_gaia_data_vector_store(supabase_client)
    return create_retriever_tool(
        retriever=vector_store.as_retriever(),
        name="Question Search",
        description="A tool to retrieve similar questions from a vector store.",
    )



if __name__ == '__main__':
    # similar_question = vector_store.as_retriever().invoke("Box Office Mojo's 2020 Worldwide Box Office list, how")
    # retriever_tool.invoke("movie")
    supabase_client = get_supabase_client()
    # __upload_gaia_data(supabase_client)
    # retriever_tool = get_gaia_data_retriever_tool(supabase_client)
    # print(retriever_tool.invoke("Box Office Mojo's 2020 Worldwide Box Office list, how"))
    vector_store = get_gaia_data_vector_store(supabase_client)
    print(vector_store.similarity_search("Box Office Mojo's 2020 Worldwide Box Office list, how"))
    # pass