# %%
from typing import Optional, List
from dataclasses import field, dataclass


@dataclass
class OpusV1Turn:
    role: str
    content: str
    names: List[str] = field(default_factory=list)
    # If set to true, will not append <|im_end|>, so the model will continue the turn.
    # In RP you can for example use the following to force a specific character response:
    # role="text"
    # names=["Jack"]
    # open="true"
    open: bool = False


@dataclass
class OpusV1Character:
    name: str
    description: str


@dataclass
class OpusV1StorySystemPrompt:
    format: str = "prose"
    plot_description: str = ""
    style_description: str = ""
    characters: List[OpusV1Character] = field(default_factory=list)


@dataclass
class OpusV1Prompt:
    story: Optional[OpusV1StorySystemPrompt] = None
    turns: List[OpusV1Turn] = field(default_factory=list)


def format_opus_v1_prompt(prompt) -> str:
    turns = prompt.turns
    if prompt.story is not None:
        system = format_opus_v1_system_prompt(prompt.story)
        turns = [OpusV1Turn(role="system", content=system)] + turns

    parts = []
    for i, turn in enumerate(turns):
        assert turn.role in ["user", "text", "system", "assistant"]
        assert turn.role != "system" or i == 0

        is_last = i == len(turns) - 1
        open = is_last and turn.open
        parts.append(format_turn(turn.role, turn.content, turn.names, open=open))
    return "".join(parts)


def format_turn(
    role: str, content: str, names: List[str] = [], open: bool = False
) -> str:
    im_start = "<|im_start|>"
    im_end = "<|im_end|>"

    body = im_start + role
    if len(names) > 0:
        body += f" names= {'; '.join(names)}"

    body += "\n"
    if open:
        return body + content.lstrip()
    else:
        return body + content.strip() + im_end + "\n"


def format_opus_v1_system_prompt(prompt) -> str:
    format_text = "story" if prompt.format == "prose" else "role-play"
    system = f"""
You are an intelligent, skilled, versatile writer.

Your task is to write a {format_text} based on the information below.

Write the {format_text} as if it's a book.
    """.strip()

    if len(prompt.plot_description) > 0:
        system += "\n\n\n## Plot description:\n\n"
        system += prompt.plot_description.strip()
    if len(prompt.style_description) > 0:
        system += "\n\n\n## Style description:\n\n"
        system += prompt.style_description.strip()
    if len(prompt.characters) > 0:
        system += "\n\n\n## Characters:\n\n"
        for character in prompt.characters:
            system += f"### {character.name}\n\n"
            system += character.description.strip()
            system += "\n\n"

    return system.strip()