from pathlib import Path

from harbor.models.package.reference import PackageReference
from harbor.models.registry import DatasetFileInfo, DatasetMetadata, DatasetSummary
from harbor.models.task.id import GitTaskId, LocalTaskId, PackageTaskId
from harbor.registry.client.base import BaseRegistryClient
from harbor.registry.db import RegistryDB


class PackageDatasetClient(BaseRegistryClient):
    def __init__(self):
        super().__init__()
        self._db = RegistryDB()

    async def _get_dataset_metadata(self, name: str) -> DatasetMetadata:
        ref = PackageReference.parse(name)
        _package, dataset_version = await self._db.resolve_dataset_version(
            ref.org, ref.short_name, ref.ref
        )

        # Get tasks in this dataset version
        tasks_data = await self._db.get_dataset_version_tasks(dataset_version["id"])

        task_ids: list[GitTaskId | LocalTaskId | PackageTaskId] = []
        for row in tasks_data:
            tv = row["task_version"]
            pkg = tv["package"]
            task_ids.append(
                PackageTaskId(
                    org=pkg["org"]["name"],
                    name=pkg["name"],
                    digest=tv["content_hash"],
                )
            )

        # Get dataset-level files (e.g., metric.py)
        files_data = await self._db.get_dataset_version_files(dataset_version["id"])
        files = [
            DatasetFileInfo(
                path=row["path"],
                storage_path=row["storage_path"],
                content_hash=row["content_hash"],
            )
            for row in files_data
        ]

        return DatasetMetadata(
            name=ref.name,
            version=ref.ref,
            description=dataset_version.get("description") or "",
            task_ids=task_ids,
            metrics=[],
            files=files,
            dataset_version_content_hash=dataset_version.get("content_hash"),
        )

    async def download_dataset_files(
        self,
        metadata: DatasetMetadata,
        overwrite: bool = False,
        output_dir: Path | None = None,
    ) -> dict[str, Path]:
        """Download dataset-level files. Returns {filename: local_path}."""
        from harbor.constants import DATASET_CACHE_DIR
        from harbor.storage.supabase import SupabaseStorage

        if not metadata.files or not metadata.dataset_version_content_hash:
            return {}

        if output_dir is not None:
            cache_dir = output_dir
        else:
            org, name = metadata.name.split("/", 1)
            version = metadata.dataset_version_content_hash
            cache_dir = DATASET_CACHE_DIR / org / name / version

        storage = SupabaseStorage()
        result: dict[str, Path] = {}

        for file_info in metadata.files:
            local_path = cache_dir / file_info.path
            if not local_path.exists() or overwrite:
                await storage.download_file(file_info.storage_path, local_path)
            result[file_info.path] = local_path

        return result

    async def list_datasets(self) -> list[DatasetSummary]:
        raise NotImplementedError("Listing all package datasets is not yet supported")
