from dataclasses import dataclass
from pathlib import Path

from harbor.models.job.config import RegistryDatasetConfig
from harbor.models.task.id import GitTaskId, LocalTaskId
from harbor.tasks.client import TasksClient


@dataclass
class DownloadedTask:
    task_id: LocalTaskId | GitTaskId
    local_path: Path


class DatasetClient:
    def download_dataset_from_config(
        self, config: RegistryDatasetConfig
    ) -> list[DownloadedTask]:
        client = TasksClient()

        task_ids = [task_id.get_task_id() for task_id in config.get_task_configs()]

        paths = client.download_tasks(
            task_ids=task_ids,
            overwrite=config.overwrite,
            output_dir=config.download_dir,
        )

        return [
            DownloadedTask(task_id=task_id, local_path=path)
            for task_id, path in zip(task_ids, paths)
        ]
