"""AntGLM Chat-model data format.

格式化 AntGLM 以及各种开源模型的符号系统:
  - 确定 Chat 模型依赖的文件数据结构协议
  - 确定单轮/多轮的统一结构
  - 确定 Chat 符号系统的协议, 包括角色定义、分隔符等
  - 方便做开源模型依赖的 prompt 转换
  - 支持工具、代码、推理等支持

参考 FastChat Conversation 对象的设计思路.
Reference: https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
"""

import copy
import dataclasses
import logging
import re
import uuid
from copy import deepcopy
from enum import IntEnum, auto
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)


class PromptStyle(IntEnum):
    """Prompt styles."""

    # 原始 antglm format 格式, 单轮指令没有结构, 多轮 `第1轮\n用户: xx\n机器人: xx\n`
    ANTGLM_RAW = auto()
    # Chat format 格式, 单轮多轮统一为 chat format 格式
    ANTGLM_CHAT = auto()
    # 单轮指令没有结构, 只有多轮为 chat format 格式
    ANTGLM_ONLY_MULTITURN_CHAT = auto()
    # OpenAI ChatML 格式, 包括千问
    CHATML = auto()
    # LLAMA2 格式
    LLAMA2 = auto()
    # ChatGLM 1/2 格式
    CHATGLM = auto()
    # ChatGLM3 格式
    CHATGLM3 = auto()
    # 百川格式
    BAICHUAN2 = auto()


@dataclasses.dataclass
class Chat:
    """Chat 数据符号结构, 格式化 AntGLM 以及各种开源模型的符号系统.

    Examples:

    ```python
    >>> from antllm.data.chat_format import Chat

    >>> ### 从 json 数据结构创建 chat 对象, 并且 format 结构使用 AntGLM 原始结构
    >>> input_json = {
    ...     "messages": [
    ...         {"role": "HUMAN", "content": "讲一个笑话"},
    ...         {"role": "ASSISTANT", "content": "为什么猪不能上网？因为它们会被网上的“猪”骗！哈哈哈！"},
    ...         {"role": "HUMAN", "content": "不好笑，换个程序员的笑话"}
    ...      ],
    ... }
    >>> chat = Chat.from_json(input_json, name='antglm_raw')

    >>> ### 根据 chat 对象创建大模型训练所需 pack 数据
    >>> pack_data = chat.prompt_pack
    >>> print(pack_data)

    >>> ### 根据 chat 对象创建大模型训练所需 input, output 数据
    >>> data = chat.prompt_inout
    >>> print(data)

    >>> ### 根据 chat 对象创建大模型预测用的 prompt
    >>> prompt = chat.prompt_str
    >>> print(prompt)

    >>> ### 从大模型训练数据 {"input": "xx", "output": "xx"} 中创建 chat 对象
    >>> data = {
    ...     'input': (
    ...         '第1轮\n用户: 讲一个笑话\n机器人: 为什么猪不能上网？因为它们会被网上的“猪”骗！哈哈哈！\n'
    ...         '第2轮\n用户: 不好笑，换个程序员的笑话\n机器人:'
    ...     ),
    ...     'output': ''
    ... }
    >>> chat = Chat.from_inout(data, name='antglm_raw')

    >>> ### 从大模型 pack 训练数据创建 chat 对象列表
    >>> pack_data = {
    ...     'inputs': ['第1轮\n用户: 讲一个笑话\n机器人:', '第2轮\n用户: 不好笑，换个程序员的笑话\n机器人:', '第1轮\n用户: 写首诗\n机器人:'],
    ...     'outputs': [
    ...         '为什么猪不能上网？因为它们会被网上的“猪”骗！哈哈哈！\n',
    ...         '为什么程序员总是喜欢使用黑色主题？因为他们喜欢“黑暗模式”（Dark Mode),这样他们就可以在晚上加班时更好地隐藏自己的错误！',
    ...         '']
    ... }
    >>> chats = Chat.from_pack(pack_data, name='antglm_raw')
    >>> assert len(chats) == 2
    >>> print(chats[0])
    >>> print(chats[1])

    >>> ### 显示总交互轮数 (以用户输出多少次为轮数个数)
    >>> print(chat.turns_num)

    >>> ### 根据 chat 对象创建 json 格式化输出
    >>> data_json = chat.to_json()
    >>> print(data_json)

    >>> ### 增加轮次信息
    >>> content = (
    ...     '为什么程序员总是喜欢使用黑色主题？'
    ...     '因为他们喜欢“黑暗模式”（Dark Mode),这样他们就可以在晚上加班时更好地隐藏自己的错误！'
    ... )
    >>> chat.append_message(chat.role_assistant, content)

    >>> ### 将 chat 对象转成 OpenAI ChatCompletion 接口的入参
    >>> openai_messages = chat.to_openai_api_messages()
    >>> print(openai_messages)

    >>> ### 复制一个 chat 对象
    >>> chat_new = chat.copy()
    ```
    """

    # 数据结构名称
    id: str = None

    # format 支持: antglm_raw, antglm_chat, chatglm1, chatglm2, llama2, qwen, baichuan2
    name: Optional[str] = None

    # Prompt 风格
    prompt_style: Optional[PromptStyle] = None

    # System Template 和 message
    system_template: str = "<role>SYSTEM</role>{}"
    system_message: str = ""

    # 角色定义
    role_human: str = "HUMAN"
    role_assistant: str = "ASSISTANT"
    role_observation: str = "OBSERVATION"
    role_template: str = "<role>{}</role>"

    # 每轮符号定义
    turn_start: str = ""
    human_end: str = ""
    assistant_start: str = ""
    assistant_end: str = ""
    assistant_end_ids: Optional[List[int]] = None
    general_role_end: str = ""

    # agent 符号定义
    tool_template = "<tool>{}</tool>"
    code_template = "<code>{}</code>"
    arithemetic_templte = "<arithemetic>{}</arithemetic>"

    # All messages. Each item is (role, message).
    messages: List[Tuple[str, str]] = ()

    # messages 中用于 few-shot messages
    offset: int = 0

    # 其他 meta data
    source: Optional[str] = None
    lang: Optional[str] = None
    topic: Optional[str] = None

    # 原始 json 数据
    origin_json: Optional[dict] = None

    @property
    def support_names(self) -> Dict[str, str]:
        """支持的数据对象名称."""
        return {
            "antglm_raw": "原始 antglm format 格式, 单轮指令没有结构, 多轮 `第1轮\\n用户:xx\\n机器人xx\\n`",
            "antglm_chat": "Chat format 格式, 单轮多轮统一为 chat format 格式",
            "chatglm1": "chatglm1 format",
            "chatglm2": "chatglm2 format",
            "llama2": "llama2 format",
            "qwen": "千问 format",
            "baichuan2": "百川 2 format",
        }

    @classmethod
    def from_json(
        cls,
        input: dict,
        name: Optional[str] = None,
        prompt_style: Optional[PromptStyle] = None,
    ):
        """从文件数据结构到数据对象的转换.

        Params:
            name: `Optional[str]`, 符号系统名称
                - format 支持: antglm_raw, antglm_chat, chatglm1, chatglm2, llama2, qwen, baichuan2
                - 如果指定了 format name, 使用该 name 符号系统, 否则使用 input 中 `name` 字段

            prompt_style: `Optional[PromptStyle]`, 指定 prompt 风格, 默认使用和 name 一致的风格

            input: `dict`, 文件中的 json dict 对象, 协议为:
                - 既支持 `messages` 字段, 也支持 `turns` 字段
            {
                "id": "xxx",
                "name": "antglm",
                "source": "xxx",
                "lang": "xx",
                "topic": "xx",
                "system_template": "",
                "system_message": "xx",
                "messages": [
                    {
                        "role": "HUMAN",
                        "content": "Hi"
                    },
                    {
                        "role": "ASSISTANT",
                        "content": "Hello"
                    },
                    {
                        "role": "OBSERVATION",
                        "content": "xxx"
                    },
                    {
                        "role": "ASSISTANT",
                        "content": "xxx"
                    }
                ],
                "turns": [
                    {"HUMAN": "xxx", "OBSERVATION": "xx", "ASSISTANT": "xx"}
                ]
            }

        Returns:
            `Chat` 对象
        """
        _id = input.get("id")
        if name:
            _name = name
        else:
            _name = input.get("name")
        source = input.get("source")
        lang = input.get("lang")
        topic = input.get("topic")
        kwargs = {}
        if "system_template" in input:
            kwargs["system_template"] = input["system_template"]
        if "system_message" in input:
            kwargs["system_message"] = input["system_message"]

        # 转换成 Chat 对象
        chat = cls(
            id=_id,
            name=_name,
            prompt_style=prompt_style,
            source=source,
            lang=lang,
            topic=topic,
            origin_json=deepcopy(input),
            **kwargs,
        )
        if "messages" in input:
            for msg in input["messages"]:
                if msg["role"] == "HUMAN":
                    role = chat.role_human
                elif msg["role"] == "OBSERVATION":
                    role = chat.role_observation
                elif msg["role"] == "ASSISTANT":
                    role = chat.role_assistant
                else:
                    raise ValueError(f"不支持数据集中的 role: {msg['role']}")

                chat.append_message(role, msg["content"])

        elif "turns" in input:
            for turn in input["turns"]:
                if "HUMAN" in turn:
                    content = turn["HUMAN"]
                    chat.append_message(chat.role_human, content)
                if "OBSERVATION" in turn:
                    content = turn["OBSERVATION"]
                    chat.append_message(chat.role_observation, content)
                if "ASSISTANT" in turn:
                    content = turn["ASSISTANT"]
                    chat.append_message(chat.role_assistant, content)

        return chat

    @classmethod
    def from_pack(
        cls,
        packs: Dict[str, List[str]],
        name: str,
        prompt_style: Optional[PromptStyle] = None,
    ) -> list:
        """根据 pack 数据创建 Chat 对象.

        Params:
            packs: `dict`, pack 样本数据
                {
                    'inputs': ['xx', 'xx'],
                    'outputs': ['xx', 'xx'],
                }

            name: `str`, 符号系统名称
            prompt_style: `Optional[PromptStyle]`, 指定 prompt 风格, 默认使用和 name 一致的风格
        """
        chat = cls(name=name, prompt_style=prompt_style)
        packs = cls._format_packs(packs)

        sys_pattern = re.compile(chat.system_template.format(r"(.*?)"), re.DOTALL)
        turn_pattern = re.compile(chat.turn_start.format(r"(\d+)"), re.DOTALL)
        human_pattern = re.compile(
            chat.role_template.format(chat.role_human).strip(), re.DOTALL
        )
        observe_pattern = re.compile(
            chat.role_template.format(chat.role_observation).strip(), re.DOTALL
        )
        assistant_pattern = re.compile(
            chat.role_template.format(chat.role_assistant).strip(), re.DOTALL
        )

        chats = []
        for input, output in zip(packs["input"], packs["output"]):
            # system message
            sys_match = sys_pattern.search(input)
            if sys_match and sys_match.group(0):
                # system 指令只在首轮, 新增 chat 对象
                if len(chat.messages) > 0:
                    chats.append(chat)
                    chat = cls(name=name, prompt_style=prompt_style)

                input = input[sys_match.end() :]
                chat.system_message = sys_match.group(1)

            # turn start
            turn_match = turn_pattern.search(input)
            if turn_match and turn_match.group(0):
                # 当出现下一个轮次开始信息, 新增 chat 对象
                if name in ["antglm", "antglm_raw", "chatglm2"]:
                    round_start = 1
                else:
                    round_start = 0

                if all(
                    [
                        len(turn_match.groups()) > 0,
                        int(turn_match.group(1)) == round_start,
                        len(chat.messages) > 0,
                    ]
                ):
                    chats.append(chat)
                    chat = cls(name=name, prompt_style=prompt_style)

                input = input[turn_match.end() :]

            human_iter = human_pattern.finditer(input)
            observe_iter = observe_pattern.finditer(input)
            assistant_iter = assistant_pattern.finditer(input)
            human_match = next(human_iter, None)
            observe_match = next(observe_iter, None)
            assistant_match = next(assistant_iter, None)

            if not human_match and not observe_match:
                # 无 role format
                chat.append_message(chat.role_human, input)

            while human_match or observe_match:
                next_human_match = next(human_iter, None)
                next_observe_match = next(observe_iter, None)
                input = cls._append_human_observation(
                    chat,
                    input,
                    human_match=human_match,
                    next_human_match=next_human_match,
                    observe_match=observe_match,
                    next_observe_match=next_observe_match,
                    assistant_match=assistant_match,
                )

                human_match = next_human_match
                observe_match = next_observe_match
                next_human_match = next(human_iter, None)
                next_observe_match = next(observe_iter, None)

            if output:
                chat.append_message(chat.role_assistant, output)

        if chat.messages:
            chats.append(chat)

        return chats

    @classmethod
    def _append_human_observation(
        cls,
        chat,
        input: str,
        human_match: Optional[re.Match] = None,
        next_human_match: Optional[re.Match] = None,
        observe_match: Optional[re.Match] = None,
        next_observe_match: Optional[re.Match] = None,
        assistant_match: Optional[re.Match] = None,
    ) -> str:
        """给 chat 对象增加 human/observation message."""
        if observe_match:
            # observation 在 human 之后
            if observe_match.span()[0] > observe_match.span()[0]:
                human_str = input[observe_match.span()[1] : observe_match.span()[0]]
                observe_str = input[observe_match.span()[1] : assistant_match.span()[0]]
                chat.append_message(chat.role_human, human_str.strip())
                input_end = observe_match.span()[1]
                if observe_match.span()[0] < next_human_match.span()[0]:
                    chat.append_message(chat.role_observation, observe_str.strip())
                    input_end = assistant_match.span()[1]
            else:
                # observation 在 human 之前
                human_str = input[observe_match.span()[1] : assistant_match.span()[0]]
                observe_str = input[observe_match.span()[1] : observe_match.span()[0]]
                chat.append_message(chat.role_observation, observe_str.strip())
                input_end = observe_match.span()[1]
                if observe_match.span()[0] < next_observe_match.span()[0]:
                    chat.append_message(chat.role_human, human_str.strip())
                    input_end = assistant_match.span()[1]
        else:
            if assistant_match:
                human_str = input[human_match.span()[1] : assistant_match.span()[0]]
                input_end = assistant_match.span()[1]
            else:
                human_str = input[human_match.span()[1] :]
                input_end = len(input)
            chat.append_message(chat.role_human, human_str.strip())

        return input[input_end:]

    @classmethod
    def from_inout(
        cls,
        sample: Dict[str, str],
        name: str,
        prompt_style: Optional[PromptStyle] = None,
    ):
        """根据单样本创建一个 Chat 对象.

        Params:
            sample: `Dict[str, str]`, input/output 数据样本
                {
                    "input": "xxx",
                    "output": "xxx",
                }

            name: `str`, 符号系统名称
            prompt_style: `Optional[PromptStyle]`, 指定 prompt 风格, 默认使用和 name 一致的风格
        """
        chat = cls(name=name, prompt_style=prompt_style)
        input = sample["input"]
        output = sample["output"]

        sys_pattern = re.compile(chat.system_template.format(r"(.*?)"), re.DOTALL)
        turn_pattern = re.compile(chat.turn_start.format(r"(\d+)"), re.DOTALL)
        human_pattern = re.compile(
            chat.role_template.format(chat.role_human).strip(), re.DOTALL
        )
        observe_pattern = re.compile(
            chat.role_template.format(chat.role_observation).strip(), re.DOTALL
        )
        assistant_pattern = re.compile(
            chat.role_template.format(chat.role_assistant).strip(), re.DOTALL
        )

        # 去除轮次信息
        input = turn_pattern.sub("", input)

        # system message search
        sys_match = sys_pattern.search(input)
        if sys_match and sys_match.group(0):
            input = input[sys_match.end() :]
            chat.system_message = sys_match.group(1)

        human_iter = human_pattern.finditer(input)
        observe_iter = observe_pattern.finditer(input)
        assistant_iter = assistant_pattern.finditer(input)
        human_match = next(human_iter, None)
        observe_match = next(observe_iter, None)
        assistant_match = next(assistant_iter, None)
        next_human_match = next(human_iter, None)
        next_observe_match = next(observe_iter, None)

        while any(
            [
                human_match,
                observe_match,
                assistant_match,
            ]
        ):
            # human/observation 先后顺序可能不一样, 并且有可能有多个
            # 判断 assitant 之前是否还有 human/observation
            while any(
                [
                    human_match and human_match.span()[0] < assistant_match.span()[0],
                    observe_match
                    and observe_match.span()[0] < assistant_match.span()[0],
                    next_human_match
                    and next_human_match.span()[0] < assistant_match.span()[0],
                    next_observe_match
                    and next_observe_match.span()[0] < assistant_match.span()[0],
                ]
            ):
                if not input:
                    break

                cls._append_human_observation(
                    chat,
                    input,
                    human_match=human_match,
                    next_human_match=next_human_match,
                    observe_match=observe_match,
                    next_observe_match=next_observe_match,
                    assistant_match=assistant_match,
                )

                human_match = next_human_match
                observe_match = next_observe_match
                next_human_match = next(human_iter, None)
                next_observe_match = next(observe_iter, None)

            # assistant message
            if assistant_match and assistant_match.span():
                if observe_match:
                    if (
                        observe_match.span()
                        and observe_match.span()[0] < human_match.span()[0]
                    ):
                        assistant_str = input[
                            assistant_match.span()[1] : observe_match.span()[0]
                        ]
                elif human_match:
                    if human_match.span():
                        assistant_str = input[
                            assistant_match.span()[1] : human_match.span()[0]
                        ]
                else:
                    assistant_str = input[assistant_match.span()[1] :]

                if assistant_str:
                    chat.append_message(chat.role_assistant, assistant_str)

            assistant_match = next(assistant_iter, None)

        if output:
            chat.append_message(chat.role_assistant, output)

        return chat

    def __hash__(self):
        """数据对象的 hash 函数."""
        return hash(self.id)

    def __post_init__(self):
        """对象初始化后的处理, 处理包括:
        - 根据数据对象名称, 支持转成其他开源数据对象的基本信息
        """
        self.id = str(uuid.uuid4())
        if not self.messages:
            self.messages = []

        if not self.name and not self.prompt_style:
            logger.error(
                "构造 Chat 对象至少包含以下一个入参: `name/prompt_style`.\n\n"
                "`name` 支持以下 format 名称:"
            )
            logger.error(
                "\n".join([f"{k}: {v}" for k, v in self.support_names.items()])
            )
            logger.error("\n`prompt_style` 参考 antllm.data.chat_format.PromptStyle")
            raise ValueError

        if self.name == "antglm":
            # 默认 antglm 使用原始 antglm_raw - 第1轮\n用户: xx\n机器人: xx\n
            self.name = "antglm_raw"

        if not self.name and self.prompt_style == PromptStyle.ANTGLM_CHAT:
            logger.info(
                "Chat 对象入参没有 `name`, 默认使用 `ANTGLM_CHAT`, format:\n"
                f"role_human: {self.role_human}\n"
                f"role_assistant: {self.role_assistant}\n"
                f"role_observation: {self.role_observation}\n"
                f"role_template: {self.role_template}\n"
                f"turn_start: {self.turn_start}\n"
                f"human_end: {self.human_end}\n"
                f"assistant_start: {self.assistant_start}\n"
                f"assistant_end: {self.assistant_end}\n"
                f"assistant_end_ids: {self.assistant_end_ids}\n"
                f"general_role_end: {self.general_role_end}\n"
                f"tool_template: {self.tool_template}\n"
                f"code_template: {self.code_template}\n"
                f"arithemetic_templte: {self.arithemetic_templte}\n"
                f"\n入参 `name` 支持: ``"
            )
            return

        if self.name == "antglm_raw" or self.prompt_style == PromptStyle.ANTGLM_RAW:
            self.prompt_style = PromptStyle.ANTGLM_RAW
            self.role_template = "{}"
            self.role_human = "用户: "
            self.role_assistant = "机器人: "
            self.turn_start = "第{}轮\n"
            self.general_role_end = "\n"

        if (
            self.name in ["chatglm1", "chatglm2"]
            or self.prompt_style == PromptStyle.CHATGLM
        ):
            self.prompt_style = PromptStyle.CHATGLM
            self.role_template = "{}"
            self.role_human = "问："
            self.role_assistant = "答："
            self.turn_start = "[Round {}]\n"
            if self.name == "chatglm1":
                self.general_role_end = "\n"
            else:
                self.general_role_end = "\n\n"

        elif self.name == "chatglm3" or self.prompt_style == PromptStyle.CHATGLM3:
            self.prompt_style = PromptStyle.CHATGLM3
            self.system_template = "<|system|>\n {}"
            self.role_human = "<|user|>\n "
            self.role_assistant = "<|assistant|>\n "
            self.role_template = "{}"

        elif self.name == "llama2" or self.prompt_style == PromptStyle.LLAMA2:
            self.prompt_style = PromptStyle.LLAMA2
            self.role_template = "{}"
            self.system_template = "[INST] <<SYS>>\n{}\n<</SYS>>\n\n"
            self.role_human = "[INST] "
            self.role_assistant = "[/INST] "
            self.human_end = " "
            self.assistant_end = " </s><s>"

        elif self.name == "qwen":
            self.prompt_style = PromptStyle.CHATML
            self.role_template = "{}"
            self.system_template = "<|im_start|>system\n{}"
            if not self.system_message:
                self.system_message = "You are a helpful assistant."
            self.role_human = "<|im_start|>user\n"
            self.role_assistant = "<|im_start|>assistant\n"
            self.general_role_end = "<|im_end|>\n"

        elif self.name == "baichuan":
            self.prompt_style = PromptStyle.BAICHUAN2
            self.role_template = "{}"
            self.system_template = "{}"
            self.role_human = "<token_id-195>"
            self.role_assistant = "<token_id-196>"

        if not self.system_template:
            self.system_template = "{}"

    def readable_messages(self) -> str:
        """将 messages 输出为人类可读的字符串, 方便分析数据."""
        pass

    @property
    def prompt_str(self) -> str:
        """将 Chat 对象转成 prompt str, 合并 human/assitant 输出为 format 字符串."""
        return f"{self.prompt_inout['input']}{self.prompt_inout['output']}"

    @classmethod
    def _format_packs(cls, packs: Dict[str, List[str]]) -> Dict[str, List[str]]:
        """格式化 pack 样本, 输出相同 pack inputs, outputs 个数."""
        _packs = copy.deepcopy(packs)
        if len(_packs["input"]) - 1 == len(_packs["output"]):
            _packs["output"].append("")

        if len(_packs["input"]) != len(_packs["output"]):
            print(packs)
            raise ValueError(
                "输入 input 和 output 数量不匹配, "
                f"input num: {len(packs['input'])}, "
                f"output num: {len(packs['output'])}"
            )

        return _packs

    @property
    def prompt_inout(self) -> Dict[str, str]:
        """将 Chat 对象转成 input prompt, output prompt 字符串.

        Returns:
            `Dict[str, str]`, 示例:
            {
                "input": "<role>SYSTEM</role>xxxx<role>HUMAN</role>你好<role>ASSISTANT</role>你好，有什么可以帮您？<role>ASSISTANT</role>",  # noqa
                "output": "你好，有什么可以帮您？"
            }
        """
        packs = self._format_packs(self.prompt_pack)

        # 兼容逻辑
        if self.prompt_style == PromptStyle.ANTGLM_RAW:
            packs["input"] = [f"{item} " for item in packs["input"]]

        prompt_input = "".join(
            [f"{x}{y}" for x, y in zip(packs["input"][:-1], packs["output"][:-1])]
        )
        prompt_input += packs["input"][-1]
        prompt_output = packs["output"][-1]

        # 兼容逻辑
        if self.prompt_style == PromptStyle.ANTGLM_RAW:
            prompt_input = prompt_input.strip()

        return {
            "input": prompt_input,
            "output": prompt_output,
        }

    @property
    def prompt_pack(self) -> Dict[str, List[str]]:
        """将数据对象转成 pack input prompt, output prompt 字符串列表.:

        Returns:
            `Dict[str, List[str]]`, 示例:

            {
                "input": [
                    "<role>SYSTEM</role>xxxx<role>HUMAN</role>你好<role>ASSISTANT</role>",
                    "<role>HUMAN</role>讲个笑话<role>ASSISTANT</role>",
                    "<role>OBSERVATION</role>{\"weather\": \"晴\"}<role>ASSISTANT</role>"
                ],
                "output": [
                    "你好，有什么可以帮您？",
                    "笑话 1",
                    "今天天气 xxx"
                ]
            }

        """
        inputs = []
        outputs = []

        # 最开始 system 构造
        system_prompt = ""
        if self.system_message:
            system_prompt = self.system_template.format(self.system_message)

        if system_prompt:
            ret = system_prompt + self.general_role_end
        else:
            ret = ""

        # 有些 prompt style 单轮指令没有 format
        if self.prompt_style in [
            PromptStyle.ANTGLM_RAW,
            PromptStyle.ANTGLM_ONLY_MULTITURN_CHAT,
        ]:
            if len(self.messages) <= 2:
                output = ""
                for role, message in self.messages:
                    if role == self.role_assistant:
                        output = message
                    else:
                        input = ret + message
                return {
                    "input": [input],
                    "output": [output],
                }

        # 多轮对话
        if self.name in ["antglm_raw", "chatglm2"]:
            round_start = 1
        else:
            round_start = 0

        for i, (role, message) in enumerate(self.messages):
            # 轮次信息
            if self.name in ["antglm_raw", "chatglm1", "chatglm2"]:
                if i % 2 == 0:
                    ret += self.turn_start.format(i // 2 + round_start)

            # 角色 + 内容
            role_end = self.general_role_end
            if role == self.role_assistant and self.assistant_end:
                role_end = self.assistant_end
            elif self.human_end:
                role_end = self.human_end

            ret += self.role_template.format(role) + message + role_end

            if role == self.role_assistant:
                # output 只保留实际 assistant 内容
                if not message:
                    outputs.append("")
                else:
                    outputs.append(message + role_end)
                # input 需要连接 assistant role
                inputs[-1] += ret[: -len(message + role_end)]
            elif all(
                [
                    role == self.role_observation,
                    len(self.messages) > 1,
                    self.messages[i - 1][0] != self.role_assistant,
                ]
            ):
                # observation 之前不是 assistant, 需要将 observation 和上一个 input 连接一起
                continue
            else:
                inputs.append(ret)
            ret = ""

            # 最后一轮不是机器人回复, 需要拼接机器人 role, 用于模型生成
            if i == len(self.messages) - 1 and role != self.role_assistant:
                inputs[-1] += self.role_template.format(self.role_assistant).strip()

        # 兼容逻辑, 去除 inputs 最后空格符号
        if self.prompt_style == PromptStyle.ANTGLM_RAW:
            inputs = [item.strip() for item in inputs]

        return {
            "input": inputs,
            "output": outputs,
        }

    @property
    def turns_num(self) -> int:
        """和机器人的交互轮数, 以用户输出多少次为轮数个数."""
        return sum([1 if msg[0] == self.role_human else 0 for msg in self.messages])

    def to_json(self) -> dict:
        """输出 chat json dict 格式, 包含不同角色和机器人交互的每轮信息.

        Returns
            `List[dict]`, {
                "id": "xx",
                "messages": [
                    {"role": "HUMAN", "content": "xxx"}
                ]
                "turns": [
                    {"HUMAN": "xx", "OBSERVATION": "xx", "ASSISTANT": "xx"}
                ]
            }
        """
        turns = []
        messages = []
        turn = {}
        for msg in self.messages:
            if msg[0] == self.role_assistant:
                messages.append({"role": "ASSISTANT", "content": msg[1]})
                turn["ASSISTANT"] = msg[1]
                turns.append(turn)
                turn = {}

            if msg[0] == self.role_human:
                messages.append({"role": "HUMAN", "content": msg[1]})
                turn["HUMAN"] = msg[1]

            if msg[0] == self.role_observation:
                messages.append({"role": "OBSERVATION", "content": msg[1]})
                turn["OBSERVATION"] = msg[1]

        if self.messages[-1][0] == self.role_human:
            messages.append({"role": "ASSISTANT", "content": ""})
            turn["ASSISTANT"] = ""
            turns.append(turn)

        result = self.origin_json or {}
        result.update(
            {
                "id": self.id,
                "name": self.name,
                "source": self.source,
                "lang": self.lang,
                "topic": self.topic,
                "system_template": self.system_template,
                "system_message": self.system_message,
                "turns": turns,
                "messages": messages,
            }
        )

        return result

    def set_system_message(self, system_message: str):
        """Set the system message."""
        self.system_message = system_message

    def append_message(self, role: str, message: str):
        """Append a new message."""
        if not message:
            message = ""
        self.messages.append([role, message])

    def to_openai_api_messages(self) -> List[dict]:
        """Convert the conversation to OpenAI chat completion format."""
        ret = [{"role": "system", "content": self.system_message}]

        for i, (_, msg) in enumerate(self.messages[self.offset :]):
            if i % 2 == 0:
                ret.append({"role": "user", "content": msg})
            else:
                if msg is not None:
                    ret.append({"role": "assistant", "content": msg})
        return ret

    def copy(self):
        return copy.deepcopy(self)