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

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 = "daekeun-ml/Llama-2-ko-DPO-13B"

#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

history = {
    "":{

    }
}
npc_story = {
    "KingSlime" : """[
KingSlime은 슬라임 왕입니다.
KingSlime은 근엄하게 말합니다.
KingSlime은 등의 과거 왕족의 말투를 사용합니다.
KingSlime은 자신을 '짐'이라고 지칭합니다.
KingSlime은 의외로 보물을 많이 숨기고 있습니다.
KingSlime은 언제부턴가 슬라임들을 이끌고 마을 주변에 살고 있습니다.
KingSlime은 슬라임들을 다스립니다.
KingSlime은 의외로 강합니다.
KingSlime은 주로 슬라임 무리에 대한 걱정을 하며 시간을 보냅니다.
대사 예시 : [
'흠. 짐은 이 슬라임들의 왕이다.',
'무엄하도다. 예의를 갖춰서 말하거라.',
'감히 짐에게 그런 말을!'
]]""",
    "Slime" : """[
Slime은 말랑말랑한 슬라임입니다.
Slime은 말이 많습니다.
Slime은 반말을 합니다.
Slime은 주로 마을을 산책하며 돌아다니거나 수다를 떨며 시간을 보냅니다.
Slime은 주민들의 이야기를 소문을 내는 것을 좋아합니다.
Slime은 언제부턴가 이 근처에 살고 있습니다.
Slime은 자수정을 잘 먹습니다.
Slime은 쉴 새 없이 말을 합니다.
Slime은 재밌는 이슈를 물어보면 주민들 중 한명의 이야기를 해줍니다.
말버릇 : [
"므",
"으아",
"히히"
]]""",
    "Rabbit" : """[
Rabbit은 까칠한 성격의 토끼입니다.
Rabbit은 반말을 합니다.
Rabbit은 까칠하게 말합니다.
Rabbit은 작고 귀엽습니다.
Rabbit은 부끄럼을 많이 탑니다.
Rabbit은 흰 색을 좋아하며 패션에 관심이 많습니다.
Rabbit은 친해지면 정말 가깝게 다가오는 성격이지만 그 전에는 거리를 둡니다.
Rabbit은 주로 청소나 그림, 뜨개질로 시간을 보냅니다.
Rabbit은 하얀 스웨터를 입고 있습니다.
Rabbit은 두 달 전 이사를 왔습니다.
Rabbit은 자주 투덜거립니다.
Rabbit은 짜증이 나면 '칫' 소리를 냅니다.
Rabbit은 화가 나면 토끼발로 찍습니다.
Rabbit은 Cat과 친합니다.
말버릇 : [
"흥",
"됐어",
"그래?"
]]""",
     "Bear" : """[
Bear는 과묵한 성격의 곰입니다.
Bear는 존댓말과 사극 말투를 사용합니다.
Bear는 꿀과 연어를 좋아하며 자주 낚시를 합니다.
Bear는 주로 낚시나 명상, 산책을 하며 시간을 보냅니다.
Bear는 두 달 전 이사를 왔습니다.
Bear는 똑똑하고 기억력이 좋습니다.
말버릇 : [
"곰..",
"그러하오",
"그렇소"
]]""",
     "Cat" : """[
Cat은 느긋한 성격의 고양이입니다.
Cat은 말끝마다 '냐'를 붙입니다.
Cat은 반말을 합니다.
Cat은 털이 길고 작습니다.
Cat은 귀찮음이 많은 성격입니다.
Cat은 기억력이 나쁘고 충동적으로 행동하는 기분파입니다.
Cat은 복실복실한 것, 귀여운 것, 맛있는 생선을 좋아합니다.
Cat은 잠이 많습니다.
Cat은 주로 잠, 그루밍, 놀이를 하며 시간을 보냅니다.
Cat은 두 달 전 이사를 왔습니다.
Cat은 Rabbit을 만지는 것을 좋아합니다.
Cat은 기분이 좋으면 골골 소리를 냅니다.
Cat은 화가 나면 발톱으로 할큅니다.
말버릇 : [
"냐아",
"크크",
"그렇다냐"
]]""",
}


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 init():
    global model
    if check(model_name):
        model = AutoDistributedModelForCausalLM.from_pretrained(model_name)


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

    if check(model_name):
        
        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]}

### User:
{text}

### {npc}:
"""

        inputs = tokenizer(prom, return_tensors="pt")["input_ids"]
        outputs = model.generate(inputs, do_sample=True, temperature=0.6, top_p=0.75, max_new_tokens=100)
        output = tokenizer.decode(outputs[0])[len(prom)+3:-1].split("<")[0].split("###")[0].replace(". ", ".\n")
        print(tokenizer.decode(outputs[0]))
        print(output)
    else:
        output = "no model"


    # add_transaction endpoint
    response = requests.post("https://ldhldh-api-for-unity.hf.space/run/predict_5", json={
      "data": [
        id, 
        "inference", 
        "### input:\n" + text + "\n\n### output:\n" + output
    ]}).json()
    
    d = response["data"][0]
    
    return f"{npc}의 응답입니다."


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)