src/dataset_samples.py

from __future__ import annotations

import os
import zipfile
from dataclasses import asdict, dataclass
from pathlib import Path, PurePosixPath
from typing import Callable, Iterable

from huggingface_hub import HfApi, hf_hub_download

from .config import (
    DATASET_SAMPLE_ROOT,
    MAX_DATASET_SAMPLE_ARCHIVE_BYTES,
    MAX_DATASET_SAMPLE_FILE_BYTES,
)

_AUDIO_EXTENSIONS = {".wav", ".flac", ".mp3", ".ogg", ".opus", ".m4a", ".aiff", ".aif"}
_ARCHIVE_EXTENSIONS = {".zip"}


@dataclass(frozen=True)
class DatasetSampleSpec:
    sample_id: str
    label: str
    repo_id: str
    repo_type: str
    revision: str
    license_id: str
    attribution: str
    source_url: str
    filename: str | None = None
    selector: str | None = None
    first_download_note: str = ""


@dataclass(frozen=True)
class ResolvedDatasetSample:
    sample_id: str
    label: str
    local_path: str
    repo_id: str
    repo_type: str
    revision: str
    repository_filename: str
    resolved_audio_filename: str
    license_id: str
    attribution: str
    source_url: str
    archive_extracted: bool
    size_bytes: int

    def to_json(self) -> str:
        return json.dumps(asdict(self), ensure_ascii=False, sort_keys=True)


DATASET_SAMPLE_SPECS: tuple[DatasetSampleSpec, ...] = (
    DatasetSampleSpec(
        sample_id="sonicsets-demo-mix",
        label="Dataset: SonicSets ensemble mix (CC BY-NC 4.0; first download may be large)",
        repo_id="sonicsets-data/Stems-Evaluation-Kit",
        repo_type="dataset",
        revision="main",
        license_id="CC-BY-NC-4.0",
        attribution="SonicSets High-Fidelity Stems Evaluation Kit",
        source_url="https://huggingface.co/datasets/sonicsets-data/Stems-Evaluation-Kit",
        selector="ensemble-mix",
        first_download_note=(
            "The repository is about 333 MB in total. SESA downloads one matching audio file when possible; "
            "if the repository exposes only an archive, the first download can approach the archive size."
        ),
    ),
    DatasetSampleSpec(
        sample_id="legacy-symphony-01",
        label="Hosted fallback: Symphony of Automation 01 (small)",
        repo_id="set-soft/audio_separation",
        repo_type="model",
        revision="main",
        filename="Audio_Examples/01-The_Symphony_of_Automation.mp3",
        license_id="MIT repository metadata",
        attribution="set-soft/audio_separation hosted audio example",
        source_url="https://huggingface.co/set-soft/audio_separation/tree/main/Audio_Examples",
        first_download_note="Small hosted fallback. The repository does not separately document audio-asset provenance.",
    ),
    DatasetSampleSpec(
        sample_id="legacy-symphony-04",
        label="Hosted fallback: Symphony of Automation 04 (small)",
        repo_id="set-soft/audio_separation",
        repo_type="model",
        revision="main",
        filename="Audio_Examples/04-The_Symphony_of_Automation.mp3",
        license_id="MIT repository metadata",
        attribution="set-soft/audio_separation hosted audio example",
        source_url="https://huggingface.co/set-soft/audio_separation/tree/main/Audio_Examples",
        first_download_note="Small hosted fallback. The repository does not separately document audio-asset provenance.",
    ),
    DatasetSampleSpec(
        sample_id="legacy-symphony-05",
        label="Hosted fallback: Symphony of Automation 05 (small)",
        repo_id="set-soft/audio_separation",
        repo_type="model",
        revision="main",
        filename="Audio_Examples/05-The_Symphony_of_Automation.mp3",
        license_id="MIT repository metadata",
        attribution="set-soft/audio_separation hosted audio example",
        source_url="https://huggingface.co/set-soft/audio_separation/tree/main/Audio_Examples",
        first_download_note="Small hosted fallback. The repository does not separately document audio-asset provenance.",
    ),
)

_SPEC_BY_ID = {item.sample_id: item for item in DATASET_SAMPLE_SPECS}
DEFAULT_DATASET_SAMPLE_ID = DATASET_SAMPLE_SPECS[0].sample_id


def dataset_sample_choices() -> list[tuple[str, str]]:
    return [(item.label, item.sample_id) for item in DATASET_SAMPLE_SPECS]


def dataset_sample_markdown(sample_id: str | None) -> str:
    spec = _SPEC_BY_ID.get(str(sample_id or ""))
    if not spec:
        return "Select a curated hosted sample, or upload a file to override it."
    note = f"  \n{spec.first_download_note}" if spec.first_download_note else ""
    return (
        f"**Source:** `{spec.repo_id}` · **license:** `{spec.license_id}`  \n"
        f"Attribution: {spec.attribution}.{note}"
    )


def _repo_token() -> str | bool:
    return os.environ.get("HF_TOKEN") or False


def _score_audio_candidate(filename: str, selector: str | None) -> int:
    path = PurePosixPath(filename)
    if path.suffix.lower() not in _AUDIO_EXTENSIONS:
        return -10_000
    text = str(path).lower().replace("-", "_").replace(" ", "_")
    score = 0
    if selector == "ensemble-mix":
        preferred = (
            ("full_ensemble_mix", 500),
            ("ensemble_mix", 450),
            ("full_mix", 400),
            ("mixture", 350),
            ("mix", 250),
            ("ensemble", 180),
        )
        for token, points in preferred:
            if token in text:
                score += points
        for token in ("isolated", "stem", "vocal", "vocals", "bass", "drum", "guitar", "harmony"):
            if token in text:
                score -= 220
        if "demo_01" in text or "demo01" in text or "/01" in text:
            score += 35
    score -= len(path.parts) * 2
    return score


def _choose_audio_filename(filenames: Iterable[str], selector: str | None) -> str | None:
    ranked = sorted(
        ((-_score_audio_candidate(name, selector), len(name), name) for name in filenames),
    )
    if not ranked:
        return None
    best_score = -ranked[0][0]
    if best_score <= 0:
        return None
    return ranked[0][2]


def _choose_archive_filename(filenames: Iterable[str]) -> str | None:
    archives = [name for name in filenames if PurePosixPath(name).suffix.lower() in _ARCHIVE_EXTENSIONS]
    if not archives:
        return None
    return sorted(archives, key=lambda value: (len(PurePosixPath(value).parts), len(value), value))[0]


def _safe_extract_zip(archive_path: Path, destination: Path) -> None:
    destination.mkdir(parents=True, exist_ok=True)
    root = destination.resolve()
    total_size = 0
    with zipfile.ZipFile(archive_path) as archive:
        for member in archive.infolist():
            if member.is_dir():
                continue
            total_size += max(0, int(member.file_size))
            if total_size > MAX_DATASET_SAMPLE_ARCHIVE_BYTES:
                raise RuntimeError("Dataset sample archive expands beyond the configured safety limit.")
            target = (destination / member.filename).resolve()
            if root not in target.parents:
                raise RuntimeError("Unsafe path found in dataset sample archive.")
        archive.extractall(destination)


def _resolve_from_archive(archive_path: Path, spec: DatasetSampleSpec) -> Path:
    extracted = archive_path.parent / "extracted"
    marker = extracted / ".complete"
    if not marker.is_file():
        _safe_extract_zip(archive_path, extracted)
        marker.write_text("ok\n", encoding="utf-8")
    candidates = [
        str(path.relative_to(extracted).as_posix())
        for path in extracted.rglob("*")
        if path.is_file() and path.suffix.lower() in _AUDIO_EXTENSIONS
    ]
    selected = _choose_audio_filename(candidates, spec.selector)
    if not selected:
        raise RuntimeError("No suitable mixture audio was found inside the dataset archive.")
    return extracted / selected


def resolve_dataset_sample(
    sample_id: str,
    *,
    api: HfApi | None = None,
    download_fn: Callable[..., str] = hf_hub_download,
) -> ResolvedDatasetSample:
    spec = _SPEC_BY_ID.get(str(sample_id or ""))
    if not spec:
        raise ValueError("Unknown dataset sample selection.")

    local_dir = DATASET_SAMPLE_ROOT / spec.sample_id
    local_dir.mkdir(parents=True, exist_ok=True)
    repository_filename = spec.filename
    if repository_filename is None:
        client = api or HfApi(token=_repo_token())
        filenames = client.list_repo_files(
            repo_id=spec.repo_id,
            repo_type=spec.repo_type,
            revision=spec.revision,
        )
        repository_filename = _choose_audio_filename(filenames, spec.selector)
        if repository_filename is None:
            repository_filename = _choose_archive_filename(filenames)
        if repository_filename is None:
            raise RuntimeError("No supported audio file or ZIP archive was found in the dataset repository.")

    downloaded = Path(
        download_fn(
            repo_id=spec.repo_id,
            repo_type=spec.repo_type,
            filename=repository_filename,
            revision=spec.revision,
            local_dir=local_dir,
            token=_repo_token(),
            etag_timeout=15,
        )
    )
    if not downloaded.is_file() or downloaded.stat().st_size <= 0:
        raise RuntimeError("Dataset sample download did not produce a readable file.")
    if downloaded.stat().st_size > MAX_DATASET_SAMPLE_FILE_BYTES:
        raise RuntimeError("Dataset sample download exceeds the configured file-size limit.")

    archive_extracted = downloaded.suffix.lower() in _ARCHIVE_EXTENSIONS
    local_audio = _resolve_from_archive(downloaded, spec) if archive_extracted else downloaded
    if not local_audio.is_file() or local_audio.suffix.lower() not in _AUDIO_EXTENSIONS:
        raise RuntimeError("Resolved dataset sample is not a supported audio file.")

    return ResolvedDatasetSample(
        sample_id=spec.sample_id,
        label=spec.label,
        local_path=str(local_audio),
        repo_id=spec.repo_id,
        repo_type=spec.repo_type,
        revision=spec.revision,
        repository_filename=repository_filename,
        resolved_audio_filename=local_audio.name,
        license_id=spec.license_id,
        attribution=spec.attribution,
        source_url=spec.source_url,
        archive_extracted=archive_extracted,
        size_bytes=local_audio.stat().st_size,
    )