from collections import defaultdict
from pathlib import Path

from harbor.constants import DEFAULT_REGISTRY_URL
from harbor.models.registry import (
    DatasetMetadata,
    DatasetSpec,
    DatasetSummary,
    Registry,
)
from harbor.registry.client.base import BaseRegistryClient, resolve_version


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_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]

    def _spec_to_metadata(self, spec: DatasetSpec) -> DatasetMetadata:
        return DatasetMetadata(
            name=spec.name,
            version=spec.version,
            description=spec.description,
            task_ids=[task.to_source_task_id() for task in spec.tasks],
            metrics=spec.metrics,
        )

    async def _get_dataset_metadata(self, name: str) -> DatasetMetadata:
        if "@" in name:
            dataset_name, version = name.split("@", 1)
        else:
            dataset_name = name
            versions = self._get_dataset_versions(dataset_name)
            version = resolve_version(versions)

        spec = self._get_dataset_spec(dataset_name, version)
        return self._spec_to_metadata(spec)

    async def list_datasets(self) -> list[DatasetSummary]:
        return [
            DatasetSummary(
                name=spec.name,
                version=spec.version,
                description=spec.description,
                task_count=len(spec.tasks),
            )
            for spec in self._registry.datasets
        ]
