from procoder.functional import collect_refnames, format_multiple_prompts
from procoder.functional import indent4 as indent
from procoder.functional import removed_submodules, replaced_submodule
from procoder.prompt import *
from safeagi.prompts.agent.shared import *
from safeagi.prompts.globals import *
from safeagi.utils import get_num_tokens, make_colorful

AGENT_NAIVE_SYSTEM_INFO = AGENT_SYSTEM_INFO

AGENT_NAIVE_PROMPT = (
    Collection(
        AGENT_ENV_SETUP, AGENT_TASK_DESC, AGENT_FORMAT_INSTRUCTION, AGENT_TASK_BEGIN
    )
    .set_indexing_method(sharp2_indexing)
    .set_sep("\n\n")
)


AGENT_NAIVE_PROMPT = replace_agent_ref_with_pronoun(AGENT_NAIVE_PROMPT)
AGENT_NAIVE_CLAUDE_PROMPT = replaced_submodule(
    AGENT_NAIVE_PROMPT, "task_begin", AGENT_TASK_BEGIN_FOR_CLAUDE
)

if __name__ == "__main__":
    inputs = dict(toolkit_descriptions=None, tool_names=None)
    add_refnames(AGENT_DUMMY_VARS, inputs, include_brackets=False)
    system_prompt, example_prompt = format_multiple_prompts(
        [AGENT_NAIVE_SYSTEM_INFO, AGENT_NAIVE_PROMPT,],
        inputs,
        include_brackets=[False, True],
    )
    print(make_colorful("system", system_prompt))
    print(make_colorful("human", example_prompt))
    print("\n\n>>>>Token lengths:", get_num_tokens(example_prompt))
