import json
from pathlib import Path

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

TASK_PROMT = r"""'''
{final_html}
'''

Based on the HTML webpage above, try to complete the following task:
Task: {confirmed_task}

Previous actions:
{previous_actions}

What should be the next action? Please select from the following choices
(If the correct action is not in the page above, please select A. 'None of the above'):

{choices_text}
"""


def load_llm_prompt() -> list[ChatMessage]:
    """Load LLM prompt from JSON file.

    This function loads the few-shot examples and system instructions from
    the llm_prompt.json file and converts them into ChatMessage objects.
    These messages are used to provide context and examples to the language
    model during the Mind2Web evaluation.

    Returns:
        A list of ChatMessage objects containing system instructions and
        few-shot examples for the Mind2Web task.
    """
    prompt_path = Path(__file__).parent / "llm_prompt.json"
    with open(prompt_path) as f:
        data = json.load(f)

    messages: list[ChatMessage] = []
    for msg in data:
        if msg["role"] == "system":
            messages.append(ChatMessageSystem(content=msg["content"], source="input"))
        elif msg["role"] == "user":
            messages.append(ChatMessageUser(content=msg["content"], source="input"))
        elif msg["role"] == "assistant":
            messages.append(
                ChatMessageAssistant(content=msg["content"], source="input")
            )
    return messages
