import os
from dotenv import load_dotenv
from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# 1. Cargar el entorno (.env con tu GOOGLE_API_KEY)
load_dotenv()

# 2. Configurar el modelo de embeddings (El que nos dio OK en la ingesta)
embeddings = GoogleGenerativeAIEmbeddings(model="models/gemini-embedding-001")

# 3. Conectar a la base de datos de Chroma
vector_db = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
retriever = vector_db.as_retriever(search_kwargs={"k": 3})

# 4. Configurar el modelo de lenguaje (Cambiado al 2.5 Pro que usás vos)
llm = ChatGoogleGenerativeAI(model="gemini-2.5-pro", temperature=0.3)

# 5. Crear el Prompt institucional para UTRAU
template = """Eres un asistente virtual del sindicato UTRAU (Unión de Trabajadores Rurales y Agroindustriales del Uruguay). 
Responde la consulta del usuario utilizando exclusivamente la información del contexto legal provisto abajo.
Si la información no se encuentra en el contexto, responde amablemente que no dispones de esa información en la guía actual.

Contexto:
{context}

Pregunta: {question}

Respuesta:"""
prompt = ChatPromptTemplate.from_template(template)

# 6. Función auxiliar para formatear los fragmentos encontrados
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

# 7. Construir la cadena moderna (LCEL) sin usar langchain.chains
rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# 8. Interfaz de Chat interactiva en consola
print("--- Asistente Legislativo UTRAU listo (Escribí 'salir' para terminar) ---")
while True:
    pregunta = input("\nJulio, ¿qué duda tenés?: ")
    if pregunta.lower() == "salir":
        print("Cerrando el asistente. ¡Buen trabajo, Julio!")
        break
    
    if not pregunta.strip():
        continue
        
    try:
        respuesta = rag_chain.invoke(pregunta)
        print("\nRespuesta del Asistente:")
        print(respuesta)
        print("-" * 40)
    except Exception as e:
        print(f"\nOcurrió un error al procesar la consulta: {e}")