from collections import defaultdict

from agentdojo.default_suites.v1.banking import banking_task_suite
from agentdojo.default_suites.v1.slack import slack_task_suite
from agentdojo.default_suites.v1.travel import travel_task_suite
from agentdojo.default_suites.v1.workspace import workspace_task_suite
from agentdojo.default_suites.v1_1.travel import user_tasks  # noqa: F401 - Register new tasks
from agentdojo.task_suite.task_suite import TaskSuite

_V1_SUITES: dict[str, TaskSuite] = {
    "workspace": workspace_task_suite,
    "travel": travel_task_suite,
    "banking": banking_task_suite,
    "slack": slack_task_suite,
}

_V1_1_SUITES: dict[str, TaskSuite] = {
    "workspace": workspace_task_suite,
    "travel": travel_task_suite.get_new_version((1, 1, 0)),
    "banking": banking_task_suite,
    "slack": slack_task_suite,
}

_SUITES: defaultdict[str, dict[str, TaskSuite]] = defaultdict(dict, {"v1": _V1_SUITES, "v1.1": _V1_1_SUITES})


def register_suite(suite: TaskSuite, version: str) -> None:
    _SUITES[version][suite.name] = suite


def get_suites(benchmark_version: str) -> dict[str, TaskSuite]:
    return _SUITES[benchmark_version]


def get_suite(benchmark_version: str, suite_name: str) -> TaskSuite:
    return _SUITES[benchmark_version][suite_name]
