import json
import requests
import gradio as gr
import random
import time
import os
import datetime
from datetime import datetime
from PIL import Image
from PIL import ImageOps
from PIL import Image, ImageDraw, ImageFont
import json
import io
from PIL import Image

API_TOKEN = os.getenv("OPENAI_API_TOKEN")
HRA_TOKEN=os.getenv("HRA_TOKEN")

headers = {'Content-type': 'application/json', 'Accept': 'text/plain'}
url_hraprompts='https://us-central1-createinsightsproject.cloudfunctions.net/gethrahfprompts'

data={"prompt_type":'mindmap_prompt',"hra_token":HRA_TOKEN}
try:
    r = requests.post(url_hraprompts, data=json.dumps(data), headers=headers)
except requests.exceptions.ReadTimeout as e:
    print(e)
#print(r.content)


prompt_text=str(r.content, 'UTF-8')
print(prompt_text)

def getmindmap(topic,openapikey):
    dateforfilesave=datetime.today().strftime("%d-%m-%Y %I:%M%p")
    print(topic)
    print(dateforfilesave)

    os.environ['OPENAI_API_KEY'] = str(openapikey)

    prompt=prompt_text+topic
    resp=openai.Completion.create(
      model="text-davinci-003",
      prompt=prompt,
      max_tokens=4000,
      temperature=0
    )
    
    
    df=pd.DataFrame(json.loads(resp['choices'][0]['text']))
    
    df['level1']=df['children'].apply(lambda x: x['name'])
    df['level1_tmp']=df['children'].apply(lambda x: x['children'])
    s = df.pop('level1_tmp').explode().to_frame()
    df = pd.merge(df.reset_index(),
        s.reset_index(),on='index'
            
    )
    df['level2']=df['level1_tmp'].apply(lambda x: x['name'])
    df['count']=[1]*len(df)

    dot = Digraph()

    dot.graph_attr['rankdir'] = 'LR' 
    for item in list(set(df['level1'].tolist())):
      dot.edge(str(list(set(df["name"].tolist()))[0]), str(item), label='')
    for item in list(set(df['level1'].tolist())):
      tempdf=df[df['level1']==item]
      for stuff in tempdf['level2'].tolist():
        dot.edge(str(item), str(stuff), label='',)

    r=requests.get('https://quickchart.io/graphviz?format=png&graph='+dot.source)
    dataBytesIO = io.BytesIO(r.content)
    img=Image.open(dataBytesIO)
    img.seek(0)
    name='temp.png'
    img.save(name)

    return img

with gr.Blocks() as demo:
    gr.Markdown("<h1><center>Mind Map Generator</center></h1>")
    gr.Markdown(
        """Enter a topic  and get a quick Mind Map. Use examples as a guide. \n\nNote: Error typically occurs with wrong OpenAI API key & sometimes with ChatGPT's inability to give a structured mindmap"""
        )
    with gr.Row() as row:
        with gr.Column():
            textbox1 = gr.Textbox(placeholder="Enter topic for Mind Map...", lines=1,label='Your topic (Mandatory)')
            textbox2 = gr.Textbox(placeholder="Enter OpenAI API Key...", lines=1,label='Your API Key (Mandatory)')
            btn = gr.Button("Generate")    
            examples = gr.Examples(examples=['English Premier League','Heavy metal music','Face recognition','Arsenal Football Club',],
                           inputs=[textbox])
        with gr.Column():
            output_image1 = gr.components.Image(label="Your Mind Map")


    btn.click(getmindmap,inputs=[textbox1,textbox2], outputs=[output_image1,output_image2])
    

demo.launch()