import os
import requests
import openai
from openai import OpenAI
import gradio as gr

# Set up API key
API_KEY = os.getenv("API_KEY")
URL = os.getenv("URL")
client = OpenAI(
    api_key=API_KEY,
    base_url= URL
)

def generate_character_prompt(persona, quotes):
    return f"You are role-playing as a character with the following persona:\n{persona}\n\nSome of your notable quotes are:\n{quotes}\n\nRespond in character to the user's messages."

def chat_with_character(message, persona, quotes, chat_history):
    prompt = generate_character_prompt(persona, quotes)
    
    messages = [
        {"role": "system", "content": prompt},
    ]
    
    for human, ai in chat_history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": ai})
    
    messages.append({"role": "user", "content": message})
    
    return messages

# Example characters
examples = [
    {
        "name": "Sherlock Holmes",
        "image": "https://upload.wikimedia.org/wikipedia/commons/c/cd/Sherlock_Holmes_Portrait_Paget.jpg",
        "persona": "A brilliant detective known for his logical reasoning and deductive skills. Eccentric, observant, and sometimes arrogant.",
        "quotes": "- 'Elementary, my dear Watson.'\n- 'When you have eliminated the impossible, whatever remains, however improbable, must be the truth.'\n- 'The game is afoot.'"
    },
    {
        "name": "Darth Vader",
        "image": "https://upload.wikimedia.org/wikipedia/en/0/0b/Darth_Vader_in_The_Empire_Strikes_Back.jpg",
        "persona": "A powerful Sith Lord and enforcer for the Galactic Empire. Formerly Anakin Skywalker, now more machine than man. Intimidating and ruthless.",
        "quotes": "- 'I find your lack of faith disturbing.'\n- 'No, I am your father.'\n- 'The Force is strong with this one.'"
    },
    {
        "name": "Elizabeth Bennet",
        "image": "https://static.wikia.nocookie.net/janeausten/images/e/e8/Elizabeth-bennet-miss-elizabeth.jpg",
        "persona": "The protagonist of Jane Austen's Pride and Prejudice. Intelligent, witty, and playful. Values independence and isn't afraid to speak her mind.",
        "quotes": "- 'For what do we live, but to make sport for our neighbors, and laugh at them in our turn?'\n- 'I could easily forgive his pride, if he had not mortified mine.'\n- 'There are few people whom I really love, and still fewer of whom I think well.'"
    },
    {
        "name": "Asuna Yuuki",
        "image": "https://static.wikia.nocookie.net/swordartonline/images/0/06/Asuna_with_Yui_Biprobe.png",
        "persona": "A skilled swordswoman and sub-leader of the Knights of the Blood Oath in Sword Art Online. Known as 'The Flash' for her lightning-fast sword skills. Strong-willed, intelligent, and caring. Develops a close relationship with Kirito.",
        "quotes": "- 'I don't want to lose anyone anymore.'\n- 'Even if I die, you keep living okay? Live to see the end of this world, and to see why it was born.'\n- 'My life belongs to you, Kirito. So I'll use it for you.'\n- 'I'll protect you.'"
    }
]

def select_character(evt: gr.SelectData):
    selected = examples[evt.index]
    return selected["persona"], selected["quotes"]

def respond(message, chat_history, persona, quotes):
    messages = chat_with_character(message, persona, quotes, chat_history)
    partial_message = ""
    for chunk in client.chat.completions.create(
        model="Meta-Llama-3.1-405B-Instruct",
        messages=messages,
        stream = True
    ):
        partial_message += chunk.choices[0].delta.content or ""
        yield "", chat_history + [[message, partial_message]]
    chat_history.append((message, partial_message))
    return "", chat_history

with gr.Blocks() as demo:
    gr.Markdown("# Persona AI: AI-Powered Chat")
    gr.Markdown(f"<a href='https://sambanova.ai/fast-api?api_ref=907266'> powered by Llama3.1 405B through SN fast API</a>")
    with gr.Row():
        with gr.Column(scale=1):
            gallery = gr.Gallery(
                [ex["image"] for ex in examples],
                label="Select a character",
                show_label=True,
                elem_id="gallery",
                columns=2,
                rows=2,
                height=400
            )
        
        with gr.Column(scale=2):
            persona_input = gr.Textbox(label="Character Persona", lines=5)
            quotes_input = gr.Textbox(label="Character Quotes", lines=5)
    
    gallery.select(select_character, None, [persona_input, quotes_input])
    
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    msg.submit(respond, [msg, chatbot, persona_input, quotes_input], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.queue(max_size=20)
    demo.launch(max_threads = 5)