"""Parallel S10 training tasks: architecture, environment and recipe.

A task owns the observation contract (the part that cannot be changed by a
resume) and the training defaults that go with its terrain revision. Shared
learner, reward and world-model code lives outside this package, so adding a
new task means adding one module and registering it in ``tasks/__init__.py``.
"""

from dataclasses import dataclass, field
from typing import Any, Callable


@dataclass(frozen=True)
class ActorSpec:
  """Actor observation contract used by training, export and playback."""
  name: str
  terms: tuple[int, ...]
  frame_terms: tuple[dict[str, Any], ...]
  wheel_vel: bool

  @property
  def frame_dim(self) -> int:
    return sum(self.terms)

  def history_dim(self, history: int) -> int:
    if history < 1:
      raise ValueError("history must be positive")
    return self.frame_dim * history

  def metadata(self, history: int) -> dict[str, Any]:
    return {
      "actor_schema": self.name,
      "actor_wheel_vel": self.wheel_vel,
      "frame_dim": self.frame_dim,
      "observation_dim": self.history_dim(history),
      "frame_terms": [dict(term) for term in self.frame_terms],
    }


@dataclass(frozen=True)
class TaskSpec:
  """One selectable training task.

  ``defaults`` are argparse destinations applied before parsing, so explicit
  command-line flags still override them. That keeps a task reproducible while
  allowing a one-off run to change length or batch size.
  """
  name: str
  description: str
  terrain_revision: int
  actor: ActorSpec
  defaults: dict[str, Any] = field(default_factory=dict)
  env_cfg: Callable[..., Any] | None = None

  @property
  def contacts(self) -> bool:
    return self.terrain_revision >= 4

  @property
  def courses(self) -> bool:
    return self.terrain_revision >= 5

  @property
  def collisions(self) -> bool:
    return self.terrain_revision >= 7

  def build_env_cfg(self, play: bool, history: int, seed: int):
    if self.env_cfg is None:
      raise ValueError(f"Task {self.name} has no environment builder")
    return self.env_cfg(play=play, history=history, seed=seed)


_REGISTRY: dict[str, TaskSpec] = {}


def register(task: TaskSpec) -> TaskSpec:
  if task.name in _REGISTRY:
    raise ValueError(f"Duplicate task name: {task.name}")
  _REGISTRY[task.name] = task
  return task


def get_task(name: str) -> TaskSpec:
  try:
    return _REGISTRY[name]
  except KeyError as exc:
    raise KeyError(f"Unknown task {name!r}; available: {', '.join(task_names())}") from exc


def list_tasks() -> tuple[TaskSpec, ...]:
  return tuple(_REGISTRY.values())


def task_names() -> tuple[str, ...]:
  return tuple(_REGISTRY)