import os
from langchain_community.vectorstores import Pinecone
from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import FlashrankRerank
from langchain.retrievers import EnsembleRetriever, BM25Retriever
from langchain_core.tools import tool
from typing import Any

from apps.agent.utils import load_and_split_docs
from apps.agent.constant import (
    INDEX_NAME_WEWEB, 
    INDEX_NAME_XANO,
    URLS_WEWEB,
    URLS_XANO,
)

embeddings = FastEmbedEmbeddings(model_name="jinaai/jina-embeddings-v2-small-en")
compressor = FlashrankRerank(model="ms-marco-MiniLM-L-12-v2")

def ensemble_retriever(index_name: str, docs: Any, embeddings, compressor):
    # retriever
    vectorstore = Pinecone.from_existing_index(embedding=embeddings, index_name=index_name)
    retriever = vectorstore.as_retriever()

    # bm25
    bm25 = BM25Retriever.from_documents(docs)
    bm25.k = 6
    ensemble_retriever = EnsembleRetriever(retrievers=[retriever, bm25],
                                       weights=[0.6, 0.4])
    
    # reranker
    reranker_retriever = ContextualCompressionRetriever(
        base_compressor=compressor, base_retriever=ensemble_retriever
    )
    return reranker_retriever

# load data
data_xano = load_and_split_docs(URLS_XANO)
data_weweb = load_and_split_docs(URLS_WEWEB)

# create retriever
retriever_xano = ensemble_retriever(INDEX_NAME_XANO, data_xano, embeddings, compressor)
retriever_weweb = ensemble_retriever(INDEX_NAME_WEWEB, data_weweb, embeddings, compressor)

@tool
def tool_xano(query: str):
    """
    Searches and returns excerpts from the Xano documentation
    """
    return retriever_xano.invoke(query)
    

@tool
def tool_weweb(query: str):
    """
    Searches and returns excerpts from the Weweb documentation
    """
    return retriever_weweb.invoke(query)