tasks/registry.py

from __future__ import annotations

from typing import Any, Callable, Dict, List

from env.models import EnvironmentConfig
from tasks import easy, hard, medium


TaskFactory = Callable[[int | None], EnvironmentConfig]


TASK_REGISTRY: Dict[str, Dict[str, Any]] = {
	easy.TASK_NAME: {
		"definition": easy.get_definition(),
		"factory": easy.get_config,
	},
	medium.TASK_NAME: {
		"definition": medium.get_definition(),
		"factory": medium.get_config,
	},
	hard.TASK_NAME: {
		"definition": hard.get_definition(),
		"factory": hard.get_config,
	},
}

TASK_DIFFICULTY_ORDER: List[str] = [easy.TASK_NAME, medium.TASK_NAME, hard.TASK_NAME]


def _ordered_task_names() -> List[str]:
	return [name for name in TASK_DIFFICULTY_ORDER if name in TASK_REGISTRY]


def list_tasks() -> List[str]:
	return _ordered_task_names()


def get_task_registry() -> Dict[str, Dict[str, Any]]:
	return {
		name: dict(item["definition"])
		for name in _ordered_task_names()
		for item in [TASK_REGISTRY[name]]
	}


def get_task_definition(task_name: str) -> Dict[str, Any]:
	item = TASK_REGISTRY.get(task_name)
	if item is None:
		raise ValueError(f"Unknown task '{task_name}'. Available: {list_tasks()}")
	return dict(item["definition"])


def get_task_config(task_name: str, seed: int | None = None) -> EnvironmentConfig:
	item = TASK_REGISTRY.get(task_name)
	if item is None:
		raise ValueError(f"Unknown task '{task_name}'. Available: {list_tasks()}")

	factory = item["factory"]
	if not callable(factory):
		raise RuntimeError(f"Task '{task_name}' has invalid factory")
	return factory(seed)


def get_all_task_metadata() -> List[Dict[str, Any]]:
	return [
		dict(TASK_REGISTRY[name]["definition"])
		for name in _ordered_task_names()
	]


def list_tasks_with_metadata() -> List[Dict[str, Any]]:
	"""Compatibility alias for tooling that expects this function name."""
	return get_all_task_metadata()