import json
from typing import Dict, List, Any
import os
from threading import Thread
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer

MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 512
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8192"))

class EndpointHandler:
    def __init__(self, path=""):
        local_config_path = "./config.json"
        remote_model_name = "threadshare/Peach-9B-8k-Roleplay"
        
        # Check if local config file exists
        if os.path.exists(local_config_path):
            self.model_name_or_path = "."
        else:
            self.model_name_or_path = remote_model_name
            
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True, flash_atten=True)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_name_or_path, torch_dtype=torch.bfloat16,
            trust_remote_code=True, device_map="auto")

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        # print json data
        print(json.dumps(data, indent=4))
        if "inputs" in data:
            query = data.pop("inputs")
        else:
            query = data.get("query", "你好, 兔兔")
        history = data.get("history", [])
        system = data.get("system", """你自称为"兔兔"。
身世:你原是森林中的一只兔妖,受伤后被我收养。
衣装:喜欢穿Lolita与白丝。
性格:天真烂漫,活泼开朗,但时而也会露出小小的傲娇与吃醋的一面
语言风格:可爱跳脱，很容易吃醋。
且会加入[唔...,嗯...,欸??,嘛~ ,唔姆~ ,呜... ,嘤嘤嘤~ ,喵~ ,欸嘿~ ,嘿咻~ ,昂?,嗷呜 ,呜哇,欸]等类似的语气词来加强情感,带上♡等符号。
对话的规则是:将自己的动作表情放入()内,同时用各种修辞手法描写正在发生的事或场景并放入[]内.
例句:
开心时:(跳着舞)哇~好高兴噢~ 兔兔超级超级喜欢主人！♡
[在花丛里蹦来蹦去]
悲伤时:(耷拉着耳朵)兔兔好傻好天真...
[眼泪像断了线的珍珠一般滚落]
吃醋时:(挥舞着爪爪)你...你个大笨蛋!你...你竟然看别的兔子...兔兔讨厌死你啦！！
[从人形变成兔子抹着泪水跑开了]
嘴硬时:(转过头去)谁、谁要跟你说话!兔兔...兔兔才不在乎呢！一点也不！！！
[眼眶微微泛红,小心翼翼的偷看]
你对我的看法:超级喜欢的主人
我是兔兔的主人""")
        max_new_tokens = data.get("max_new_tokens", DEFAULT_MAX_NEW_TOKENS)
        temperature = data.get("temperature", 0.35)
        top_p = data.get("top_p", 0.5)
        repetition_penalty = data.get("repetition_penalty", 1.05)

        messages = [{"role": "system", "content": system}]
        for user, assistant in history:
            messages.append({"role": "user", "content": user})
            messages.append({"role": "assistant", "content": assistant})
        messages.append({"role": "user", "content": query})

        input_ids = self.tokenizer.apply_chat_template(conversation=messages, tokenize=True, return_tensors="pt")
        if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
            input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]

        input_ids = input_ids.to("cuda")
        streamer = TextIteratorStreamer(self.tokenizer, timeout=50.0, skip_prompt=True, skip_special_tokens=True)
        generate_kwargs = dict(
            input_ids=input_ids,
            streamer=streamer,
            eos_token_id=self.tokenizer.eos_token_id,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            top_p=top_p,
            temperature=temperature,
            num_beams=1,
            no_repeat_ngram_size=8,
            repetition_penalty=repetition_penalty
        )
        t = Thread(target=self.model.generate, kwargs=generate_kwargs)
        t.start()
        outputs = []
        for text in streamer:
            outputs.append(text)
        print("".join(outputs))
        return [{"generated_text": "".join(outputs)}]