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}\n\nUser: {message}\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,
    )
    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