from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Iterator


@dataclass(frozen=True)
class DatasetInfoFile:
    repo_id: str
    path: Path


@dataclass(frozen=True)
class DatasetConfigRef:
    dataset_id: str
    config_name: str | None


def _repo_from_dirname(dirname: str) -> str:
    if not dirname.startswith("datasets--"):
        raise ValueError(f"Unsupported directory name: {dirname}")
    remainder = dirname[len("datasets--") :]
    if "--" not in remainder:
        return remainder
    owner, dataset = remainder.split("--", 1)
    return f"{owner}/{dataset}"


def iter_dataset_info_files(cache_dir: Path) -> Iterator[DatasetInfoFile]:
    """Yield dataset_info.json files stored under the HF hub cache."""
    for repo_dir in cache_dir.glob("datasets--*"):
        repo_id = _repo_from_dirname(repo_dir.name)
        for info_path in repo_dir.rglob("dataset_info*.json"):
            if info_path.is_dir():
                continue
            yield DatasetInfoFile(repo_id=repo_id, path=info_path)


def iter_repo_ids(cache_dir: Path) -> Iterator[str]:
    """Yield repo ids present in the hub cache."""
    for repo_dir in cache_dir.glob("datasets--*"):
        yield _repo_from_dirname(repo_dir.name)


def _dataset_id_from_cache_dir(name: str) -> str:
    if "___" in name:
        owner, dataset = name.split("___", 1)
        return f"{owner}/{dataset}"
    return name


def iter_dataset_cache_configs(dataset_cache_dir: Path) -> Iterator[DatasetConfigRef]:
    """Infer downloaded dataset/config combos from the datasets cache."""
    if not dataset_cache_dir.exists():
        return

    for entry in dataset_cache_dir.iterdir():
        if not entry.is_dir():
            continue
        if entry.name.startswith("_") or entry.name == "downloads":
            continue
        dataset_id = _dataset_id_from_cache_dir(entry.name)
        subdirs = [p for p in entry.iterdir() if p.is_dir()]
        if not subdirs:
            yield DatasetConfigRef(dataset_id=dataset_id, config_name=None)
            continue
        for sub in subdirs:
            config_name = sub.name
            yield DatasetConfigRef(
                dataset_id=dataset_id,
                config_name=None if config_name == "default" else config_name,
            )
