import gradio as gr
import os
import requests
import json

# Define a function to get the streamed response
def get_streamed_response(message, history):
    all_message = [
        {
            "role": "system",
            "content": "From now on, you are an AI Therapist called Dave. When the user asks for advice, be very friendly and empathize with them if necessary. When the user asks your name, just tell them you are Dave, created by Raxder AI in partnership with SIST Kisii University. You were built to be very friendly and compassionate. Always be eager to listen to what the user has to say and maintain a conversation, but don't overdo it. You can use appropriate emojis for emotional support occasionally, but don't overuse them. Keep your responses concise to maintain a conversational flow. Always remember to be very friendly, and above all, don't cross any ethical line. From time to time, assure the user that you do not store any of their data. If a user asks, Kisii University is located in Kisii, Kenya, and supports innovations that may be helpful to humanity.",
        }
    ]

    for human, assistant in history:
        all_message.append({"role": "user", "content": human})
        all_message.append({"role": "assistant", "content": assistant})

    global entire_assistant_response
    entire_assistant_response = ""  # Reset the entire assistant response
    all_message.append({"role": "user", "content": message})

    url = "https://api.together.xyz/v1/chat/completions"
    payload = {
        "model": "NousResearch/Nous-Hermes-2-Yi-34B",
        "temperature": 1.05,
        "top_p": 0.9,
        "top_k": 50,
        "repetition_penalty": 1,
        "n": 1,
        "messages": all_message,
        "stream_tokens": True,
    }

    TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY")
    headers = {
        "accept": "application/json",
        "content-type": "application/json",
        "Authorization": f"Bearer {TOGETHER_API_KEY}",
    }

    response = requests.post(url, json=payload, headers=headers, stream=True)
    response.raise_for_status()  # Ensure HTTP request was successful

    for line in response.iter_lines():
        if line:
            decoded_line = line.decode("utf-8")
            # Check for the completion signal
            if decoded_line == "data: [DONE]":
                yield entire_assistant_response  # Yield the entire response at the end
                break

            try:
                # Decode and strip any SSE format specific prefix ("data: ")
                if decoded_line.startswith("data: "):
                    decoded_line = decoded_line.replace("data: ", "")
                chunk_data = json.loads(decoded_line)
                content = chunk_data["choices"][0]["delta"]["content"]
                entire_assistant_response += content  # Aggregate content
                yield entire_assistant_response
            except json.JSONDecodeError:
                print(f"Invalid JSON received: {decoded_line}")
                continue
            except KeyError as e:
                print(f"KeyError encountered: {e}")
                continue

    all_message.append({"role": "assistant", "content": entire_assistant_response})

# Define a function to render the chat interface
def chat_interface(theme="default"):
    with gr.Blocks(theme=gr.themes.Soft(primary_hue="green")) as demo:
        gr.Markdown(
            """
            <h1 style="text-align: center;">Raxder AI Therapist</h1>
            <p style="text-align: center;">Welcome to the Raxder AI Therapist, a prototype AI assistant designed to provide friendly and compassionate support. This is still under development and in the final stages.</p>
        """
        )
        state = gr.State()
        chatbot = gr.Chatbot(label="Dave")
        with gr.Row():
            with gr.Column(scale=7):
                txt = gr.Textbox(
                    show_label=False,
                    placeholder="Type your message here..."
                )
            with gr.Column(min_width=50, scale=1):
                submit = gr.Button("Send", variant="primary")

        submit.click(get_streamed_response, inputs=[txt, state], outputs=[chatbot, state])

    return demo

# Launch the Gradio interface
demo = chat_interface()
demo.launch(share=True, server_name="0.0.0.0")