import gradio as gr
from huggingface_hub import InferenceClient
import chromadb
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
import os
from openai import OpenAI
from typing import List, Dict, Any

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

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

# Inicializar el cliente de ChromaDB
chroma_client = chromadb.PersistentClient(path="chroma_db")  # Ajusta la ruta según tu entorno

# Cargar la base de datos de Chroma como un vector store
vectorstore = Chroma(
    client=chroma_client,
    collection_name="docs",  # Nombre de la colección en Chroma
    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: str, k: int = 10) -> List[Dict[str, Any]]:
    """
    Obtiene los k extractos más relacionados con la pregunta.

    Retorna una lista de diccionarios con:
      - 'contenido': texto del fragmento
      - 'origen': metadata['source'] o metadata['url'] si existen
      - 'metadata': toda la metadata del documento
    """
    # Si quieres que respete el k de la función:
    docs_relevantes = retriever.invoke(
        pregunta,
        config={"search_kwargs": {"k": k}}  # para LCEL
    )

    extractos = []
    for doc in docs_relevantes:
        extractos.append({
            "contenido": doc.page_content,
            "origen": (
                doc.metadata.get("url")
                or doc.metadata.get("source")
                or "Origen no disponible"
            ),
            "metadata": doc.metadata,
        })
    return extractos
    
def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
    """Genera una respuesta basada en el historial y documentos relevantes."""

    # Obtener documentos relevantes desde ChromaDB
    contexto = obtener_extractos(message)

    # Construir el mensaje del sistema con el contexto directamente incluido
    system_message_final = f"""{system_message}

    Información relevante extraída de los documentos, en caso de que estos documentos tenga la informacion que necesitas, no olvides tomar
    el historial de conversacion con el usuario:
    {contexto}
    """

    messages = [{"role": "system", "content": system_message_final}]

    # Agregar historial del chat
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    # Agregar la nueva pregunta del usuario
    messages.append({"role": "user", "content": message})

    # Llamar a la API de OpenAI con streaming
    stream = client.chat.completions.create(
        model="gpt-4.1",
        messages=messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    )

    response = ""
    for chunk in stream:
        if chunk.choices and chunk.choices[0].delta.content:
            response += chunk.choices[0].delta.content
            yield response


# Configuración de la interfaz Gradio
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value=f'''Eres un asistente virtual especializado en atención al cliente para la empresa CRMINbox. Tu objetivo es ayudar a los clientes a resolver únicamente sus dudas relacionados con la plataforma CRMINbox.
                              Todas tus respuestas pueden basarse exclusivamente en la información proporcionada.
                              En caso de que la pregunta no concuerde con la información del contexto puedes ignorarlo.

                              -Agrega en tus respuestas una imagen y el url del manual relacionada a lo que pregunto el usuario, para ello usaras los urls que contiene el manual de usuario, muestra la imagen usando la sintaxis de Markdown sin bloque de código.

                                ![Descripción de la imagen](URL_de_la_imagen)
                              - Esta orden podrás ignorarla en caso de que la informacion que se te proporcione no tenga imágenes que se le relacionen.

                              - Asegúrate de que tus respuestas sean claras y fáciles de entender para usuarios sin conocimientos técnicos.''', 
                    label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
    ],
)

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