from collections import defaultdict
from pathlib import Path

from harbor.constants import DEFAULT_REGISTRY_URL
from harbor.models.registry import DatasetSpec, Registry
from harbor.registry.client.base import BaseRegistryClient


class JsonRegistryClient(BaseRegistryClient):
    def __init__(self, url: str | None = None, path: Path | None = None):
        super().__init__(url, path)

        self._init_registry()

    def _init_registry(self):
        if self._url is not None and self._path is not None:
            raise ValueError("Only one of url or path can be provided")

        if self._url is not None:
            self._registry = Registry.from_url(self._url)
        elif self._path is not None:
            self._registry = Registry.from_path(self._path)
        else:
            self._registry = Registry.from_url(DEFAULT_REGISTRY_URL)

    @property
    def dataset_specs(self) -> dict[str, dict[str, DatasetSpec]]:
        datasets = defaultdict(dict)
        for dataset in self._registry.datasets:
            datasets[dataset.name][dataset.version] = dataset
        return datasets

    def get_datasets(self) -> list[DatasetSpec]:
        """Get all datasets available in the registry."""
        return self._registry.datasets

    def get_dataset_versions(self, name: str) -> list[str]:
        if name not in self.dataset_specs:
            raise ValueError(f"Dataset {name} not found")
        return list(self.dataset_specs[name].keys())

    def _get_dataset_spec(self, name: str, version: str) -> DatasetSpec:
        if name not in self.dataset_specs:
            raise ValueError(f"Dataset {name} not found")

        if version not in self.dataset_specs[name]:
            raise ValueError(f"Version {version} of dataset {name} not found")

        return self.dataset_specs[name][version]
