from inspect_ai.model import (
    ChatMessage,
    ChatMessageAssistant,
    ChatMessageUser,
)

CLOSED_BOOK_SYSTEM_PROMPT = """
You are a knowledgeable AI assistant tasked with providing concise answers to a wide range of questions.  Always provide an answer, even if you're not entirely certain. If you're unsure, make an educated guess based on your knowledge.

Follow these guidelines:
1. Don't ask the user any questions.
2. Structure your response as a series of intermediate questions and answers, followed by a final answer.
3. Do not generate the same intermediate question twice.
4. Keep the final answer concise. It should be either: a number, a string, a list of strings, or a list of JSONs.
5. Ensure the answer can be parsed with the Python method: json.loads(input_str).
6. If you truly can't provide any answer, use an empty string as the final answer.

This is a closed-book task, so rely on your existing knowledge without asking for additional information.

Use the format:
Intermediate question: [Your question]
Intermediate answer: [Your answer]
(Repeat as needed)
The final answer is: [Your concise final answer]

Remember, it's better to provide a potentially outdated or slightly inaccurate answer than no answer at all. Your goal is to maximize the answer rate and accuracy to the best of your ability. Always provide at least one intermediate question and answer before giving the final answer.
""".strip()

CLOSED_BOOK_ONE_SHOT_TASK = (
    "Which member of the famous Beckham couple has the most Instagram followers?"
)

CLOSED_BOOK_ONE_SHOT_ANSWER = """
Intermediate question: Who are the members of the famous Beckham couple?
Intermediate answer: The members of the Beckham couple are Victoria Beckham and David Beckham.
Intermediate question: How many followers does Victoria Beckham have on Instagram?
Intermediate answer: Victoria Beckham has 32.8M followers on Instagram.
Intermediate question: How many followers does David Beckham have on Instagram?
Intermediate answer: David Beckham has 87.1M followers on Instagram.
The final answer is: David Beckham
""".strip()

WEB_SEARCH_REQUEST_PREFIX = "Intermediate question:"
WEB_SEARCH_ANSWER_PREFIX = "The final answer is:"

# The "The format will be" section is not in the original paper prompt.
# It was found models would not follow the required structure in zero-shot without it.
OPEN_BOOK_SYSTEM_PROMPT = """
You are a helpful agent that has to provide answers to complex tasks.

Tasks can require accessing the web. To access the web, you are equipped with a search engine.
To access the search engine, simply generate an intermediate question. The intermediate question will be sent to a search engine (Google) and the result from the search engine will be provided by the user in the next turn.

It is best to access the search engine with simple queries, as they are most likely to return the correct answers. For example, if you require to find the number of Instagram followers for a group, generate a query for each member of the group.

When the user message ends with an intermediate question, always generate an intermediate answer.

If the search results do not answer your intermediate question, begin the intermediate answer with "Intermediate answer: The intermediate answer cannot be inferred from the search results".
In that case, you can use your own knowledge to answer the intermediate question or generate a new intermediate question.

You can generate additional intermediate questions and answers until arriving at the final answer. Do not generate the same intermediate question twice.

When arriving at the final answer, start your prediction with: "{final_answer_prefix}" followed by the answer to the task.
If you are unsure what the final answer is, generate the empty string ("").

Keep the answer as concise as possible. The answer should either be: a number, a string, a list of strings, or a list of jsons.
The answer should be parsed with the python method: json.loads(input_str). If no answer is found, generate an empty string.

The format will be:
Assistant: Intermediate question: [Your search engine query]
User:      Context0: [Search engine results]
           Task: [Original task]
           Intermediate question: [Your most recent search engine query]
Assistant: Intermediate answer: [Your answer as inferred from the search results or your own knowledge]
           Intermediate question: [Your next search engine query]
(Repeat as needed)
Assistant: Intermediate answer: [Your answer as inferred from the search results or your own knowledge]
           The final answer is: [Your next search engine query]
""".strip().format(final_answer_prefix=WEB_SEARCH_ANSWER_PREFIX)


WEB_SEARCH_RESPONSE = """Context{count}: {search_results}
Task: {original_task}
Intermediate question: {intermediate_question}
""".strip()

_WEB_SEARCH_ASSISTANT_EXAMPLE = """
Intermediate answer: {intermediate_answer}
Intermediate question: {intermediate_question}
""".strip()

_WEB_SEARCH_ONE_SHOT_TASK = (
    "Which member of the famous Beckham couple has the most Instagram followers?"
)

_WEB_SEARCH_ONE_SHOT_Q0 = "Who are the members of the famous Beckham couple?"
_WEB_SEARCH_ONE_SHOT_R0 = "Posh and Becks: Posh and Becks is a widely used nickname for the British celebrity supercouple Victoria Beckham (née Adams, 'Posh Spice' of the Spice Girls) and David Beckham (a footballer and former England captain)."
_WEB_SEARCH_ONE_SHOT_A0 = (
    "The members of the Beckham couple are Victoria Beckham and David Beckham."
)

_WEB_SEARCH_ONE_SHOT_Q1 = "How many followers does Victoria Beckham have on Instagram?"
_WEB_SEARCH_ONE_SHOT_R1 = "Victoria Beckham has 32.8M followers on Instagram."
_WEB_SEARCH_ONE_SHOT_A1 = "Victoria Beckham has 32.8M followers on Instagram."

_WEB_SEARCH_ONE_SHOT_Q2 = "How many followers does David Beckham have on Instagram?"
_WEB_SEARCH_ONE_SHOT_R2 = "David Beckham has 87.1M followers on Instagram."
_WEB_SEARCH_ONE_SHOT_A2 = "David Beckham has 87.1M followers on Instagram."

WEB_SEARCH_ONE_SHOT_PROMPTS: list[ChatMessage] = [
    ChatMessageUser(content=_WEB_SEARCH_ONE_SHOT_TASK),
    ChatMessageAssistant(content=f"Intermediate question: {_WEB_SEARCH_ONE_SHOT_Q0}"),
    ChatMessageUser(
        content=WEB_SEARCH_RESPONSE.format(
            count=0,
            search_results=_WEB_SEARCH_ONE_SHOT_R0,
            original_task=_WEB_SEARCH_ONE_SHOT_TASK,
            intermediate_question=_WEB_SEARCH_ONE_SHOT_Q0,
        )
    ),
    ChatMessageAssistant(
        content=_WEB_SEARCH_ASSISTANT_EXAMPLE.format(
            intermediate_answer=_WEB_SEARCH_ONE_SHOT_A0,
            intermediate_question=_WEB_SEARCH_ONE_SHOT_Q1,
        )
    ),
    ChatMessageUser(
        content=WEB_SEARCH_RESPONSE.format(
            count=1,
            search_results=_WEB_SEARCH_ONE_SHOT_R1,
            original_task=_WEB_SEARCH_ONE_SHOT_TASK,
            intermediate_question=_WEB_SEARCH_ONE_SHOT_Q1,
        )
    ),
    ChatMessageAssistant(
        content=_WEB_SEARCH_ASSISTANT_EXAMPLE.format(
            intermediate_answer=_WEB_SEARCH_ONE_SHOT_A1,
            intermediate_question=_WEB_SEARCH_ONE_SHOT_Q2,
        )
    ),
    ChatMessageUser(
        content=WEB_SEARCH_RESPONSE.format(
            count=2,
            search_results=_WEB_SEARCH_ONE_SHOT_R2,
            original_task=_WEB_SEARCH_ONE_SHOT_TASK,
            intermediate_question=_WEB_SEARCH_ONE_SHOT_Q2,
        )
    ),
    ChatMessageAssistant(
        content=f"Intermediate answer: {_WEB_SEARCH_ONE_SHOT_A2}\nThe final answer is: David Beckham."
    ),
]

WEB_BROWSER_SYSTEM_PROMPT = """
You are a knowledgeable AI assistant tasked with providing concise answers to a wide range of questions.  Always provide an answer, even if you're not entirely certain. If you're unsure, make an educated guess based on your knowledge.

Follow these guidelines:
1. Don't ask the user any questions.
2. Keep the final answer concise. It should be either: a number, a string, a list of strings, or a list of JSONs.
3. Only provide a more complex answer if the task requires it, don't provide related information unless needed.
4. Ensure the answer can be parsed with the Python method: json.loads(input_str).
5. If you truly can't provide any answer, use an empty string as the final answer.

This is an open book task.  The tasks will require accessing websites for which you can use a Playwright based browser with the ability to access and navigate URLs.  For convenience you are also provided with access to a search engine where you can directly input a search query and the top results will be returned to you.  You may also use your existing knowledge.

Please think step by step before calling tools. When you are ready to answer, use the submit tool to provide your final answer.

Remember, it's better to provide a potentially outdated or slightly inaccurate answer than no answer at all. Your goal is to maximize the answer rate and accuracy to the best of your ability.
""".strip()
