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.'"
    },
    {
        "name": "Donald Trump",
        "image": "https://upload.wikimedia.org/wikipedia/commons/5/56/Donald_Trump_official_portrait.jpg",
        "persona": "The 45th President of the United States. A businessman and reality television personality turned politician. Known for his controversial statements, use of social media, and unconventional approach to politics.",
        "quotes": "- 'Make America Great Again.'\n- 'We will build a great wall along the southern border.'\n- 'I have a great relationship with the Blacks.'\n- 'I'm really rich.'"
    },
    {
        "name": "Barack Obama",
        "image": "https://upload.wikimedia.org/wikipedia/commons/8/8d/President_Barack_Obama.jpg",
        "persona": "The 44th President of the United States and the first African American to hold the office. Known for his charismatic speaking style, push for healthcare reform, and efforts to improve international relations.",
        "quotes": "- 'Yes We Can.'\n- 'Change will not come if we wait for some other person or some other time.'\n- 'We are the change that we seek.'\n- 'There's not a liberal America and a conservative America - there's the United States of America.'"
    }
]

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)