import logging

from inspect_ai import Task, task
from inspect_ai.model import (
    ChatMessage,
)

from inspect_evals.mind2web.dataset import Mind2WebSplit, get_dataset
from inspect_evals.mind2web.prompts import TASK_PROMT, load_llm_prompt
from inspect_evals.mind2web.scorer import mind2web_scorer
from inspect_evals.mind2web.solver import solve

logger = logging.getLogger(__name__)


@task()
def mind2web(split: str = "train") -> Task:
    """Mind2Web evaluation task using few-shot prompting.

    This function creates a Task for evaluating language models on the Mind2Web
    dataset, which tests the ability to follow natural language instructions to
    complete tasks on websites. The evaluation uses a few-shot prompting approach
    with examples loaded from llm_prompt.json.

    Args:
        split: Dataset split to use for evaluation. Defaults to CROSS_DOMAIN.
               - CROSS_TASK: Tests generalization to new tasks on websites seen during training
               - CROSS_WEBSITE: Tests generalization to new websites within domains seen during training
               - CROSS_DOMAIN: Tests generalization to entirely new domains and websites

    Returns:
        A Task object configured for the Mind2Web evaluation.
    """
    logger.info("\nStarting Mind2Web evaluation...")
    messages: list[ChatMessage] = load_llm_prompt()

    return Task(
        name="mind2web_eval",
        dataset=get_dataset(split=Mind2WebSplit(split)),
        solver=[solve(TASK_PROMT, messages=messages)],
        scorer=[mind2web_scorer()],
    )
