from langchain_openai import ChatOpenAI
from langchain_core.prompts.chat import PromptTemplate
from langchain_milvus import Milvus
from langchain_openai import OpenAIEmbeddings
import statistics
import concurrent.futures


def milvus_busqueda(collection_name: str, db_name: str, uri: str):
    embeddings = OpenAIEmbeddings(
        model="text-embedding-3-large",
        openai_api_key="sk-agrota-v7nqf1DdSh1EYdGJic5ST3BlbkFJr5HoyVHSYxf462AMV6hs",
    )
    
    conexion = Milvus(
        embedding_function=embeddings,
        collection_name=collection_name,
        connection_args={
            "db_name": db_name,
            "uri":uri,
            "token":"root:Milvus"
        },
    )
    return conexion


def procesar_colecciones(colecciones, db_name, uri, pregunta, k):
    lista_colecciones = []

    def procesar_coleccion(coleccion):
        coleccion_docs = milvus_busqueda(
            coleccion, db_name, uri
        ).similarity_search_with_score(pregunta, k=k)
        print(coleccion)
        for doc in coleccion_docs:
            puntaje = doc[1]
            if puntaje > 0.0:
                print(puntaje)
                return coleccion_docs
        return None

    with concurrent.futures.ThreadPoolExecutor() as executor:
        futuros = {
            executor.submit(procesar_coleccion, coleccion): coleccion
            for coleccion in colecciones
        }
        for futuro in concurrent.futures.as_completed(futuros):
            resultado = futuro.result()
            if resultado:
                lista_colecciones.append(resultado)
                print(len(lista_colecciones))
    return lista_colecciones

def procesar_colecciones_codigo(colecciones, db_name, uri, pregunta, k):
    lista_colecciones = []

    def procesar_coleccion_codigo(coleccion):
        coleccion_docs = milvus_busqueda(
            coleccion, db_name, uri
        ).similarity_search_with_score(pregunta, k=k)
        print(coleccion)
        for doc in coleccion_docs:
            puntaje = doc[1]
            if puntaje > 0.1:
                print(puntaje)
                return coleccion_docs
        return None

    with concurrent.futures.ThreadPoolExecutor() as executor:
        futuros = {
            executor.submit(procesar_coleccion_codigo, coleccion): coleccion
            for coleccion in colecciones
        }
        for futuro in concurrent.futures.as_completed(futuros):
            resultado = futuro.result()
            if resultado:
                lista_colecciones.append(resultado)
                print(len(lista_colecciones))
    return lista_colecciones

# def vectores(pregunta: str):
#     colecciones = [
#         "listado_productos",
#         "Catalogo_Maquinaria",
#         "links_catalogos",
#         "catalogos_quimicos",
#         "listados_quimicos",
#         "informacion_garantia",
#     ]

#     for coleccion in colecciones:
#         coleccion_docs = milvus_busqueda(coleccion).similarity_search_with_score(pregunta, k=10)
#         print(coleccion)
#         for documento, puntaje in coleccion_docs:
#             puntajes = []
#             puntajes.append(puntaje)
#             if statistics.mean(puntajes) > 0.65:
#                 print(statistics.mean(puntajes))
#                 return milvus_busqueda(coleccion)
#     return milvus_busqueda("info_completa")