from .. import ty
from ..models import (
    GroupDetail,
    TaskAgentPopulating,
    TaskCreating,
    TaskDeleting,
    TaskDetail,
    TaskReading,
    TaskSummary,
    TaskUpdating,
)
from .base import ModelClient


class Task(
    ModelClient[
        TaskCreating,
        TaskReading,
        TaskUpdating,
        TaskDeleting,
        TaskSummary,
        TaskDetail,
    ]
):
    Creating = TaskCreating
    Reading = TaskReading
    Updating = TaskUpdating
    Deleting = TaskDeleting
    Summary = TaskSummary
    Detail = TaskDetail

    def add_agent(self, task_id: ty.UUID, agent_id: ty.UUID) -> GroupDetail:
        data = TaskAgentPopulating(id=task_id, agent_id=agent_id)
        obj = self.request(
            "POST", endpoint="add-agent", data=data, return_model=GroupDetail
        )
        assert isinstance(obj, GroupDetail)
        return obj
