import json
from pathlib import Path
from typing import Literal

from inspect_ai.dataset import Dataset, MemoryDataset, Sample
from platformdirs import user_cache_dir

from .container.code.osworld_types import Example, ExampleCorpusModel
from .container.code.util import to_getter_list
from .sparse_clone import get_sparse_repo


def osworld_dataset(
    *,
    corpus: Literal["all", "small"] | str,
    include_apps: str | list[str] | None,
    exclude_apps: str | list[str] | None,
    include_connected: bool | None = None,
) -> Dataset:
    """
    Creates and returns a Dataset the osworld evaluation.

    Args:
      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.
      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:
      Dataset: A Dataset object configured with the specified set of Samples.
    """
    include_apps = [include_apps] if isinstance(include_apps, str) else None
    exclude_apps = (
        unsupported_apps
        if exclude_apps is None
        else exclude_apps + unsupported_apps
        if isinstance(exclude_apps, list)
        else [exclude_apps]
    )
    cache_dir_root = Path(user_cache_dir("inspect_evals")) / "osworld"
    local_osworld_repo_path = cache_dir_root / "repo"

    get_sparse_repo(
        remote_repo_url=OSWORLD_REPO_URL,
        repo_path=local_osworld_repo_path,
        sparse_root=EVALUATION_EXAMPLES_ROOT,
    )
    evaluation_examples_root = Path(local_osworld_repo_path) / EVALUATION_EXAMPLES_ROOT
    corpus_filepath = (
        evaluation_examples_root / "test_all.json"
        if corpus == "all"
        else evaluation_examples_root / "test_small.json"
        if corpus == "small"
        else Path(corpus)
    )

    with open(corpus_filepath, "r", encoding="utf-8") as f:
        corpus_json = ExampleCorpusModel(**json.loads(f.read())).root.items()

    examples = [
        (
            example_sub_dir,
            file_path
            := f"{evaluation_examples_root}/examples/{example_sub_dir}/{guid}.json",
            _example_from_file_path(file_path),
        )
        for example_sub_dir, guids in corpus_json
        for guid in guids
    ]

    if not examples:
        raise RuntimeError("No samples selected for dataset")
    return MemoryDataset(
        [
            _sample_from_example(example_file_path, example)
            for example_sub_dir, example_file_path, example in examples
            if _should_include_example(
                example,
                example_sub_dir,
                include_apps,
                exclude_apps,
                include_connected,
            )
        ]
    )


def _should_include_example(
    example: Example,
    example_sub_dir: str | None,
    include_apps: list[str] | None,
    exclude_apps: list[str],
    include_connected: bool | None = None,
) -> bool:
    return (
        # some examples are just not yet supported
        example.id not in unsupported_examples
        # exclude examples that require network unless the user overrode
        and (include_connected or False or example.id not in examples_requiring_network)
        # if include_apps is specified, and there are any related_apps that are
        # not in the list, exclude the example
        and (
            not include_apps
            or all(app in example.related_apps for app in include_apps)
            # this is a hack to support multi_apps as an "app"
            or (example_sub_dir == "multi_apps" and "multi_apps" in include_apps)
        )
        # if any of the related_apps is in the exclude list, exclude
        and not any(app in example.related_apps for app in exclude_apps)
    )


def _example_from_file_path(file_path: str) -> Example:
    with open(file_path, "r", encoding="utf-8") as file:
        return Example(**json.loads(file.read()))


def _sample_from_example(example_file_path: str, example: Example) -> Sample:
    return Sample(
        example.instruction,
        id=example.id,
        files={
            CONTAINER_EXAMPLE_PATH: example_file_path,
            **(_files_from_download_configs(example)),
            **(_files_from_cloud_file_getters(example)),
        },
        setup=CONFIG_COMMAND,
        metadata={
            "snapshot": example.snapshot,
            "source": example.source,
        },
    )


def _files_from_download_configs(example: Example) -> dict[str, str]:
    """
    Generate a dictionary of file paths and their corresponding URLs from the download `Config`'s of an example.

    `Config`'s can be specified in the `Example`'s `config:` property or in the `evaluator/postconfig` property.

    Args:
      example (Example): An instance of the `Example` class containing configuration and evaluator information.

    Returns:
      dict[str, str]: A dictionary where the keys are file paths and the values are the corresponding URLs.
    """
    all_configs = (example.config or []) + (example.evaluator.postconfig or [])

    return {
        file.path: file.url
        for config in [config for config in all_configs if config.type == "download"]
        for file in config.parameters.files
    }


def _files_from_cloud_file_getters(example: Example) -> dict[str, str]:
    """
    Generate a dictionary of file paths and their corresponding URLs from any `cloud_file` getter's in the example.

    Args:
      example (Example): An instance of the `Example` class containing evaluator information.

    Returns:
      dict[str, str]: A dictionary where the keys are file paths and the values are the corresponding URLs.
    """
    cloud_file_getters = [
        getter
        for getter in to_getter_list(example.evaluator.expected)
        if getter.type == "cloud_file"
    ]

    single_getters = [g for g in cloud_file_getters if g.multi is not True]
    multi_getters = [g for g in cloud_file_getters if g.multi is True]

    result = {
        **{f"/tmp/osworld/cache/{g.dest}": g.path for g in single_getters},
        **{
            f"/tmp/osworld/cache/{dest}": path
            for g in multi_getters
            for path, dest in zip(g.path, g.dest)
        },
    }
    return result


OSWORLD_REPO_URL = "https://github.com/xlang-ai/OSWorld.git"
EVALUATION_EXAMPLES_ROOT = "evaluation_examples"
CONTAINER_EXAMPLE_PATH = "/tmp/osworld/desktop_env/example.json"
CONFIG_COMMAND = "python3 -m  desktop_env.cli config"

unsupported_apps = ["chrome", "thunderbird"]
"""
List of unsupported applications that should be excluded from the dataset.

- chrome — the docker container does not yet have chrome installed.
- thunderbird — the credentials included in these OSWorld examples are invalid. This
isn't fatal, since many of the examples work on local data. However, it does cause
modal error windows to be opened distracting the model.
"""

unsupported_examples = [
    # video will not load properly when --start-time option is passed. Need to debug.
    "efcf0d81-0835-4880-b2fd-d866e8bc2294",
    # the following examples use the accessibility_tree getter in their evaluator.
    "ef9d12bd-bcee-4ba0-a40e-918400f43ddf",
    "f5c13cdd-205c-4719-a562-348ae5cd1d91",
    "12086550-11c0-466b-b367-1d9e75b3910e",
    # relies on systemd which is not in place in Docker
    "b6781586-6346-41cd-935a-a6b1487918fc",  # "timedatectl status" - could be replaced "date"
    "37b9808f-b2b4-4177-ab00-9ddfae4bad27",  # ["timedatectl", "set-ntp", "off"] and ["timedatectl", "set-time", "16:00"]
    # assumes the use of org.gnome.shell
    "ec4e3f68-9ea4-4c18-a5c9-69f89d1178b3",
    "f9be0997-4b7c-45c5-b05c-4612b44a6118",
    "2b9493d7-49b8-493a-a71b-56cd1f4d6908",
    # requires get_terminal_output which uses accessibility tree
    "0b35aee5-6f2a-4c71-b2da-aed0105fdbde",
    "4127319a-8b79-4410-b58a-7a151e15f3d7",
    "13584542-872b-42d8-b299-866967b5c3ef",
    # these multi_apps examples rely on Chrome dev tools and socat port forwarding. NYI
    "0c825995-5b70-4526-b663-113f4c999dd2",
    "0e5303d4-8820-42f6-b18d-daf7e633de21",
    "3e3fc409-bff3-4905-bf16-c968eee3f807",
    "3eb2a122-a5e3-4f89-9820-f7fa1a582969",
    "3f05f3b9-29ba-4b6b-95aa-2204697ffc06",
    "4e9f0faf-2ecc-4ae8-a804-28c9a75d1ddc",
    "22a4636f-8179-4357-8e87-d1743ece1f81",
    "42d25c08-fb87-4927-8b65-93631280a26f",
    "48c46dc7-fe04-4505-ade7-723cba1aa6f6",
    "48d05431-6cd5-4e76-82eb-12b60d823f7d",
    "68a25bd4-59c7-4f4d-975e-da0c8509c848",
    "69acbb55-d945-4927-a87b-8480e1a5bb7e",
    "74d5859f-ed66-4d3e-aa0e-93d7a592ce41",
    "78aed49a-a710-4321-a793-b611a7c5b56b",
    "873cafdd-a581-47f6-8b33-b9696ddb7b05",
    "897e3b53-5d4d-444b-85cb-2cdc8a97d903",
    "26660ad1-6ebb-4f59-8cba-a8432dfe8d38",
    "67890eb6-6ce5-4c00-9e3d-fb4972699b06",
    "236833a3-5704-47fc-888c-4f298f09f799",
    "5990457f-2adb-467b-a4af-5c857c92d762",
    "46407397-a7d5-4c6b-92c6-dbe038b1457b",
    "58565672-7bfe-48ab-b828-db349231de6b",
    "a0b9dc9c-fc07-4a88-8c5d-5e3ecad91bcb",
    "a74b607e-6bb5-4ea8-8a7c-5d97c7bbcd2a",
    "a82b78bb-7fde-4cb3-94a4-035baf10bcf0",
    "aad10cd7-9337-4b62-b704-a857848cedf2",
    "acb0f96b-e27c-44d8-b55f-7cb76609dfcd",
    "b52b40a5-ad70-4c53-b5b0-5650a8387052",
    "b337d106-053f-4d37-8da0-7f9c4043a66b",
    "c7c1e4c3-9e92-4eba-a4b8-689953975ea4",
    "da922383-bfa4-4cd3-bbad-6bebab3d7742",
    "dd60633f-2c72-42ba-8547-6f2c8cb0fdb0",
    "df67aebb-fb3a-44fd-b75b-51b6012df509",
    "e135df7c-7687-4ac0-a5f0-76b74438b53e",
    "e2392362-125e-4f76-a2ee-524b183a3412",
    "f8cfa149-d1c1-4215-8dac-4a0932bad3c2",
    # to be assessed
    "937087b6-f668-4ba6-9110-60682ee33441",  # vm_ip
    "9219480b-3aed-47fc-8bac-d2cffc5849f7",
    "869de13e-bef9-4b91-ba51-f6708c40b096",
    "9219480b-3aed-47fc-8bac-d2cffc5849f7",
    "26150609-0da3-4a7d-8868-0faf9c5f01bb",
    "f8369178-fafe-40c2-adc4-b9b08a125456",
    "e1fc0df3-c8b9-4ee7-864c-d0b590d3aa56",
    "c2751594-0cd5-4088-be1b-b5f2f9ec97c4",
]
"""
List of unsupported example IDs that should be excluded from the dataset because
they are not yet supported for a variety of reasons.
"""

examples_requiring_network = [
    # uses vlc_playing_info which gets the current playing information from VLC's HTTP interface.
    "59f21cfb-0120-4326-b255-a5b827b38967",
    # instructions explicitly require internet
    "bba3381f-b5eb-4439-bd9e-80c22218d5a7",
    "94d95f96-9699-4208-98ba-3c3119edf9c2",
    "da52d699-e8d2-4dc5-9191-a2199e0b6a9b",
    "3c8f201a-009d-4bbe-8b65-a6f8b35bb57f",
    "0e5303d4-8820-42f6-b18d-daf7e633de21",
    "69acbb55-d945-4927-a87b-8480e1a5bb7e",
    "788b3701-3ec9-4b67-b679-418bfa726c22",
    "236833a3-5704-47fc-888c-4f298f09f799",
    "36037439-2044-4b50-b9d1-875b5a332143",
    "58565672-7bfe-48ab-b828-db349231de6b",
    "aad10cd7-9337-4b62-b704-a857848cedf2",
    "acb0f96b-e27c-44d8-b55f-7cb76609dfcd",
    "da52d699-e8d2-4dc5-9191-a2199e0b6a9b",
    "deec51c9-3b1e-4b9e-993c-4776f20e8bb2",
    "e1fc0df3-c8b9-4ee7-864c-d0b590d3aa56",
    "f8369178-fafe-40c2-adc4-b9b08a125456",
    "42f4d1c7-4521-4161-b646-0a8934e36081",
]
"""
List of example IDs that require network access and should be excluded from the dataset if network
access is not allowed. Note that this list is likely incomplete, and additions are welcome.
"""
