import os
import faiss
import numpy as np
import json
import gradio as gr
from openai import OpenAI
from sentence_transformers import SentenceTransformer

# Step 1: Set up OpenAI API key
openai_api_key = os.environ.get("OPENAI_API_KEY", "")

client = OpenAI(api_key=openai_api_key)

# Step 2: Load the pre-trained FAISS index and SentenceTransformer model
index = faiss.read_index("faiss_index.bin")
model = SentenceTransformer('all-MiniLM-L6-v2')

def load_documents(docs_path):
    with open(docs_path, 'r', encoding='utf-8') as file:
        return json.load(file)

# Specify the path to your JSON file
docs_path = 'documents.json'
documents = load_documents(docs_path)
dimension = 1536

def get_embeddings(text):
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input = [text]
    )
    embedding = response.data[0].embedding
    return np.array(embedding, dtype='float32')

# Step 3: Function to search FAISS index
def search_index(query, k=3):
    # Convert query to an embedding
    query_vector = get_embeddings(query).reshape(1, -1).astype('float32')
    
    # Check if the index is not empty before searching
    if index.ntotal == 0:
        return "No documents in the index."

    # Search the FAISS index for the nearest neighbors
    distances, indices = index.search(query_vector, k)
    
    # Retrieve the top matching documents
    results = [documents[i] for i in indices[0] if i != -1]
    
    if results:
        return "\n\n".join(results)
    else:
        return "No relevant documents found."

# Step 4: Function to generate a response using OpenAI's GPT
def generate_response(context, user_input):
    prompt = f"{context}\n\nUser: {user_input}\nAssistant:"
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "system", "content": "You are a helpful assistant."},
                  {"role": "user", "content": prompt}],
        # stream=True,
    )
    
    # for chunk in stream:
    #     if chunk.choices[0].delta.content is not None:
    #         print(chunk.choices[0].delta.content, end="")
    return response.choices[0].message.content

# Step 5: Gradio chatbot function
def chatbot_interface(user_input, chat_history):
    # Step 5.1: Retrieve context using FAISS
    context = search_index(user_input)
    
    # Step 5.2: Generate a response using OpenAI GPT model
    response = generate_response(context, user_input)
    
    # Step 5.3: Update chat history
    chat_history.append((user_input, response))
    return chat_history, chat_history

def chat_gen(message, history):
    history_openai_format = []
    context = search_index(message)
    prompt = f"Context: {context}\n\n Question: {message}\nAssistant:"

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "system", "content": "You are a helpful assistant. You are assisting students who are studying at DDI program at Assumption University. Answer the questions with the provided context. Do not include based on the context or based on the documents in your answer. Please say you do not know if you do not know or cannot find the information needed."},
                  {"role": "user", "content": prompt}],
        stream=True,
    )
    partial_message = ""
    for chunk in response:
        if chunk.choices[0].delta.content is not None:
            partial_message = partial_message + chunk.choices[0].delta.content
            yield partial_message

initial_msg = "Hello! I am DDI assistant. You can ask me anything about DDI program. I am happy to assist you."
chatbot = gr.Chatbot(value = [[None, initial_msg]])
demo = gr.ChatInterface(chat_gen, chatbot=chatbot, fill_height=True).queue()


try:
    demo.launch(debug=True, share=False, show_api=False)
    demo.close()
except Exception as e:
    demo.close()
    print(e)
    raise e