from __future__ import annotations

from pydantic import BaseModel, Field

from lite_dist2.curriculum_models.study_portables import StudyRegistry
from lite_dist2.curriculum_models.trial import TrialModel


class BaseParam(BaseModel):
    pass


class StudyRegisterParam(BaseParam):
    study: StudyRegistry = Field(description="`Study` to register")


class TrialReserveParam(BaseParam):
    retaining_capacity: set[str] = Field(description="List of capabilities that the worker node has.")
    max_size: int = Field(description="The maximum size of parameter space reserving.")
    worker_node_name: str | None = Field(default=None, description="Name of the worker node. ")
    worker_node_id: str = Field(description="ID of the worker node")


class TrialRegisterParam(BaseModel):
    trial: TrialModel = Field(description="Registering trial to the table node.")
