import os
import nltk
import certifi
from langchain_community.chat_models import ChatOllama
from langchain.prompts import ChatPromptTemplate, PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain.retrievers.multi_query import MultiQueryRetriever
from get_vector_db import get_vector_db

# Especifica una ruta donde tengas permisos de escritura
nltk_data_dir = "/app/nltk_data"
os.makedirs(nltk_data_dir, exist_ok=True)

# Configura NLTK para usar la nueva ruta
nltk.data.path.append(nltk_data_dir)

# Descarga el recurso 'punkt'
nltk.download('punkt', download_dir=nltk_data_dir)

# Configuración del certificado SSL
os.environ['SSL_CERT_FILE'] = certifi.where()

LLM_MODEL = os.getenv('LLM_MODEL', 'mistral')

# Función para obtener las plantillas de prompt
def get_prompt():
    QUERY_PROMPT = PromptTemplate(
        input_variables=["question"],
        template="""You are an AI language model assistant. Your task is to generate five
        different versions of the given user question to retrieve relevant documents from
        a vector database. By generating multiple perspectives on the user question, your
        goal is to help the user overcome some of the limitations of the distance-based
        similarity search. Provide these alternative questions separated by newlines.
        Original question: {question}""",
    )

    template = """Answer the question based ONLY on the following context:
    {context}
    Question: {question}
    """

    prompt = ChatPromptTemplate.from_template(template)

    return QUERY_PROMPT, prompt

# Función principal para manejar las consultas
def query(input):
    if input:
        # Inicializa el modelo de lenguaje
        llm = ChatOllama(model=LLM_MODEL)
        # Obtiene la base de datos vectorial
        db = get_vector_db()
        # Obtiene las plantillas de prompt
        QUERY_PROMPT, prompt = get_prompt()

        # Configura el retriever para generar múltiples consultas
        retriever = MultiQueryRetriever.from_llm(
            db.as_retriever(),
            llm,
            prompt=QUERY_PROMPT
        )

        # Define la cadena de procesamiento para recuperar el contexto, generar la respuesta y analizar la salida
        chain = (
            {"context": retriever, "question": RunnablePassthrough()}
            | prompt
            | llm
            | StrOutputParser()
        )

        response = chain.invoke(input)

        return response

    return None