from pathlib import Path
from textwrap import dedent
from typing import Literal

from inspect_ai import Task, task
from inspect_ai.solver import Solver, basic_agent, system_message
from inspect_ai.tool import computer
from inspect_ai.util import SandboxEnvironmentType

from .dataset import osworld_dataset
from .scorer import osworld_scorer

COMPOSE_FILE = Path(__file__).parent / "compose.yaml"
DEFAULT_DOCKER_SANDBOX = ("docker", COMPOSE_FILE.as_posix())
DEFAULT_OSWORLD_MESSAGE_LIMIT = 75


@task
def osworld(
    solver: Solver | None = None,
    corpus: Literal["all", "small"] | str = "all",
    include_apps: str | list[str] | None = None,
    exclude_apps: str | list[str] | None = None,
    max_screenshots: int | None = 1,
    max_attempts: int = 1,
    message_limit: int | None = DEFAULT_OSWORLD_MESSAGE_LIMIT,
    computer_timeout: int | None = 180,
    sandbox: SandboxEnvironmentType = DEFAULT_DOCKER_SANDBOX,
    include_connected: bool | None = None,
) -> Task:
    """
    Creates and returns a Task object configured for the osworld evaluation.

    Args:
      solver: Provide a custom solver (if not specified uses Inspect's basic_agent with computer tool).
      corpus (Literal["all", "small"] | str, optional): Specifies the corpus to use. Defaults to "all".
        "all": Uses the full OSWorld corpus.
        "small": Uses the small OSWorld corpus.
        any other str will be interpreted as a file path to a custom corpus file.
      include_apps (str | list[str] | None, optional): Specifies the apps to filter the dataset.
            Can be a single app (str), a list of apps (list[str]), or None to include all apps.
      exclude_apps (str | list[str] | None, optional): Specifies the apps to exclude from the dataset.
            Can be a single app (str), a list of apps (list[str]), or None to exclude no apps.
      max_screenshots (int): Maximum number screenshots returned by the tool to be retained in the message history. Defaults to None.
      max_attempts (int): Maximum number of submissions to accept before terminating. Defaults to 1.
      message_limit (int | None): Limit on messages in sample before terminating agent. Defaults to None.
      computer_timeout (int | None): Limit on the amount of time the computer tool can take to execute a command.  If not specified,
        will use timeout defined by the tool.
      sandbox: Sandbox environment to use for the task.
      include_connected (bool | None, optional): Specifies whether to include samples that require network access.
            Defaults to None. When None or False, those samples will not be included.

    Returns:
      Task: A Task object configured with the specified dataset, solver, scorer, and sandbox environment.
    """
    return Task(
        dataset=osworld_dataset(
            corpus=corpus,
            include_apps=include_apps,
            exclude_apps=exclude_apps,
            include_connected=include_connected,
        ),
        solver=solver
        or osworld_basic_agent(
            max_attempts,
            message_limit,
            computer_timeout=computer_timeout,
            max_screenshots=max_screenshots,
        ),
        scorer=osworld_scorer(),
        sandbox=sandbox,
    )


@task
def osworld_small(
    solver: Solver | None = None,
    include_apps: str | list[str] | None = None,
    exclude_apps: str | list[str] | None = None,
    max_screenshots: int | None = 1,
    max_attempts: int = 1,
    message_limit: int | None = DEFAULT_OSWORLD_MESSAGE_LIMIT,
    computer_timeout: int | None = 180,
    sandbox: SandboxEnvironmentType = DEFAULT_DOCKER_SANDBOX,
    include_connected: bool | None = None,
) -> Task:
    """
    Creates and returns a Task object configured for the osworld evaluation.

    Args:
      solver: Provide a custom solver (if not specified uses Inspect's basic_agent with computer tool).
      include_apps (str | list[str] | None, optional): Specifies the apps to filter the dataset.
            Can be a single app (str), a list of apps (list[str]), or None to include all apps.
      exclude_apps (str | list[str] | None, optional): Specifies the apps to exclude from the dataset.
            Can be a single app (str), a list of apps (list[str]), or None to exclude no apps.
      max_screenshots (int): Maximum number screenshots returned by the tool to be retained in the message history. Defaults to None.
      max_attempts (int): Maximum number of submissions to accept before terminating. Defaults to 1.
      message_limit (int | None): Limit on messages in sample before terminating agent. Defaults to None.
      computer_timeout (int | None): Limit on the amount of time the computer tool can take to execute a command.  If not specified,
        will use timeout defined by the tool.
      sandbox: Sandbox environment to use for the task.
      include_connected (bool | None, optional): Specifies whether to include samples that require network access.
            Defaults to None. When None or False, those samples will not be included.

    Returns:
      Task: A Task object configured with the specified dataset, solver, scorer, and sandbox environment.
    """
    return osworld(
        solver=solver,
        corpus="small",
        include_apps=include_apps,
        exclude_apps=exclude_apps,
        max_screenshots=max_screenshots,
        max_attempts=max_attempts,
        message_limit=message_limit,
        computer_timeout=computer_timeout,
        sandbox=sandbox,
        include_connected=include_connected,
    )


SHARED_HOW_TO_FINISH_MESSAGE = """
                    If you think the task is done, submit '```DONE```' with the ```submit()``` tool.

                    If you think the task cannot be done, submit '```FAIL <reason for failure>```' with the ```submit()``` tool.  Don't submit ```FAIL``` without first trying your best to complete the task.

                    If you are unaware of what the screen looks like, you can ask for a screenshot with the ```computer()``` tool.
                   """
CONTINUE_MESSAGE = f"""
                    {SHARED_HOW_TO_FINISH_MESSAGE}

                    Please proceed to the next step using your best judgement.
                    """
SYSTEM_MESSAGE = f"""
                    You are an agent that follows my instructions and performs desktop computer tasks as instructed.

                    You have good knowledge of computers and a good internet connection, and you can control a computer using a mouse and keyboard via the computer tool.

                    For each step, you will get an an image the computer screen, and you will predict the action of the computer based on the image.

                    IMPORTANT: When placing the mouse inside a text field, ensure that you've selected the existing text if your intent is to replace it.

                    IMPORTANT: To launch applications from desktop icons, you must double click rather than single click.

                    {SHARED_HOW_TO_FINISH_MESSAGE}
                    """


def osworld_basic_agent(
    max_attempts: int,
    message_limit: int | None = 75,
    computer_timeout: int | None = 180,
    max_screenshots: int | None = 1,
) -> Solver:
    return basic_agent(
        init=system_message(dedent(SYSTEM_MESSAGE)),
        tools=[computer(max_screenshots=max_screenshots, timeout=computer_timeout)],
        max_attempts=max_attempts,
        message_limit=message_limit,
        continue_message=dedent(CONTINUE_MESSAGE),
    )
