import os
import base64
import tempfile
from inference import Chat, get_conv_template
import torch

def save_base64_to_tempfile(base64_str, suffix):
    header_removed = base64_str
    # 去除可能的data:image/...;base64,前缀
    if ',' in base64_str:
        header_removed = base64_str.split(',', 1)[1]

    data = base64.b64decode(header_removed)
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    tmp.write(data)
    tmp.close()
    return tmp.name

class EndpointHandler:
    def __init__(self, model_path: str):
        device = "cuda" if torch.cuda.is_available() else "cpu"
        self.chat = Chat(
            model_path=model_path,
            device=device,
            num_gpus=1,
            max_new_tokens=1024,
            load_8bit=False,
        )
        self.vision_feature = None
        self.modal_type = "text"
        self.chat.conv = get_conv_template("husky").copy()

    def __call__(self, data: dict) -> dict:
        # reset conversation if specified
        if data.get("clear_history"):
            self.chat.conv = get_conv_template("husky").copy()
            self.vision_feature = None
            self.modal_type = "text"

        prompt = data.get("inputs", "")
        image_input = data.get("image", None)
        video_input = data.get("video", None)

        # 判断image输入是路径还是base64字符串
        if image_input:
            if os.path.exists(image_input):
                # 直接用路径
                self.vision_feature = self.chat.get_image_embedding(image_input)
            else:
                # base64字符串，保存临时文件再处理
                tmp_path = save_base64_to_tempfile(image_input, suffix=".jpg")
                self.vision_feature = self.chat.get_image_embedding(tmp_path)
                os.unlink(tmp_path)  # 删除临时文件
            self.modal_type = "image"
            self.chat.conv = get_conv_template("husky").copy()

        elif video_input:
            if os.path.exists(video_input):
                self.vision_feature = self.chat.get_video_embedding(video_input)
            else:
                tmp_path = save_base64_to_tempfile(video_input, suffix=".mp4")
                self.vision_feature = self.chat.get_video_embedding(tmp_path)
                os.unlink(tmp_path)
            self.modal_type = "video"
            self.chat.conv = get_conv_template("husky").copy()

        else:
            self.modal_type = "text"
            self.vision_feature = None

        conversations = self.chat.ask(prompt, self.chat.conv, modal_type=self.modal_type)
        output = self.chat.answer(conversations, self.vision_feature, modal_type=self.modal_type)

        # 更新对话历史
        self.chat.conv.messages[-1][1] = output.strip()

        return {"output": output.strip()}