import gradio as gr
import openai
from prisma import Prisma
from prisma.models import public_users as Public_Users, Assistants, Chats, Messages
import asyncio

# Initialize Prisma client
prisma = Prisma()

async def chat_with_atlas(api_key, assistant_id, user_message):
    if not api_key or not assistant_id:
        return "Please provide both an API key and an Assistant ID."
     
    # Connect to the database
    if not prisma.is_connected():
        await prisma.connect()

    # Retrieve user and assistant settings from the database
    user: Public_Users = await prisma.public_users.find_unique(where={'apiKey': api_key})
    if not user:
        return "Invalid API key"

    assistant_settings: Assistants = await prisma.assistants.find_unique(where={'id': assistant_id})
    if not assistant_settings:
        return "Assistant not found"

    # Connect to OpenAI Assistant
    openai.api_key = user.openai_api_key
    openai_assistant_id = assistant_settings.openai_assistant_id
    if not openai_assistant_id:
        return "OpenAI Assistant ID not found for this assistant"

    # Send the user's message to OpenAI and get a response
    assistant = openai.Assistant.create(
        name="Atlas Assistant",
        instructions="You are a personal website assistant. You answer user questions.",
        tools=[{"type": "code_interpreter"}],
        model=assistant_settings.model or "gpt-4-1106-preview"
    )

    # Create a Thread for the conversation
    thread = openai.Thread.create()

    # Add a Message to the Thread with the user's input
    message = openai.Message.create(
        thread_id=thread.id,
        role="user",
        content=user_message
    )

    # Run the Assistant to process the user's message and generate a response
    run = openai.Run.create(
        thread_id=thread.id,
        assistant_id=assistant.id
    )

    # Check the Run status and wait for it to complete
    run_status = openai.Run.retrieve(
        thread_id=thread.id,
        run_id=run.id
    )
    while run_status['status'] != 'completed':
        # Optionally, implement a delay here before checking the status again
        run_status = openai.Run.retrieve(
            thread_id=thread.id,
            run_id=run.id
        )

    # Retrieve the Assistant's response messages
    messages = openai.Message.list(
        thread_id=thread.id
    )

    # Extract the Assistant's latest response content
    response_content = next(
        (msg['content'] for msg in messages['data'] if msg['role'] == 'assistant'),
        "No response from the assistant."
    ).strip()

    return response_content

iface = gr.Interface(
    fn=chat_with_atlas,
    inputs=["text", "text", "text"],
    outputs="text",
    live=True
)
iface.launch()