import gradio as gr
import requests
import json
import os
import uuid
import traceback

ENDPOINT = os.environ.get("ENDPOINT", None)
TOKEN = os.environ.get("TOKEN", None)

PLACEHOLDER = """
<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
   <img src="https://cdn.cosmosytu.online/images/cosmos_transparent.png" style="width: 80%; max-width: 550px; height: auto; opacity: 0.55;  "> 
   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">Cosmos LLaMa</h1>
   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Bana herhangi bir şey sorabilirsiniz...</p>
</div>
"""

css = """
h1 {
  text-align: center;
  display: block;
}
"""

def convert_to_messages(message, history):
    messages = []

    for item in history:
        messages.append({"role": "user", "content": item[0]})
        messages.append({"role": "assistant", "content": item[1]})

    messages.append({"role": "user", "content": message})

    return messages

def chat_cosmosllama(message: str, history: list, conversation_id: str, parent_message_id: str) -> (str, str):
    messages = convert_to_messages(message, history)

    headers = {
        "Cosmos-Token": TOKEN,
        "Content-Type": "application/json"
    }

    body = {"messages": messages}
    body["conversation_id"] = conversation_id
    
    if parent_message_id:
        body["parent_message_id"] = parent_message_id

    response = requests.post(f"https://model.cosmosytu.online{ENDPOINT}", json=body, headers=headers, stream=True)

    full_response = ""
    new_parent_message_id = parent_message_id

    if response.status_code == 200:
        for line in response.iter_lines():
            if line:
                decoded_line = line.decode('utf-8')
                try:
                    json_line = json.loads(decoded_line[6:])

                    if "assistantMessageId" in json_line:
                        new_parent_message_id = json_line["assistantMessageId"]

                    if "CosmosLlama" in json_line:
                        full_response += json_line['CosmosLlama']
                        yield full_response, new_parent_message_id
                except Exception:
                    traceback.print_exc()
    else:
        yield "Bir hata oluştu.", new_parent_message_id

def get_conversation_id():
    return str(uuid.uuid4())

def clear_conversation():
    return get_conversation_id(), []

chatbot = gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Cosmos LLaMa')

with gr.Blocks(fill_height=True, css=css) as demo:
    conversation_id_state = gr.State(value=get_conversation_id())
    parent_message_history = gr.State(value=[])
    
    def handle_user_message(message, history, conversation_id_state, parent_message_history):
        is_parent_message_set = False
        response_generator = chat_cosmosllama(message, history, conversation_id_state, parent_message_history[-1] if parent_message_history else None)
        for response, updated_parent_message_id in response_generator:
            if not is_parent_message_set:
                parent_message_history.append(updated_parent_message_id)
                is_parent_message_set = True
            yield response

    def undo_message(history, parent_message_history):
        if parent_message_history:
            parent_message_history.pop()
            history.pop()
            return history, parent_message_history
        return history, parent_message_history

    def retry_message(history, conversation_id_state, parent_message_history):
        if parent_message_history and history:
            parent_message_history.pop()
            history.pop()
            return history, parent_message_history
        return history, parent_message_history

    interface = gr.ChatInterface(
        fn=handle_user_message,
        chatbot=chatbot,
        fill_height=True,
        cache_examples=False,
        additional_inputs=[conversation_id_state, parent_message_history],
        retry_btn="🔄 Tekrar Cevapla",
        undo_btn="↩️ Geri Al",
        clear_btn="🗑️ Temizle",
        submit_btn="Gönder",
        stop_btn="Durdur"
    )

    demo.load(fn=clear_conversation, inputs=[], outputs=[conversation_id_state, parent_message_history])

    interface.undo_btn.click(
        fn=undo_message,
        inputs=[chatbot, parent_message_history],
        outputs=[chatbot, parent_message_history]
    )

    interface.retry_btn.click(
        fn=retry_message,
        inputs=[chatbot, conversation_id_state, parent_message_history],
        outputs=[chatbot, parent_message_history]
    )

    interface.clear_btn.click(fn=clear_conversation, inputs=[], outputs=[conversation_id_state, parent_message_history])
    interface.textbox.placeholder="Mesajınızı buraya giriniz..."

if __name__ == "__main__":
    demo.launch()