import pandas as pd
import openai
import faiss
import numpy as np
import time
import os
import pickle
import gradio as gr
from langchain.embeddings.openai import OpenAIEmbeddings
from io import StringIO
from huggingface_hub import hf_hub_download
from huggingface_hub import login

openai.api_key = os.getenv("OPENAI_API_KEY")
hf_token = os.getenv("HF_TOKEN")

login(token=hf_token)

def load_embeddings_and_faiss():
    embeddings_path = hf_hub_download(repo_id="chukbert/embedding-faq-medquad", filename="embeddings.pkl",repo_type="dataset", token=hf_token)
    faiss_index_path = hf_hub_download(repo_id="chukbert/embedding-faq-medquad", filename="faiss.index",repo_type="dataset", token=hf_token)

    faiss_index = faiss.read_index(faiss_index_path)
    
    with open(embeddings_path, 'rb') as f:
        question_embeddings = pickle.load(f)

    return faiss_index, question_embeddings

def retrieve_answer(question, faiss_index, embedding_model, answers, log_output, threshold=0.2):
    question_embedding = embedding_model.embed_query(question)
    distances, indices = faiss_index.search(np.array([question_embedding]), k=1)

    closest_distance = distances[0][0]
    closest_index = indices[0][0]
    log_output.write(f"closest_distance: {closest_distance}")

    if closest_distance > threshold:
        return "No good match found in dataset. Using GPT-4o-mini to generate an answer."

    return answers[closest_index]

def ask_openai_gpt4(question):
    response = openai.chat.completions.create(
        messages=[
            {"role": "user", "content": f"Answer the following medical question: {question}"}
        ],
        model="gpt-4o-mini",
        max_tokens=150
    )
    return response.choices[0].message.content

def chatbot(user_input):
    log_output = StringIO() 

    faiss_index, question_embeddings = load_embeddings_and_faiss()
    embedding_model = OpenAIEmbeddings(openai_api_key=openai.api_key)
    
    start_time = time.time()
    
    log_output.write("Retrieving answer from FAISS...\n")
    response_text = retrieve_answer(user_input, faiss_index, embedding_model, answers, log_output, threshold=0.3)
    
    if response_text == "No good match found in dataset. Using GPT-4o-mini to generate an answer.":
        log_output.write("No good match found in dataset. Using GPT-4o-mini to generate an answer.\n")
        response_text = ask_openai_gpt4(user_input)

    end_time = time.time() 
    response_time = end_time - start_time  

    return response_text, f"Response time: {response_time:.4f} seconds", log_output.getvalue()

demo = gr.Interface(
    fn=chatbot,  
    inputs="text", 
    outputs=[
        gr.Textbox(label="Chatbot Response"),
        gr.Textbox(label="Response Time"),   
        gr.Textbox(label="Logs") 
    ],
    title="Medical Chatbot with Custom Knowledge About Medical FAQ",
    description="A chatbot with custom knowledge using FAISS for quick responses or fallback to GPT-4o-mini when no relevant answer is found. Response time is also tracked."
)

if __name__ == "__main__":
    df = pd.read_csv("medquad.csv")
    questions = df['question'].tolist()
    answers = df['answer'].tolist()

    print(f"Loaded questions and answers. Number of questions: {len(questions)}, Number of answers: {len(answers)}")
    demo.launch()