from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import gradio as gr
import gradio.components as grc
import random

repo_id = "TheBloke/Mistral-7B-Instruct-v0.1-GGUF"
filename = "mistral-7b-instruct-v0.1.Q4_K_M.gguf"
model_file_path = hf_hub_download(repo_id=repo_id, filename=filename)


def generate_characters(generate_backstories=False,
                        num_characters=1,
                        num_multiclass_characters=1):
    class_list = ["Barbarian", "Bard", "Cleric", "Druid", "Fighter",
                  "Monk", "Paladin", "Ranger", "Rogue", "Sorcerer", "Warlock", "Wizard"]
    sub_class_list = ["Barbarian - Beserker", "Barbarian - Wildheart", "Barbarian - Wild Magic", "Bard - College of Lore", "Bard - College of Valour", "Bard - College of Swords", "Cleric - Life Domain", "Cleric -  Light Domain", "Cleric - Trickery Domain", "Cleric - Knowledge Domain", "Cleric - Nature Domain", "Cleric - Tempest Domain", "Cleric - War Domain", "Druid - Circle of the moon", "Druid - Circle of the Land", "Druid - Circle of the Spores", "Fighter - Battle Master", "Fighter - Eldritch Knight", "Fighter - Champion", "Monk - Way of the open hand", "Monk - Way of shadow", "Monk - Way of the four elements",
                      "Paladin - Oath of Devotion", "Paladin - Oath of the ancients", "Paladin - Oath of Vengence", "Paladin - Oathbreaker", "Ranger - Beat Master", "Ranger - Hunter", "Ranger - Gloom Stalker", "Rogue - Thief", "Rogue - Arcane Trickster", "Rogue - Assassin", "Sorcerer - Draconic Bloodline", "Sorcerer - Wild Magic", "Sorcerer - Storm Sorcery", "Warlock - The Fiend", "Warlock - the great old one", "Warlock - Archfey", "Wizard - Abjuration", "Wizard - Conjuration", "Wizard - Divination", "Wizard - Enchantment", "Wizard - Evocation", "Wizard - Necromancer", "Wizard - Illusion", "Wizard - Transmutation"]
    race_list = ["Dragonborn", "Drow", "Dwarf", "Elf", "Githyanki",
                 "Gnome", "Half-Elf", "Half-Orc", "Halfling", "Human", "Tiefling"]
    sub_race_list = ["Black Dragonborn", "Blue Dragonborn", "Brass Dragonborn", "Bronze Dragonborn", "Copper Dragonborn", "Gold Dragonborn", "Green Dragonborn", "Red Dragonborn", "Silver Dragonborn", "White Dragonborn", "Lolth-Sworn Drow", "Seldarine Drow", "Gold Dwarf", "Shield Dwarf",
                     "Duegar (Dwarf)", "High Elf - Elf", "Wood Elf - Elf", "Githyanki", "Deep Gnome", "Forest Gnome", "Rock Gnome", "High Half-Elf", "Wood Half-Elf", "Drow Half-Elf", "Half-Orc", "Lightfoot Halfling", "Strongheart Halfling", "Human", "Asmodeus Tiefling", "Mephistopheles Tiefling", "Zariel Tiefling"]
    alignment_list = ["Chaotic Good", "Lawful Good", "Neutral Good",
                      "Chaotic Evil", "Lawful Evil", "Neutral Evil", "True Chaotic", "True Neutral"]
    background_list = ["Acolyte", "Charlatan", "Criminal", "Entertainer",
                       "Folk Hero", "Guild Artisan", "Noble", "Hermit", "Outlander", "Sage", "Soldier"]

    def generate_character():
        character = {
            'Class': random.choice(class_list),
            'Sub Class': random.choice(sub_class_list),
            'Race': random.choice(race_list),
            'Sub Race': random.choice(sub_race_list),
            'Alignment': random.choice(alignment_list),
            'Background': random.choice(background_list),
        }
        return character

    def generate_multiclass_character():
        character = generate_character()
        second_class = random.choice(
            [cls for cls in class_list if cls != character['Class']])
        levels_class1 = random.randint(1, 11)
        levels_class2 = 12 - levels_class1
        character['Class'] += f' ({levels_class1} levels), {second_class} ({levels_class2} levels)'
        return character

    characters = [generate_character() for _ in range(num_characters)]
    multiclass_characters = [generate_multiclass_character()
                             for _ in range(num_multiclass_characters)]
    characters.extend(multiclass_characters)

    def create_character_output(character, generate_backstories=False):
        output = ""
        for key, value in character.items():
            output += f"- **{key}:** {value}\n"
        output += "\n"
        if generate_backstories:
            backstory = generate_backstory(character)
            # Remove everything up to and including `[/INST]\n`
            output += backstory[backstory.find("[/INST]\n") + 8:]

        return output

    output = ""
    for idx, char in enumerate(characters, 1):
        output += f"### Character {idx}:\n"
        output += create_character_output(char, generate_backstories)

    return output


def generate_text_from_prompt(user_prompt,
                              max_tokens=4096,
                              temperature=1.0,
                              top_p=0.8,
                              echo=True,
                              stop=["<s>", "[/INST]", "Q"]):
    mistral_model = Llama(model_path=model_file_path,
                          n_ctx=4096,
                          n_gpu_layers=-1)

    # Define the parameters
    model_output = mistral_model(
        f"<s>[INST] {user_prompt} [/INST]\n",
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        echo=echo,
        stop=stop
    )

    return model_output


def generate_backstory(character):
    return generate_text_from_prompt(
        f"Please provide a name and creative, well-written backstory "
        f" for the following D&D Character: \n```{character}```\n"
    )["choices"][0]["text"].strip()


if __name__ == "__main__":
    iface = gr.Interface(fn=generate_characters,
                         title="D&D Character Generator",
                         inputs=[grc.Checkbox(label="Generate Backstories?"),
                                 grc.Slider(minimum=0, maximum=10,
                                            label="Single-class characters?",
                                            step=1),
                                 grc.Slider(minimum=0, maximum=10,
                                            label="Multi-class characters?",
                                            step=1)
                                 ],
                         description="Generates 3 random D&D characters.",
                         outputs="markdown")
    iface.launch()