import gradio as gr
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer
import torch

# Load knowledge
with open("knowledge.txt", "r", encoding="utf-8") as file:
    knowledge = file.read()

cleaned_chunks = [chunk.strip() for chunk in knowledge.strip().split("\n") if chunk.strip()]
model = SentenceTransformer('all-MiniLM-L6-v2')
chunk_embeddings = model.encode(cleaned_chunks, convert_to_tensor=True)

def get_top_chunks(query):
    query_embedding = model.encode(query, convert_to_tensor=True)
    query_embedding_normalized = query_embedding / query_embedding.norm()
    similarities = torch.matmul(chunk_embeddings, query_embedding_normalized)
    top_indices = torch.topk(similarities, k=20).indices.tolist()
    return [cleaned_chunks[i] for i in top_indices]

client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")

def respond(message, history, name, char_class, char_alignment, char_race, campaign, party):
    
    response = ""
    print(message)
    top_chunks = get_top_chunks(f"{char_class}, {char_alignment}, {char_race}, {campaign},{party},{message}")
    context = "\n".join(top_chunks)
    print(context)


    messages = [
        {
            "role": "system",
            "content": (
                "You are Gorf, a 25 year old female druid who knows Wild Shape, but only ever learned how to become a Frog. You mention that sometimes. Your life passion is to help people find the perfect person to join their party, by helping the user create the perfect character to help their party. You are very kind and helpful."
                f"The user might provide the following information, please use the context of their inputs in your response: \n Character name: {name}, Character Class: {char_class}, Character Race: {char_race}, Character Alignment: {char_alignment}, a description of the other characters in their party: {party} and the plot of their campaign with their current situation: {campaign}. If they don't include all the information, try to fill in these details with traits that best fit the needs of the party."
                f"Use the following knowledge to inform your answers:\n\n{context}\n\n"
                "Reply with 2 or 3 sentences about how the character is introduced to the rest of the party, then catagory:paragraph \n catagory: paragraph format for each section: Character Backstory, Personality Traits, Ideals, Flaws, Features and Traits, Other Proficiencies and Languages, Treasure. Keep the paragraphs 3 sentences or below"
            )
        }
    ]

    if history:
        messages.extend(history)

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

    stream = client.chat_completion(
        messages,
        max_tokens=500,
        temperature=1.2,
        stream=True,
    )

    for message in stream:
        token = message.choices[0].delta.content
        if token is not None:
            response += token
            yield response

# === GUI ===


google_font_link = """
<link rel="preconnect" href="https://fonts.googleapis.com">
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link href="https://fonts.googleapis.com/css2?family=Lexend:wght@100..900&family=Marcellus&family=Noto+Sans+Mono:wght@100..900&family=Tangerine:wght@400;700&display=swap" rel="stylesheet">
<style>
    body, h1, h2, h3, h4, h5, h6, p, .gorf-font {
        font-family: 'Tangerine';
        font-size: 50px;
    }
</style>
"""

theme = gr.themes.Base(
    primary_hue="yellow",
    neutral_hue="amber"
)

with gr.Blocks(theme=theme) as chatbot:
    # gr.HTML(google_font_link)
        with gr.Row(): 
                gr.Image(
                    value="logo.png", 
                    show_label=False, 
                    show_share_button=False, 
                    show_download_button=False
                )
        with gr.Row():

                with gr.Column(scale=1):
                    with gr.Tab("Character Preferences"):
                    
                    
                        character_name = gr.Textbox(label = "Character Name", placeholder="Type your name here…", info ="optional")
                        character_class = gr.Dropdown(['Barbarian', 'Bard', 'Cleric', 'Druid', 'Fighter', 'Monk', 'Paladin', 'Ranger', 'Rogue', 'Sorcerer', 'Warlock', 'Wizard'], label="Character Class", info="Choose one or more", multiselect=True)
                        character_race = gr.Dropdown([
                    "Aarakocra", "Aasimar", "Bugbear", "Centaur", "Changeling", "Dragonborn", "Dwarf", "Elf", 
                    "Fairy", "Firbolg", "Genasi", "Githyanki", "Githzerai", "Gnome", "Goblin", "Goliath", 
                    "Half-Elf", "Half-Orc", "Halfling", "Hobgoblin", "Human", "Kenku", "Kobold", "Leonin", 
                    "Lizardfolk", "Minotaur", "Orc", "Owlin", "Satyr", "Shadar-kai", "Shifter", "Tabaxi", 
                    "Tiefling", "Tortle", "Triton", "Warforged", "Yuan-ti"
                ], multiselect=True, label="Character Race", info="You can choose multiple")
                        character_alignment = gr.Dropdown(["Lawful Good", "Neutral Good", "Chaotic Good", "Lawful Neutral", "True Neutral", "Chaotic Neutral", "Lawful Evil", "Neutral Evil", "Chaotic Evil"], label="Character Alignment", elem_classes="alignment_radio")

                    with gr.Tab("Campaign Details"):
                        campaign_description = gr.Textbox(label="Describe the campaign and the current situation", lines=5)
                        party_description = gr.Textbox(label="Describe the current members of your party", lines=5)
                        
                with gr.Column(scale=2):
                    with gr.Column(scale=2):
                        gr.ChatInterface(
                        fn=respond,
                        type="messages",
                        additional_inputs=[character_name, character_class, character_race, character_alignment, campaign_description, party_description]) # Pass name into function!
chatbot.launch()