import gradio as gr
import os
import time
import json
from openai import OpenAI

# Initialize OpenAI client
client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))

# Load assistant and vector store from export
with open('assistant_config.json') as f:
    config = json.load(f)

assistant_id = config['assistant_id']
vector_store_id = config['vector_store_id']

# Create a new conversation thread
thread = client.beta.threads.create()

def respond(message, history):
    client.beta.threads.messages.create(
        thread_id=thread.id,
        role='user',
        content=message
    )

    run = client.beta.threads.runs.create(
        thread_id=thread.id,
        assistant_id=assistant_id
    )

    while run.status != 'completed':
        time.sleep(1)
        run = client.beta.threads.runs.retrieve(
            thread_id=thread.id,
            run_id=run.id
        )
        if run.status == 'failed':
            return 'Sorry, I encountered an error processing your request.'

    messages = client.beta.threads.messages.list(
        thread_id=thread.id,
        order='desc',
        limit=1
    )

    response = messages.data[0].content[0].text.value
    annotations = messages.data[0].content[0].text.annotations
    citations = []
    for annotation in annotations:
        if hasattr(annotation, 'file_citation'):
            citations.append(f"Source: {annotation.file_citation.file_id}")
    if citations:
        response += '\n\n**Sources:** ' + ', '.join(citations)

    return response

with gr.Blocks(title="AI Research Assistant") as demo:
    gr.Markdown("# AI Research Assistant")
    # Theme: Academic Research Helper, Clear, minimal, organized
    gr.ChatInterface(
        respond,
        chatbot=gr.Chatbot(label="Detailed answers with citations, with follow up questions"),
        textbox=gr.Textbox(placeholder="Ask questions about research papers"),
    )
    # Upload research papers (PDF, DOCX, TXT)

demo.launch(ssr_mode=False)