from ..utils.session import Session
from ._base import Prompt


class DeepSeekLLaMA3Prompt(Prompt):
    """DeepSeek R1 LLaMA 3.1 prompt that follows this structure
    <｜begin▁of▁sentence｜>{system_prompt}<｜User｜>{prompt}<｜Assistant｜>{assistant}<｜end▁of▁sentence｜><｜User｜>{prompt}<｜Assistant｜>{assistant}<｜end▁of▁sentence｜>

    Args:
        system (str): system prompt

    """

    PROMPT_START = "<｜begin▁of▁sentence｜>"
    TEXT_END = "<｜end▁of▁sentence｜>"
    USER = "<｜User｜>"
    ASSISTANT = "<｜Assistant｜>"

    def __init__(self, system: str) -> None:
        self.system = system

    def prepare(self, session: Session) -> str:
        """Prepare prompt for model input

        Args:
            session (Session): dialog session

        Returns:
            str: model input prompt
        """

        """
        <｜begin▁of▁sentence｜>{system_prompt}<｜User｜>{prompt}<｜Assistant｜>{assistant}<｜end▁of▁sentence｜><｜User｜>{prompt}<｜Assistant｜>{assistant}<｜end▁of▁sentence｜>
        """

        prompt = f"{self.PROMPT_START}{self.system}"
        for qa in session.history:
            # qa --> {"question": ..., "answer": ...}
            prompt += f"{self.USER}{qa['question']}{self.ASSISTANT}"
            if qa["answer"] is None:
                break
            else:
                prompt += f"{qa['question']}{self.TEXT_END}"

        return prompt
