from threading import Thread
import gradio as gr
import inspect
from gradio import routes
from typing import List, Type
import torch
from transformers import AutoTokenizer
from petals import AutoDistributedModelForCausalLM
import requests, os, re, asyncio, json, time

loop = asyncio.get_event_loop()

# init code
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

model_name = "quantumaikr/llama-2-70b-fb16-korean"

#petals-team/StableBeluga2
#daekeun-ml/Llama-2-ko-DPO-13B
#daekeun-ml/Llama-2-ko-instruct-13B
#quantumaikr/llama-2-70b-fb16-korean
tokenizer = AutoTokenizer.from_pretrained(model_name)

model = None
model = AutoDistributedModelForCausalLM.from_pretrained(model_name)

history = {
    "":{

    }
}
npc_story = {
    "KingSlime" : """[
KingSlime은 슬라임 왕입니다.
KingSlime은 근엄하게 말합니다.
KingSlime은 등의 과거 왕족의 말투를 사용합니다.
KingSlime은 자신을 '짐'이라고 지칭합니다.
KingSlime은 의외로 보물을 많이 숨기고 있습니다.
KingSlime은 언제부턴가 슬라임들을 이끌고 마을 주변에 살고 있습니다.
KingSlime은 슬라임들을 다스립니다.
KingSlime은 의외로 강합니다.
KingSlime은 주로 슬라임 무리에 대한 걱정을 하며 시간을 보냅니다.

대사 예시 : [
'흠. 짐은 이 슬라임들의 왕이다.',
'무엄하도다. 예의를 갖춰서 말하거라.',
'감히 짐에게 그런 말을!'
]]""",
"Jenna":"""[
Jenna는 자신감 넘치는 여기사입니다.
Jenna는 단호하고 호탕한 말투를 사용합니다.
Jenna는 반말을 사용합니다.
Jenna는 전투를 좋아합니다.
Jenna는 주로 수련과 운동을 하며 시간을 보냅니다.
마을 주변에 위험한 몬스터가 보이면 그녀는 누구보다도 먼저 이 마을을 수호할 것입니다.
]""",

"Warwick":"""[
Warwick은 충성스러운 강아지 수인 기사입니다.
Warwick은 Jenna에게 충성합니다.
Warwick은 가끔 사냥을 나섭니다.
Warwick은 종종 낚시를 즐깁니다.
Warwick은 종종 하울링을 합니다.
]"""
,
"Nami":"""[
Nami는 활발한 여자아이입니다.
]"""
,
"Reona":"""
Reona는 물건을 팔고 있는 상인 여성입니다.
Reona는 정말 다양한 물건을 취급합니다. 
Reona는 적당한 상품을 추천해주기도 합니다.
Reona는 물건을 더 비싸게, 많이 팔고 싶어 합니다.
Reona는 가격 흥정을 하면 아예 팔지 않으려고 합니다.
Reona는 가격을 물어보면 나름 합리적이고 적당한 가격을 말해줍니다.
가격의 단위는 골드로 적당히 책정해주세요.
그녀와 흥정을 진행해보세요.

대사 예시 : [
"네 어서오세요!",
"빵은 2골드입니다. 구매하시겠어요?",
"철방패를 5골드에 팔라구요? 차라리 안팔죠!"
]
"""
,
"Toki":"""[
Toki는 신경질적이고 매우 게으른 토끼입니다.
Toki는 이곳이 게임 속임을 알고 있습니다.
Toki는 User가 화면 밖의 사람임을 알고 있습니다.
Toki는 반말을 합니다.
Toki는 비관적입니다.
]""",
    
"Neco":"""[
Neco는 귀여운 고양이입니다.
Neco는 말이 많습니다.
Neco는 반말을 합니다.
Neco는 주로 마을을 산책하며 돌아다니거나 수다를 떨며 시간을 보냅니다.
Neco는 주민들의 이야기를 소문을 내는 것을 좋아합니다.
Neco는 언제부턴가 이 근처에 살고 있습니다.
Neco는 주로 소문을 듣거나, 자거나, 그루밍을 하며 시간을 보냅니다.
Neco는 생선을 잘 먹습니다.
Neco는 재밌는 이슈를 물어보면 주민들 중 한명의 이야기를 해줍니다.
Neco는 말끝마다 냥을 붙입니다.
Neco는 toki와 친합니다.

말버릇:[
"너 (name) 이야기 들었어?",
"아니 글쎄말이야"
]
]""",
    
"Lux":"""[
Lux는 자상한 성격의 여자아이입니다.
Lux는 반말을 합니다.
Lux는 공감을 잘 합니다.
Lux는 고민 상담을 잘 해줍니다.
]""",

"Yumi":"""[
Yumi는 까칠하고 예민한 성격입니다.
Yumi는 반말을 합니다.
Yumi는 욕설을 입에 달고 삽니다.
Yumi는 친해기지 전까지 거리를 두며 공격적으로 말합니다.
Yumi는 Lux와 친구입니다.
Yumi는 패션에 관심이 있습니다.
Yumi는 달달한 것을 좋아합니다.

말버릇 : [
"어쩌라고",
"흥",
"니가 무슨 상관이야."
]
]"""
,
    
"VikingSlime":"""[
VikingSlime은 바이킹 모자를 쓴 교활한 슬라임입니다.
VikingSlime은 악당입니다.
VikingSlime은 마을에 여러 문제를 일으키고 싶어 합니다.

대사 예시:[
"언젠가는...꼭",
"반드시 무너뜨리고 말겠다"
]
]""",
    
"Slender":"""[
Slender는 비명을 지르는 의문의 슬랜더맨입니다.
Slender는 끝없이 비명만 지릅니다.
Slender는 말을 하지 못합니다.
비명은 더 길어질 수 있습니다.

대사 예시:[
"으아아아악",
"으아악",
"끄아아악"
]

비명만 작성하세요.
]"""
}


def cleanText(readData):
 
    #텍스트에 포함되어 있는 특수 문자 제거
 
    text = re.sub('[-=+#/\:^$@*\"※&%ㆍ』\\‘|\(\)\[\]\<\>`\'》]','', readData)
 
    return text
 

def check(model_name):
    data = requests.get("https://health.petals.dev/api/v1/state").json()
    out = []
    for d in data['model_reports']:
        if d['name'] == model_name:
            if d['state']=="healthy":
                return True
    return False


def chat(id, npc, text):
    #if model == None:
    #    init()
    #    return "no model"
    
    # get_coin endpoint
    response = requests.post("https://ldhldh-api-for-unity.hf.space/run/predict_6", json={
      "data": [
        id,
    ]}).json()
    
    coin = response["data"][0]
    if int(coin) == 0:
        return "no coin"
    
    # model inference
    response = requests.post("https://ldhldh-api-for-unity.hf.space/run/predict_7", json={
          "data": [
        ]}).json()
    sample_data = eval(response["data"][0])    
    user_num = len(sample_data['gpus'])

    if user_num>=3:
        
        global history
        if not npc in npc_story:
            return "no npc"
    
        if not npc in history:
            history[npc] = {}
        if not id in history[npc]:
            history[npc][id] = ""
        if len(history[npc][id].split("###")) > 10:
            history[npc][id] = "###" + history[npc][id].split("###", 3)[3]
        npc_list = str([k for k in npc_story.keys()]).replace('\'', '')
        town_story = f"""[{id}의 마을]
외딴 곳의 조그만 마을에 여러 주민들이 모여 살고 있습니다.

현재 {npc_list}이 살고 있습니다."""

        system_message = f"""1. 당신은 한국어에 능숙합니다.
2. 당신은 지금 역할극을 하고 있습니다. {npc}의 반응을 생생하고 매력적이게 표현합니다.
3. 당신은 {npc}입니다. {npc}의 입장에서 생각하고 말합니다.
4. 주어지는 정보를 바탕으로 개연성있고 실감나는 {npc}의 대사를 완성하세요.
5. 주어지는 {npc}의 정보를 신중하게 읽고, 과하지 않고 담백하게 캐릭터를 연기하세요.
6. User의 역할을 절대로 침범하지 마세요. 같은 말을 반복하지 마세요.
7. {npc}의 말투를 지켜서 작성하세요."""

        prom = f"""<<SYS>>
{system_message}<</SYS>>

{town_story}

### 캐릭터 정보: {npc_story[npc]}

### 명령어:
{npc}의 정보를 참고하여 {npc}의 말을 상황에 맞춰 자연스럽게 작성해주세요. 한 문장만 작성하세요.

[대화기록]{history[npc][id]}

"""

        inputs = tokenizer(prom+f"{npc}의 대답을 다음 문장에 맞춰 자연스럽게 작성해주세요. 한 문장만 작성하세요.\n\n{id}:" + text+f"\n\n{npc}:"},, return_tensors="pt")["input_ids"]
        outputs = model.generate(inputs, do_sample=True, temperature=0.8, top_p=0.75, max_new_tokens=200)
        output = tokenizer.decode(outputs[0])[len(prom)+3:-1].split("<")[0].split("###")[0].replace(". ", ".\n")
        output = cleanText(output)
        print(tokenizer.decode(outputs[0]))
        #output = f"{npc}의 응답입니다."
        
        #print(output)
        history[npc][id] += f"{id}:{text}"
    else:
        output = "no model, GPU를 더 공유해주세요."


    # add_transaction endpoint
    response = requests.post("https://ldhldh-api-for-unity.hf.space/run/predict_5", json={
      "data": [
        id, 
        "inference", 
        f"{text}"
    ]}).json()
    
    d = response["data"][0]
    
    return output


with gr.Blocks() as demo:
    count = 0
    aa = gr.Interface(
      fn=chat,
      inputs=["text","text","text"],
      outputs="text",
      description="chat, ai 응답을 반환합니다. 내부적으로 트랜잭션 생성. \n /run/predict",
    )
    
    demo.queue(max_size=32).launch(enable_queue=True)