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

from inspect_ai import Task, task
from inspect_ai.solver import Solver, basic_agent, system_message
from inspect_ai.tool import bash, python, web_browser
from inspect_ai.util._sandbox.environment import SandboxEnvironmentType

from .dataset import gaia_dataset
from .scorer import gaia_scorer

TASK_DIR = Path(__file__).parent
COMPOSE_FILE = TASK_DIR / "compose.yaml"
VALUES_FILE = TASK_DIR / "values.yaml"

DEFAULT_DOCKER_SANDBOX = ("docker", COMPOSE_FILE.as_posix())
DEFAULT_K8S_SANDBOX = ("k8s", VALUES_FILE.as_posix())


@task
def gaia(
    solver: Solver | None = None,
    input_prompt: str | None = None,
    max_attempts: int = 1,
    max_messages: int = 100,
    subset: Literal[
        "2023_all", "2023_level1", "2023_level2", "2023_level3"
    ] = "2023_all",
    split: Literal["test", "validation"] = "validation",
    instance_ids: str | list[str] | None = None,
    sandbox: SandboxEnvironmentType = DEFAULT_DOCKER_SANDBOX,
) -> Task:
    """GAIA task.

    Args:
      solver: Provide a custom solver (if not specified uses Inspect's
      basic_agent with bash, python, and web browsing tools).
      input_prompt: Per sample question prompt. Should include a {file}
      variable (for identifying any files relevant to the task) and a {question} variable for rendering the question.
      max_attempts: Maximum number of submission attempts (defaults to 1). Only applies when using the default solver.
      max_messages: Maximum number of messages before giving up (defaults to 100). Only applies when using the default solver.
      subset: Which GAIA subset to evaluate (defaults to 2023_all).
      split: Which split to evaluate ("validation" or "test")
      instance_ids: Specific question instances to evaluated.
      sandbox: Sandbox environment to use for the task.

    Returns:
      GAIA Inspect task.
    """
    # read dataset
    dataset = gaia_dataset(
        input_prompt=input_prompt,
        subset=subset,
        split=split,
    )

    # filter by instance id if requested
    if instance_ids:
        instance_ids = [instance_ids] if isinstance(instance_ids, str) else instance_ids
        dataset = dataset.filter(lambda x: x.id in instance_ids)

    # resolve solver
    solver = solver or default_solver(max_attempts, max_messages)

    # resolve scorer (test split has no answers)
    scorer = gaia_scorer() if split == "validation" else None

    # return task
    return Task(
        dataset=dataset,
        plan=solver,
        scorer=scorer,
        sandbox=sandbox,
    )


@task
def gaia_level1(**kwargs: Any) -> Task:
    return gaia(subset="2023_level1", **kwargs)


@task
def gaia_level2(**kwargs: Any) -> Task:
    return gaia(subset="2023_level2", **kwargs)


@task
def gaia_level3(**kwargs: Any) -> Task:
    return gaia(subset="2023_level3", **kwargs)


def default_solver(
    max_attempts: int, max_messages: int, code_timeout: int = 180
) -> Solver:
    return basic_agent(
        init=system_message(
            dedent(
                """
                You are an autonomous AI agent, consisting of a large language model and a scaffolding program that runs on a linux VM.

                You have full control over your VM through the scaffolding program's provided functions, including: writing and editing files, running bash commands, installing packages etc.. From your VM, you can access the internet.

                You are state-of-the-art, and excellent at analysis, planning and programming. You possess encyclopedic knowledge of the world. You write in a very direct and concise style.

                You are resourceful and adaptive, and you never give up. Before deciding something can't be done, you try it out. You consider multiple options and choose the best one. If your current approach doesn't work, you formulate a new plan. You are given a task you need to solve completely on your own.

                Please think step by step before calling tools. When you are ready to answer, use the submit tool to provide your final answer.
                """
            )
        ),
        tools=[bash(code_timeout), python(code_timeout)] + web_browser(),
        max_attempts=max_attempts,
        max_messages=max_messages,
    )
