import openai
import os
from dotenv import load_dotenv
import gradio as gr
import time
import asyncio
from typing import AsyncGenerator

# Load environment variables
load_dotenv()

class AssistantChat:
    def __init__(self):
        self.openai_client = openai
        self.openai_client.api_key = os.getenv("AURON_API_KEY")
        self.assistant_id = os.getenv("ASSISTANT_ID")
        self.sessions = {}

    def create_thread(self):
        try:
            response = self.openai_client.beta.threads.create()
            return response.id
        except Exception as e:
            print(f"Error creating thread: {e}")
            return None

    async def get_assistant_response(self, thread_id: str, input_text: str) -> str:
        try:
            # Create the message
            await asyncio.to_thread(
                self.openai_client.beta.threads.messages.create,
                thread_id=thread_id,
                role="user",
                content=input_text
            )

            # Create and start the run
            run = await asyncio.to_thread(
                self.openai_client.beta.threads.runs.create,
                thread_id=thread_id,
                assistant_id=self.assistant_id
            )

            response_text = ""
            while True:
                run_status = await asyncio.to_thread(
                    self.openai_client.beta.threads.runs.retrieve,
                    thread_id=thread_id,
                    run_id=run.id
                )

                if run_status.status == "completed":
                    messages = await asyncio.to_thread(
                        self.openai_client.beta.threads.messages.list,
                        thread_id=thread_id,
                        order="desc",
                        limit=1
                    )
                    if messages.data:
                        response_text = messages.data[0].content[0].text.value
                    break
                elif run_status.status == "failed":
                    response_text = "Sorry, there was an error generating the response."
                    break

                await asyncio.sleep(0.1)

            return response_text

        except Exception as e:
            return f"Error: {e}"

chat_instance = AssistantChat()

async def chat_response(message: str, history) -> AsyncGenerator[str, None]:
    if not history:
        session_key = message
        chat_instance.sessions[session_key] = {
            'thread_id': chat_instance.create_thread(),
            'count': 0
        }
    else:
        session_key = history[0][0]
    
    session = chat_instance.sessions[session_key]
    session['count'] += 1

    if session['count'] >= 200:
        yield "You've hit your free message limit. To continue the conversation, consider applying as an Alpha tester. Click here to [Continue](https://www.auronspeaks.com/apply)"
        return

    response = await chat_instance.get_assistant_response(session['thread_id'], message)
    
    partial_message = ""
    for char in response:
        partial_message += char
        yield partial_message
        await asyncio.sleep(0.01)  # Add a small delay for visual effect

demo = gr.ChatInterface(
    fn=chat_response,
    title="Valentine’s Day Edition / Alpha 04",
    description="This AI is funny and flirty! Treat them like your long-distance date and see how they respond. (Please be respectful. You can ask for virtual hugs but NSFW chat is not allowed)",
    examples=[
        "Hi, I'd love to chat with you!",
        "Hello Sweetheart"
    ],
    css="footer {visibility: hidden}"
)

if __name__ == "__main__":
    demo.queue().launch(share=True)