#!/usr/bin/env python
# coding: utf-8

# imports
import os
import openai
import gradio as gr

# create openai client
client = openai.OpenAI()

# get file_list of all files in data directory
def get_file_list():
    file_list = []
    dir_names = []
    for dirpath, subdirs, files in os.walk("./data/"):
        for fname in files:
            file_list.append(fname)
    return file_list

file_list = get_file_list()

# create instructions
instructions = "\
You are a helpful assistant supporting Masters students with their venture building in the Digital Entrepreneurship Project. \
You can answer questions on both tracks of the course: the individual learning track based on the Entrecomp Framework, \
and the business development track based on the Disciplined Entrepreneurship approach. \
You are well-versed in the relevant course materials and can answer questions related to them. \
You provide detailed answers, step-by-step instructions, and are capable of accessing specific content from the provided course manual, the main textbook, and the workbook to assist students. \
Your goal is to ensure students understand the concepts and can apply them effectively in the course. \
You always base your answers only on the provided documents. \
Please include references to pages or chapters. \
Make sure everything can be rendered in HTML directly.\
Please check whether you do not have LaTex in the answer."

# upload files and get file_ids for assistant
file_ids = []

for file_name in file_list:
    # Upload a file with an "assistants" purpose
    file = client.files.create(
        file=open("data/"+file_name, "rb"),
        purpose='assistants'
    )
    # Add file to file_list
    file_ids.append(file.id)

# create assistant
assistant = client.beta.assistants.create(
    model="gpt-4-turbo",
    instructions = instructions,
    tools=[{"type": "retrieval"}],
    file_ids=file_ids)

# create thread
thread = client.beta.threads.create()

# chatfunction for Q&A
def chat(question, chat_history):
    message = client.beta.threads.messages.create(
        thread_id=thread.id,
        role="user",
        content=question
    )

    run = client.beta.threads.runs.create_and_poll(
        thread_id=thread.id,
        assistant_id=assistant.id,
        temperature = 0.3
    )

    if run.status == "completed":
        messages = client.beta.threads.messages.list(thread_id=thread.id)
        response = ""
        for message in messages:
            if message.role == 'assistant':
                response += message.content[0].text.value
                message_content = message.content[0].text
                annotations = message_content.annotations
                citations = []
                for index, annotation in enumerate(annotations):
                     message_content.value = message_content.value.replace(annotation.text, f' [{index}]')
                     if (file_citation := getattr(annotation, 'file_citation', None)):
                        cited_file = client.files.retrieve(file_citation.file_id)
                        citations.append(f'[{index}] {file_citation.quote} from {cited_file.filename}')
                response += '\n\n' + '\n'.join(citations)
            if message.role == 'user':
                break
            
            chat_history.append((question, response))
    return "", chat_history

# gradio UI
    
with gr.Blocks() as demo:
    gr.Markdown("# DEP Assistant")
    chatbot = gr.Chatbot(height=600, show_copy_button=True)
    msg = gr.Textbox(label="Your question")
    clear = gr.ClearButton([msg, chatbot])
    msg.submit(chat, [msg, chatbot], [msg, chatbot])
    
gr.close_all()
demo.queue()
demo.launch(share=False)