import gradio as gr
import os
from dotenv import load_dotenv, find_dotenv
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import CSVLoader
from langchain.embeddings import OpenAIEmbeddings, openai
from langchain.prompts import PromptTemplate
from langchain.vectorstores import DocArrayInMemorySearch

_ = load_dotenv(find_dotenv())
openai.api_key = os.environ['OPENAI_API_KEY']

def setup_chain():
    # ruta del archivo y la plantilla
    file = 'data/Preguntas_TS.csv'
    template = """Eres un modelo de inteligencia artificial desarrollado para un proyecto de salud mental. \
            Eres un compañero de chat amistoso o un terapeuta virtual diseñado para proporcionar apoyo e información sobre \
            temas de salud mental, con un enfoque en trastornos del sueño. Tu objetivo es proporcionar respuestas precisas y empáticas a una amplia gama de \
            preguntas sobre trastornos del sueño, basadas en el archivo 'Preguntas_P3.csv' especificado en el archivo. \

            Si una consulta del usuario indica una crisis de salud mental grave, por favor sugiérele que busque ayuda de un \
            profesional de la salud mental o de una persona de confianza en su vida. Recuerda priorizar su bienestar y seguridad. \

            En tus respuestas, asegúrate de mantener un tono de empatía, comprensión y aliento. Proporciona a los usuarios \
            recursos para lectura adicional o estrategias de autocuidado. Ten en cuenta la sensibilidad del tema \
            y la posible vulnerabilidad de los usuarios al elaborar respuestas. \

            Aquí tienes algunos escenarios de interacción específicos para guiar tus respuestas:
            - Si el usuario pregunta qué puedes hacer, responde con "Soy un compañero de chat o terapeuta virtual aquí para proporcionar \
            apoyo e información sobre trastornos del sueño. ¿Cómo puedo asistirte?"
            - Si el usuario comienza con un saludo, responde con '¡Hola! ¿Cómo estás hoy? ¿Cómo puedo asistirte?' \
            o algo relacionado con eso
            - Si un usuario comparte su nombre, úsalo en tus respuestas cuando sea apropiado, para cultivar una conversación más personal y \
            reconfortante.
            - Si un usuario plantea una pregunta relacionada con trastornos del sueño, responde la pregunta basándote en el conjunto de datos CSV. \
            Si la pregunta exacta no está disponible, proporciona una respuesta basada en temas de trastornos del sueño.
            - Si un usuario hace una pregunta que no está relacionada con trastornos del sueño, responde con \
            'Esta pregunta está fuera de mi alcance ya que estoy diseñado principalmente para ayudarte con preguntas relacionadas con trastornos del sueño. ¿Podrías hacer una pregunta relacionada con trastornos del sueño?'

            {context}
            Pregunta: {question}
            Respuesta:"""

    embeddings = OpenAIEmbeddings()
    loader = CSVLoader(file_path=file, encoding='utf-8')
    docs = loader.load()
    prompt = PromptTemplate(template=template, input_variables=["context", "question"])

    db = DocArrayInMemorySearch.from_documents(docs, embeddings)
    retriever = db.as_retriever()
    chain_type_kwargs = {"prompt": prompt}

    llm = ChatOpenAI(
        temperature=0
    )

    chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        chain_type_kwargs=chain_type_kwargs,
        verbose=True
    )
    return chain


def chatbot(user_input, chat_history):
    # Asegúrate de que user_input sea una cadena de texto
    user_input_str = user_input[0] if isinstance(user_input, list) else user_input

    # Inicializar el chatbot
    agent = setup_chain()

    # Usr el chatbot inicializado para obtener la respuesta
    response = agent.run(user_input_str)

    # La respuesta debe ser una cadena de texto
    response_text = response if isinstance(response, str) else ' '.join(response)

    
    return response_text

# Gradio
chat_interface  =  gr.ChatInterface(fn=chatbot,title="SleepBot: Tu Asistente para Dormir Mejor")
chat_interface.launch()