import gradio as gr
from gradio_client import Client as GrClient
import inspect
from gradio import routes
from typing import List, Type
from aiogoogletrans import Translator 

import requests, os, re, asyncio


loop = asyncio.get_event_loop()
gradio_client = GrClient(os.environ.get('GrClient_url'))
translator = Translator()
# Monkey patch
def get_types(cls_set: List[Type], component: str):
    docset = []
    types = []
    if component == "input":
        for cls in cls_set:
            doc = inspect.getdoc(cls)
            doc_lines = doc.split("\n")
            docset.append(doc_lines[1].split(":")[-1])
            types.append(doc_lines[1].split(")")[0].split("(")[-1])
    else:
        for cls in cls_set:
            doc = inspect.getdoc(cls)
            doc_lines = doc.split("\n")
            docset.append(doc_lines[-1].split(":")[-1])
            types.append(doc_lines[-1].split(")")[0].split("(")[-1])
    return docset, types
routes.get_types = get_types


# App code
def mbti(x):
    t = loop.run_until_complete(translator.translate(x, src='ko', dest='en'))
    str_trans = re.sub('[-=+,#/\?:^.@*\"※~ㆍ!』‘|\(\)\[\]`\'…》\”\“\’·]', '', t.text)
    result = gradio_client.predict(
				str_trans,	# str representing input in 'User input' Textbox component
				fn_index=2
    )
    r = sorted(eval(result), key=lambda x : x['score'], reverse=True)
    
    return r

def chat(x):
    x = f"[***너는 Assistant입니다. 상대에게 다양한 질문을 하며 대화를 이끌고 있습니다. Human에게 긍정적이고, 공감하며, 최대한 길게 대답해주세요***] {x}"
    x = x.replace('friend','Human').replace('you','Assistant')
    print("\n___________________\n" + f"{x}")
    result = gradio_client.predict(
        x,
        # str representing input in 'User input' Textbox component
		0.91,	# float, representing input in 'Top-p (nucleus sampling)' Slider component
		40,	    # int, representing input in 'Top-k (nucleus sampling)' Slider component
		0.7,	# float, representing input in 'Temperature' Slider component
		30,	    # int, representing input in 'Max New Tokens' Slider component
		1.2,	# float, representing input in 'repetition_penalty' Slider component
		fn_index=0
    )
    result = str(result)
    output = result[len(x.rsplit(':', 1)[0])+2:]
    output = re.sub('하하','ㅎㅎ', output)
    output = output.split('띓')[0]
    output = output.split('endoftext')[0]
    output = re.sub('[=+#/\:@*\"※ㆍ』‘|\\\<\>\(\)\[\]`\'…》\”\“\’·]', '', output)
    #output = re.sub('[a-zA-Z]',' ',output)

    return output

def yn(x):
    result = gradio_client.predict(
				x,	# str representing input in 'User input' Textbox component
				fn_index=1
    )
    return result

with gr.Blocks() as demo:
    count = 0
    aa = gr.Interface(
      fn=chat,
      inputs="text",
      outputs="text",
      description="chat",
    #examples= [[f"\nfriend: 너는 꿈이 뭐야? \n\n### \nyou: "],[f"\nyou: 너는 무슨 색을 가장 좋아해? \nfriend: 글쎄 너는? \n\n### \nyou: "]]
        examples= [[f"\nHuman: 너는 꿈이 뭐야? \n\n### \nAssistant:   "],[f"\nAssistant: 너는 무슨 색을 가장 좋아해? \nHuman: 나는 파란색이 좋더라. 너는 무슨 색이 좋아? \n\n### \nAssistant:   "]]
    )
    
    bb = gr.Interface(
      fn=mbti,
      inputs="text",
      outputs="text",
      description="mbti"
    )
    
    cc = gr.Interface(
      fn=yn,
      inputs="text",
      outputs="text",
      description="yn",
        examples= [[f"그래라"],[f"별론데?"]]
    )
    demo.queue(max_size=32).launch(enable_queue=True)