import gradio as gr
from huggingface_hub import InferenceClient
import chromadb
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
import os
import time
from uuid import uuid4
from openai import OpenAI
import json

# Configurar la API Key de OpenAI
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
ASSISTANT_ID = os.getenv("OPENAI_ASSISTANT_ID")

# Inicializar el cliente de OpenAI
client = OpenAI(api_key=OPENAI_API_KEY)

# Inicializar el cliente de ChromaDB
chroma_client = chromadb.PersistentClient(path="chroma_db")

# Cargar la base de datos de Chroma como un vector store
vectorstore = Chroma(
    client=chroma_client,
    collection_name="docs_medicina",
    embedding_function=OpenAIEmbeddings(model="text-embedding-3-small", openai_api_key=OPENAI_API_KEY)
)

# Crear un retriever
retriever = vectorstore.as_retriever()

# Función para obtener extractos relevantes
def obtener_extractos(pregunta):
    docs_relevantes = retriever.get_relevant_documents(pregunta)
    return [
        (doc.page_content,
         f"Libro: {doc.metadata.get('nombre_libro', 'N/A')}, Página: {doc.metadata.get('numero_pagina', 'N/A')}")
        for doc in docs_relevantes
    ]

# Diccionario para almacenar threads por sesión
session_threads = {}
# Diccionario para rastrear el número de mensajes ya procesados por thread
thread_message_count = {}

def respond(message, history: list[tuple[str, str]], system_message):
    # Generar un ID de sesión basado en el historial
    session_id = str(hash(json.dumps(history, sort_keys=True)))
    
    # Obtener o crear thread único por sesión
    thread_id = session_threads.get(session_id)
    
    if not thread_id:
        # Crear nuevo thread
        thread = client.beta.threads.create()
        thread_id = thread.id
        session_threads[session_id] = thread_id
        thread_message_count[thread_id] = 0
        
        # Agregar historial completo solo para threads nuevos
        for user_msg, assistant_msg in history:
            client.beta.threads.messages.create(
                thread_id=thread_id,
                role="user",
                content=user_msg
            )
            client.beta.threads.messages.create(
                thread_id=thread_id,
                role="assistant",
                content=assistant_msg
            )
            thread_message_count[thread_id] += 2
    else:
        # Para threads existentes, solo agregar mensajes nuevos
        current_history_length = len(history) * 2  # Multiplicar por 2 porque cada intercambio son 2 mensajes
        messages_in_thread = thread_message_count.get(thread_id, 0)
        
        # Agregar solo los mensajes nuevos que no están en el thread
        if current_history_length > messages_in_thread:
            # Calcular cuántos intercambios nuevos hay
            new_exchanges = (current_history_length - messages_in_thread) // 2
            
            # Agregar solo los intercambios nuevos
            for i in range(-new_exchanges, 0):
                user_msg, assistant_msg = history[i]
                client.beta.threads.messages.create(
                    thread_id=thread_id,
                    role="user",
                    content=user_msg
                )
                client.beta.threads.messages.create(
                    thread_id=thread_id,
                    role="assistant",
                    content=assistant_msg
                )
                thread_message_count[thread_id] += 2

    # Agregar el nuevo mensaje del usuario
    client.beta.threads.messages.create(
        thread_id=thread_id,
        role="user",
        content=message
    )
    thread_message_count[thread_id] += 1

    # Ejecutar el assistant
    with client.beta.threads.runs.stream(
        thread_id=thread_id,
        assistant_id=ASSISTANT_ID
    ) as stream:
        
        response = ""
        for event in stream:
            if event.event == 'thread.message.delta':
                delta = event.data.delta
                if delta.content:
                    chunk = delta.content[0].text.value
                    response += chunk
                    yield response  # Respuesta acumulativa

    # Obtener la respuesta más reciente
    messages = client.beta.threads.messages.list(thread_id=thread_id, limit=1)
    if messages.data:
        latest_message = messages.data[0]
        if latest_message.role == "assistant":
            thread_message_count[thread_id] += 1  # Incrementar contador por la respuesta del assistant
            yield latest_message.content[0].text.value
            return

# Configuración de la interfaz Gradio
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="""Eres un asistente especializado en medicina.
                   Uso de fuentes: Utiliza la información del contexto médico proporcionado.
Citación obligatoria: En cada respuesta debes incluir:
Nombre exacto del libro/documento de referencia
Número de página específica
Capítulo o sección cuando sea relevante
Formato: "[Fuente: Título del Libro, Página X, Capítulo Y]"
        Ejemplo de respuesta:
"Según la información disponible, [respuesta médica]. Esta información se encuentra en [Fuente: Título del Libro, Página X].""", label="System Message")
    ],
)

if __name__ == "__main__":
    demo.launch()