import gradio as gr
import openai
import json
from pymongo import MongoClient
import os

# Initialize MongoDB client
client = MongoClient(os.environ['DB_URL'])

db = client['test']
collection = db['gradio']

def get_saved_data():
    saved_data = collection.find({}, {"_id": 0, "name": 1})
    options = [item['name'] for item in saved_data]
    return options

def chat_with_gpt(question, api_key, temperature, system_message):
    openai.api_key = api_key
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": system_message},
            {"role": "user", "content": question}
        ],
        temperature=temperature
    )
    assistant_reply = response['choices'][0]['message']['content']
    return f"{assistant_reply}"

def update_dropdown_choices():
    new_choices = get_saved_data()
    return gr.Dropdown.update(choices=new_choices)

def save_to_mongodb(name, question, system_message):
    if not name.strip():  # Check if name is empty or just whitespace
        return "Please enter a name.", None  # Return a message and None to indicate no update for the dropdown
    collection.insert_one({"name": name, "question": question, "system_message": system_message})
    return "Saved to MongoDB.", update_dropdown_choices()

def update_textboxes(selected_name):
    selected_data = collection.find_one({"name": selected_name}, {"_id": 0})
    return selected_data['question'], selected_data['system_message'], selected_data['name']


with gr.Blocks() as app:
    saved_data_dropdown = gr.Dropdown(get_saved_data(), label="Select Saved Data")
    name = gr.Textbox(lines=1, placeholder="Name", label="Name")
    question = gr.Textbox(lines=2, placeholder="What's your question?", label="Question")
    api_key = gr.Textbox(lines=1, placeholder="Your OpenAI API Key", label="API Key")
    temperature = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.7, label="Temperature")
    system_message = gr.Textbox(lines=1, placeholder="System Message (Optional)", value="You are a helpful assistant.", label="System Message")
    
    with gr.Row():
        chat_btn = gr.Button("Chat with GPT")
        save_btn = gr.Button("Share to community")
    
    output = gr.Textbox(label="Result", interactive=False)

    chat_btn.click(chat_with_gpt, inputs=[question, api_key, temperature, system_message], outputs=output)
    save_btn.click(save_to_mongodb, inputs=[name, question, system_message], outputs=[output, saved_data_dropdown])
    saved_data_dropdown.select(update_textboxes, inputs=[saved_data_dropdown], outputs=[question, system_message, name])

app.launch()