from asyncio import constants
import gradio as gr
import requests
import os 
import re
import random

# GPT-J-6B API
API_URL = "https://api-inference.huggingface.co/models/EleutherAI/gpt-j-6B"
#HF_TOKEN = os.environ["HF_TOKEN"]
#headers = {"Authorization": f"Bearer {HF_TOKEN}"}

prompt = """
Bilbo is a hobbit rogue who wears a brown cloak and carries a ring.

Bremen is a human wizard, he wears a blue robe and carries a wand.
"""

examples = [["river"], ["night"], ["trees"],["table"],["laughs"]]


def npc_randomize():
    #name is a random combination of syllables
    vowels = list("aeiou")
    constants = list("bcdfghjklmnpqrstvwxyz")
    seperators=list("-'")
    name =""
    for i in range(random.randint(2,4)):
        name += random.choice(constants)
        name += random.choice(vowels)
        if random.random()<0.5:
            name += random.choice(constants)
        if random.random()<0.1:
            name += random.choice(seperators)

    #capitalize first letter
    name = name[0].upper() + name[1:]
    
    races="""Dwarf
    Elf
    Halfling
    Human
    Dragonborn
    Gnome
    Half-elf
    Half-orc
    Tiefling
    Aarakocra
    Genasi
    Goliath""".split("\n")

    races=[x.strip() for x in races]

    race=random.choice(races)

    print("foo",races,race)


    classes="""Barbarian
    Bard
    Cleric
    Druid
    Fighter
    Monk
    Paladin
    Ranger
    Rogue
    Sorcerer
    Warlock
    Wizard""".split("\n")
    
    classes=[x.strip() for x in classes]

    characterClass=random.choice(classes)

    pronoun=random.choices(["he","she","they"],weights=[0.45,0.45,0.1],k=1)[0]

    return name,race,characterClass,pronoun


    

def genericDescription():
    colors="""red
    blue
    green
    yellow
    orange
    purple
    pink
    brown
    black
    white""".split("\n")
    colors=[x.strip() for x in colors]
    outfits="""shirt
    pair of pants
    pair of shoes
    hat
    pair of glasses
    backpack
    belt
    tie
    cloak
    robe
    chain mail vest
    suit of plate armor
    suit of leather armor
    suit of studded leather armor
    suit of scale armor
    suit of chain mail armor
    suit of ring mail armor
    """.split("\n")
    outfits=[x.strip() for x in outfits]
    weapons="""sword
    dagger
    mace
    axe
    polearm
    bow
    crossbow
    sling
    club
    flail
    warhammer
    morningstar
    halberd
    war pick
    war sickle
    war hammer""".split("\n")
    weapons=[x.strip() for x in weapons]

    objects="""shield
    lantern
    sack
    severed head
    crystal""".split("\n")
    objects=[x.strip() for x in objects]
    
    desc=" wears a {color} {outfit}".format(color=random.choice(colors),outfit=random.choice(outfits))
    if random.random()<0.5:
        desc+=" and a {color} {outfit}".format(color=random.choice(colors),outfit=random.choice(outfits))
    
    if random.random()<0.5:
        desc+=" and carries a {weapon}".format(weapon=random.choice(weapons))
    elif random.random()<0.5:
        desc+=" and carries a {weapon} and a {object}".format(weapon=random.choice(weapons),object=random.choice(objects))
    else:
        desc+=" and carries two {weapon}s".format(weapon=random.choice(weapons))
        
    return desc


def npc_generate(name,race,characterClass,pronoun):

  desc="{name} is a {race} {characterClass}, {pronoun}".format(name=name,race=race,characterClass=characterClass,pronoun=pronoun)

  p = prompt + "\n"+desc
  print(f"*****Inside desc_generate - Prompt is :{p}")
  json_ = {"inputs": p,
            "parameters":
            {
            "top_p": 0.9,
          "temperature": 1.1,
          "max_new_tokens": 50,
          "return_full_text": False,
          }}
  #response = requests.post(API_URL, headers=headers, json=json_)
  response = requests.post(API_URL, json=json_)
  output = response.json()
  print(f"If there was an error? Reason is : {output}")


  #error handling
  if "error" in output:
    print("using fallback description method!")
    #fallback method
    longDescription=genericDescription()
  else:
    output_tmp = output[0]['generated_text']
    print(f"GPTJ response without splits is: {output_tmp}")
    if "\n\n" not in output_tmp:
        if output_tmp.find('.') != -1:
            idx = output_tmp.find('.')
            longDescription = output_tmp[:idx+1]
        else:
            idx = output_tmp.rfind('\n')
            longDescription = output_tmp[:idx]
    else:
        longDescription = output_tmp.split("\n\n")[0] # +"."
    longDescription = longDescription.replace('?','')
    print(f"longDescription being returned is: {longDescription}")
  return desc+longDescription

def desc_to_image(desc):
  print("*****Inside desc_to_image")
  desc = " ".join(desc.split('\n'))
  desc = desc + ", character art, concept art, artstation"
  steps, width, height, images, diversity = '50','256','256','1',15
  iface = gr.Interface.load("spaces/multimodalart/latentdiffusion")
  print("about to die",iface,dir(iface))

  prompt = re.sub(r'[^a-zA-Z0-9 ,.]', '', desc)
  print("about to die",prompt)


  img=iface(desc, steps, width, height, images, diversity)[0]
  return img

demo = gr.Blocks()

with demo:
  gr.Markdown("<h1><center>NPC Generator</center></h1>")
  gr.Markdown(
        "based on <a href=https://huggingface.co/spaces/Gradio-Blocks/GPTJ6B_Poetry_LatentDiff_Illustration> Gradio poetry generator</a>."
        "<div>first input name, race and class (or generate them randomly)</div>"
        "<div>Next, use GPT-J to generate a short description</div>"
        "<div>Finally, Generate an illustration 🎨 provided by Latent Diffusion model.</div>"
    )
  
  with gr.Row():
    b0 = gr.Button("Randomize name,race and class")
    b1 = gr.Button("Generate NPC Description")
    b2 = gr.Button("Generate Portrait")
  
  with gr.Row():  
    input_name = gr.Textbox(label="name",placeholder="Drizzt")
    input_race = gr.Textbox(label="race",placeholder="dark elf")
    input_class = gr.Textbox(label="class",placeholder="ranger")
    input_pronoun = gr.Textbox(label="pronoun",placeholder="he")

  with gr.Row():
    desc_txt = gr.Textbox(label="description",lines=7)
    output_image = gr.Image(label="portrait",type="filepath", shape=(256,256))
  
  b0.click(npc_randomize,inputs=[],outputs=[input_name,input_race,input_class,input_pronoun])
  b1.click(npc_generate, inputs=[ input_name,input_race,input_class,input_pronoun], outputs=desc_txt)
  b2.click(desc_to_image, desc_txt, output_image)
  #examples=examples

demo.launch(enable_queue=True, debug=True)