import os
import chromadb
from groq import Groq
from dotenv import load_dotenv
from chroma_setup import initialize_client
from embedding import query_chroma_db

# Carga las variables de entorno desde el archivo .env
load_dotenv()

# Inicializa el cliente de Groq con la clave de API
groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))

# Inicializa el cliente de ChromaDB
chroma_client = initialize_client()

# Verifica o crea la colección donde se guardan las incrustaciones (embeddings)
collection_name = "text_embeddings"
try:
    collection = chroma_client.get_collection(name=collection_name)
except chromadb.api.segment.InvalidCollectionException:
    collection = chroma_client.create_collection(name=collection_name)

def query_llm(query: str, user_id: str, document_id: str):
    """
    Realiza una consulta al LLM (modelo de lenguaje) usando la información
    más relevante obtenida de la base de datos ChromaDB.

    Parámetros:
    -----------
    query : str
        Pregunta o consulta que se le hará al LLM.
    user_id : str
        Identificador único de usuario.
    document_id : str
        Identificador único de documento.

    Retorna:
    --------
    str
        Respuesta generada por el LLM en base al contexto encontrado.
    """
    if not user_id or not document_id:
        raise ValueError("Se requieren tanto user_id como document_id para realizar la consulta.")

    # Consulta a la base de datos para obtener el contexto
    context = query_chroma_db(user_id, document_id, query)

    # Registra el contexto antes de enviarlo al LLM
    print(f"Contexto para el LLM: {context}")

    if not context.strip():
        return "No se encontró contexto para esta consulta."

    prompt = f"""Eres un asistente inteligente. Basado en el siguiente contexto, 
proporciona una respuesta resumida y bien interpretada a la pregunta.
Contexto: {context}

Pregunta: {query}

Por favor, ofrece una respuesta detallada y clara basada en el documento."""

    chat_completion = groq_client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
        model="llama3-8b-8192",
        max_tokens=500,
    )

    return chat_completion.choices[0].message.content