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 = "petals-team/StableBeluga2"

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

model = None

history = {
    "":{

    }
}

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(outputs)
        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" + prompt + "\n\n### output:\n" + output
    ]}).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)