# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""HuggingFace Hub download worker, spawned as a subprocess so SIGKILL stops all chunk threads.

Resume safety
-------------
Downloads here MUST be single-stream sequential writers so the parent's
SIGKILL → restart loop can rely on ``os.path.getsize(.incomplete)`` to
compute the correct resume offset.

Enforced by:
- Setting ``HF_HUB_DISABLE_XET=1`` and ``HF_HUB_ENABLE_HF_TRANSFER=0`` on
  the spawning side (see :mod:`hub.utils.download_registry`) for transport=http.
- Passing ``max_workers=1`` to ``snapshot_download`` so files download
  serially, making the at-most-one-active-`.incomplete` invariant hold
  globally and simplifying reasoning about partial state during a SIGKILL.
- Letting ``prepare_cache_for_transport`` purge any pre-existing
  ``.incomplete`` blobs not provably from the same sequential writer.

If the final byte count doesn't match what HF declared, huggingface_hub
raises ``EnvironmentError`` ("Consistency check failed: …"); we surface
that on stderr so the watcher can show the exact message to the user.
"""

from __future__ import annotations

import argparse
import os
import signal
import sys
import threading
import time
from pathlib import Path

_HERE = Path(__file__).resolve().parent
_BACKEND = _HERE.parent.parent
if str(_BACKEND) not in sys.path:
    sys.path.insert(0, str(_BACKEND))

from hub.utils.snapshot_filters import (
    SNAPSHOT_IGNORE_PATTERNS,
)
from hub.utils.gguf_plan import (
    GgufVariantPlan,
    build_gguf_variant_plans,
    plan_from_expected_files,
    sibling_sha256,
)
from hub.utils.state_dir import RepoType

HfTokenArg = str | bool | None


# Bound the metadata fetch so a stalled connection fails the worker (exit 1)
# instead of hanging at 0%. The file download itself is governed separately by
# huggingface_hub's own timeout.
_METADATA_REQUEST_TIMEOUT = 10.0
_METADATA_RETRY_TIMEOUT = 30.0
_METADATA_RETRY_DELAY = 1.0


def _on_signal(signum, frame):
    # 130 is what `classify_exit` maps to the "cancelled" job state.
    sys.exit(130)


def _install_signal_handlers() -> None:
    signal.signal(signal.SIGTERM, _on_signal)
    signal.signal(signal.SIGINT, _on_signal)
    sigpipe = getattr(signal, "SIGPIPE", None)
    if sigpipe is not None:
        signal.signal(sigpipe, _on_signal)


def _parent_poll_seconds() -> float:
    raw = os.environ.get("UNSLOTH_HF_WORKER_PARENT_POLL_SECONDS")
    if raw:
        try:
            value = float(raw)
            if value > 0:
                return value
        except ValueError:
            pass
    return 2.0


def _protected_blob_hashes() -> frozenset[str]:
    """Blob hashes a concurrent same-repo peer is writing (passed by the backend
    as a plain env list). Excluded from this worker's purge so a shared
    ``.incomplete`` (e.g. a bundled mmproj) is never deleted under the peer."""
    raw = os.environ.get("UNSLOTH_PROTECTED_BLOB_HASHES", "")
    return frozenset(h for h in raw.split(",") if h)


def _parent_is_alive(parent_pid: int) -> bool:
    """Whether the recorded parent (the backend) is still running.

    Liveness ONLY: ``os.kill(pid, 0)`` on POSIX, an ``OpenProcess`` handle on
    Windows, against the *recorded* PID (never os.getppid(), so POSIX
    reparenting to init after the backend dies still resolves as dead). Probe
    ambiguity is treated as alive so a transient error never kills a healthy
    download.

    We deliberately do NOT compare psutil ``create_time()`` for PID-reuse
    detection: it isn't stable across reads on some platforms, so an exact match
    can spuriously kill a live download. PID-reuse after parent death is covered
    by the boot-time orphan reaper.
    """
    if sys.platform == "win32":
        import ctypes
        from ctypes import wintypes

        SYNCHRONIZE = 0x00100000
        WAIT_OBJECT_0 = 0x0
        ERROR_INVALID_PARAMETER = 87
        kernel32 = ctypes.WinDLL("kernel32", use_last_error = True)
        kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
        kernel32.OpenProcess.restype = wintypes.HANDLE
        kernel32.WaitForSingleObject.argtypes = [wintypes.HANDLE, wintypes.DWORD]
        kernel32.WaitForSingleObject.restype = wintypes.DWORD
        kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
        kernel32.CloseHandle.restype = wintypes.BOOL
        ctypes.set_last_error(0)
        handle = kernel32.OpenProcess(SYNCHRONIZE, False, parent_pid)
        if not handle:
            return ctypes.get_last_error() != ERROR_INVALID_PARAMETER
        try:
            return kernel32.WaitForSingleObject(handle, 0) != WAIT_OBJECT_0
        finally:
            kernel32.CloseHandle(handle)
    try:
        os.kill(parent_pid, 0)
    except ProcessLookupError:
        return False
    except OSError:
        return True
    return True


def _terminate_orphaned_self() -> None:
    # Hard exit from the watchdog thread: a self-SIGTERM would be deferred while
    # the main thread is GIL-blocked in a C socket read. The partial .incomplete
    # resumes byte-exact and marker/manifest writes are atomic, so cancelled code
    # 130 is safe. The diagnostic is best-effort: a dead parent's closed stderr
    # pipe can raise BrokenPipeError, which must never preempt the exit.
    try:
        print(
            "Parent process exited; stopping orphaned download worker.",
            file = sys.stderr,
        )
        sys.stderr.flush()
    except Exception:
        pass
    os._exit(130)


def _install_parent_death_watchdog(parent_pid: int | None) -> None:
    if not parent_pid or parent_pid <= 0:
        return
    interval = _parent_poll_seconds()

    def _watch() -> None:
        while True:
            try:
                alive = _parent_is_alive(parent_pid)
            except Exception:
                alive = True
            if not alive:
                _terminate_orphaned_self()
                return
            time.sleep(interval)

    threading.Thread(
        target = _watch,
        name = "parent-death-watchdog",
        daemon = True,
    ).start()


def _hf_token_arg(hf_token: str | None) -> HfTokenArg:
    return hf_token if hf_token else False


def _retry_metadata_fetch(repo_id: str, fetch, *, label: str):
    for attempt, timeout in enumerate((_METADATA_REQUEST_TIMEOUT, _METADATA_RETRY_TIMEOUT)):
        try:
            return fetch(timeout)
        except Exception as e:
            if attempt == 1:
                raise
            print(
                f"{label} request failed for {repo_id} " f"({type(e).__name__}: {e}); retrying.",
                file = sys.stderr,
            )
            time.sleep(_METADATA_RETRY_DELAY)
    raise RuntimeError(f"{label} unavailable for {repo_id}")


def _model_info_with_retry(repo_id: str, hf_token: str | None):
    from huggingface_hub import model_info as hf_model_info
    return _retry_metadata_fetch(
        repo_id,
        lambda timeout: hf_model_info(
            repo_id,
            token = _hf_token_arg(hf_token),
            timeout = timeout,
            files_metadata = True,
        ),
        label = "Metadata",
    )


def _dataset_info_with_retry(repo_id: str, hf_token: str | None):
    from huggingface_hub import HfApi
    api = HfApi(token = _hf_token_arg(hf_token))
    return _retry_metadata_fetch(
        repo_id,
        lambda timeout: api.dataset_info(
            repo_id,
            timeout = timeout,
            files_metadata = True,
        ),
        label = "Dataset metadata",
    )


# Tied to drain_stderr_excerpt's 500-byte head/tail window in the parent (see
# hub/utils/download_registry.py): listing every expected file would blow past
# it and lose the diagnostic. Cap the preview so the summary line survives.
_VERIFY_PATH_LIST_CAP = 10


def _format_path_list(paths: tuple[str, ...], cap: int = _VERIFY_PATH_LIST_CAP) -> str:
    if len(paths) <= cap:
        return ", ".join(paths)
    head = ", ".join(paths[:cap])
    return f"{head}, ... and {len(paths) - cap} more"


def _verify_completed_download(
    repo_type: RepoType,
    repo_id: str,
    variant: str | None,
    snapshot_path: str,
    *,
    metadata_unavailable: bool = False,
) -> None:
    """Verify every manifest file is on disk at its declared size; exit nonzero
    with a diagnostic if not.

    No-op when no manifest exists: the manifest write is best-effort, so absence
    means "verification unavailable, trust snapshot_download's exit code".
    """
    from hub.utils import download_manifest

    manifest = download_manifest.read_manifest(repo_type, repo_id, variant)
    if manifest is None:
        return
    result = download_manifest.verify_against_disk(
        manifest,
        Path(snapshot_path),
    )
    if result.ok:
        return
    label = f"{repo_id}{f' [{variant}]' if variant else ''}"
    if metadata_unavailable:
        print(
            f"Could not reach Hugging Face for {label} and the copy on disk is "
            f"incomplete ({len(result.missing)} file(s) missing, "
            f"{len(result.size_mismatched)} the wrong size). Access to a private "
            "or restricted repo may have been lost (HF token removed or "
            "changed), the connection dropped, or Hugging Face is temporarily "
            "unavailable. Set a valid HF token or reconnect, then resume the "
            "download.",
            file = sys.stderr,
        )
    else:
        print(
            f"Verification failed for {label}: snapshot_download completed but "
            f"{len(result.missing)} expected file(s) are missing and "
            f"{len(result.size_mismatched)} have incorrect size on disk.",
            file = sys.stderr,
        )
    if result.missing:
        print(
            f"Missing: {_format_path_list(result.missing)}",
            file = sys.stderr,
        )
    if result.size_mismatched:
        print(
            f"Size mismatched: {_format_path_list(result.size_mismatched)}",
            file = sys.stderr,
        )
    sys.exit(1)


def _preflight_disk_space(repo_type: str, repo_id: str, expected_files: list) -> None:
    """Fail fast when the active HF cache filesystem can't hold what's left to
    download. Fail-open: any inability to size the work or read free space skips
    the check, so a real download is never blocked by an estimation gap."""
    import shutil

    from hub.utils.download_registry import existing_blob_bytes
    from hub.utils.hf_cache_state import hf_cache_root

    try:
        size_by_hash: dict[str, int] = {}
        unhashed_bytes = 0
        for expected in expected_files:
            size = int(getattr(expected, "size", 0) or 0)
            if size <= 0:
                continue
            blob_hash = getattr(expected, "sha256", None)
            if blob_hash:
                # Dedup by content hash: a blob listed under two filenames is
                # written once, so it must be counted once.
                size_by_hash[blob_hash] = size
            else:
                unhashed_bytes += size
        total_expected = sum(size_by_hash.values()) + unhashed_bytes
        if total_expected <= 0:
            return
        already_have = existing_blob_bytes(
            repo_type,
            repo_id,
            frozenset(size_by_hash),
        )
        remaining = max(0, total_expected - already_have)
        if remaining <= 0:
            return
        root = hf_cache_root(create = True)
        if root is None:
            return
        free = shutil.disk_usage(root).free
    except Exception:
        return

    if free < remaining:
        print(
            f"Not enough disk space to download {repo_id}: need about "
            f"{remaining / (1024 ** 3):.1f} GB free in {root}, but only "
            f"{free / (1024 ** 3):.1f} GB is available. Free up space and "
            "try again.",
            file = sys.stderr,
        )
        sys.exit(1)


def _snapshot_download_plan(info) -> tuple[list[str], list]:
    from hub.utils.download_manifest import ExpectedFile
    from hub.utils.snapshot_filters import (
        resolve_snapshot_ignore_patterns_for_files,
        snapshot_download_siblings,
    )

    filenames = [s.rfilename for s in info.siblings if isinstance(s.rfilename, str)]
    filtered = snapshot_download_siblings(info.siblings)
    expected_files = [
        ExpectedFile(
            path = s.rfilename,
            size = int(getattr(s, "size", 0) or 0),
            sha256 = sibling_sha256(s),
        )
        for s in filtered
        if isinstance(s.rfilename, str)
    ]
    return resolve_snapshot_ignore_patterns_for_files(filenames), expected_files


def _dataset_expected_files(info) -> list:
    from hub.utils.download_manifest import ExpectedFile
    return [
        ExpectedFile(
            path = s.rfilename,
            size = int(getattr(s, "size", 0) or 0),
            sha256 = sibling_sha256(s),
        )
        for s in info.siblings
        if isinstance(s.rfilename, str)
    ]


def _recover_manifest_after_download(
    repo_type: RepoType,
    repo_id: str,
    snapshot_path: str,
    mode: str,
    *,
    fetch_info,
    expected_files_from_info,
    label: str = "",
) -> None:
    """Best-effort manifest write for a download whose metadata was unavailable
    at start: re-fetch and record the expected files, else fall back to the
    on-disk file list. Shared by the model and dataset workers.

    A pre-existing manifest is authoritative and is preserved untouched. This is
    load-bearing: when access is lost on resume (token revoked/changed on a
    gated/private repo), snapshot_download returns the cached partial snapshot
    WITHOUT downloading, so rebuilding the manifest from on-disk files would record
    the partial set as expected and let _verify_completed_download certify a
    half-finished download as complete.

    The same hazard exists with NO prior manifest. When metadata is still
    unavailable here, leftover ``.incomplete`` blobs prove a cached partial was
    returned without downloading, so we fail (exit 1) instead of deriving a
    self-certifying manifest, leaving the partial intact for a later resume. That
    signal misses a file that never started (no ``.incomplete``), so a kill
    between files is accepted optimistically from the on-disk subset; a later
    metadata-bearing attempt writes the true manifest and catches any shortfall."""
    from hub.utils import download_manifest
    from hub.utils.hf_cache_state import has_active_incomplete_blobs

    if download_manifest.read_manifest(repo_type, repo_id, None) is not None:
        return

    try:
        if download_manifest.write_manifest(
            repo_type,
            repo_id,
            None,
            expected_files_from_info(fetch_info()),
            mode,
        ):
            return
        reason = "manifest write failed"
    except Exception as e:
        reason = f"{type(e).__name__}: {e}"

    if has_active_incomplete_blobs(repo_type, repo_id):
        print(
            f"{label}could not reach Hugging Face for {repo_id} and the copy on "
            "disk is still incomplete. Access to a private or restricted repo may "
            "have been lost (HF token removed or changed), the connection dropped, "
            "or Hugging Face is temporarily unavailable. Set a valid HF token or "
            "reconnect, then resume the download.",
            file = sys.stderr,
        )
        sys.exit(1)

    fallback_files = download_manifest.expected_files_from_snapshot_dir(Path(snapshot_path))
    if fallback_files and download_manifest.write_manifest(
        repo_type,
        repo_id,
        None,
        fallback_files,
        mode,
    ):
        print(
            f"{label}could not record the metadata manifest for {repo_id}, "
            "recorded one from the downloaded files so completion "
            f"is tracked ({reason})",
            file = sys.stderr,
        )
    else:
        print(
            f"{label}could not record the metadata manifest for {repo_id}, "
            f"{download_manifest.MANIFEST_DEGRADED_MARKER} ({reason})",
            file = sys.stderr,
        )


def _download_snapshot(repo_id: str, hf_token: str | None, mode: str) -> None:
    from huggingface_hub import snapshot_download
    from hub.utils.download_registry import prepare_cache_for_transport
    from hub.utils import download_manifest

    # One metadata fetch powers both the ignore-pattern decision (drop
    # consolidated.* when transformers weights exist) and the manifest's
    # expected_files. A failure is non-fatal: fall back to the legacy
    # ignore-pattern set (keeping consolidated) and skip the manifest, so
    # download proceeds without the verification + partial detection it enables.
    try:
        info = _model_info_with_retry(repo_id, hf_token)
    except Exception as e:
        print(
            f"metadata unavailable, downloading full snapshot for {repo_id} "
            f"({type(e).__name__}: {e})",
            file = sys.stderr,
        )
        info = None

    download_manifest.clear_cancel_marker("model", repo_id, None)
    if info is not None:
        ignore_patterns, expected_files = _snapshot_download_plan(info)
        # Written for every transport. The manifest verifies the finalized
        # files under snapshots/, which both transports produce identically
        # (XET also renames a full, correctly-sized blob into place). XET's
        # block-level dedup lives only in the chunk-cache the manifest never
        # inspects, so per-file size verification is valid regardless of transport.
        download_manifest.write_manifest("model", repo_id, None, expected_files, mode)
    else:
        ignore_patterns = list(SNAPSHOT_IGNORE_PATTERNS)
        expected_files = []

    purged = prepare_cache_for_transport("model", repo_id, mode)
    if purged:
        print(
            f"Purged {purged} untrusted partial blob(s) for {repo_id} "
            f"before starting {mode} download.",
            file = sys.stderr,
        )
    _preflight_disk_space("model", repo_id, expected_files)
    snapshot_path = snapshot_download(
        repo_id = repo_id,
        token = _hf_token_arg(hf_token),
        ignore_patterns = ignore_patterns,
        max_workers = 1,
    )
    if info is None:
        _recover_manifest_after_download(
            "model",
            repo_id,
            snapshot_path,
            mode,
            fetch_info = lambda: _model_info_with_retry(repo_id, hf_token),
            expected_files_from_info = lambda recovered: _snapshot_download_plan(recovered)[1],
        )
    _verify_completed_download(
        "model",
        repo_id,
        None,
        snapshot_path,
        metadata_unavailable = info is None,
    )


def _gguf_variant_target_plan(
    repo_id: str, variant: str, hf_token: str | None
) -> GgufVariantPlan | None:
    try:
        info = _model_info_with_retry(repo_id, hf_token)
    except Exception as e:
        print(
            f"metadata unavailable, cannot resolve GGUF variant '{variant}' "
            f"for {repo_id} ({type(e).__name__}: {e})",
            file = sys.stderr,
        )
        raise RuntimeError(
            f"Metadata unavailable while resolving GGUF variant '{variant}' " f"for {repo_id}"
        ) from e
    return build_gguf_variant_plans(list(info.siblings)).get(variant.lower())


def _download_gguf_variant(repo_id: str, variant: str, hf_token: str | None, mode: str) -> None:
    from huggingface_hub import snapshot_download
    from hub.utils.download_registry import prepare_cache_for_transport
    from hub.utils.hf_cache_state import has_active_incomplete_blobs
    from hub.utils import download_manifest

    metadata_unavailable = False
    try:
        plan = _gguf_variant_target_plan(repo_id, variant, hf_token)
    except RuntimeError:
        plan = None
        metadata_unavailable = True

    if not metadata_unavailable:
        if plan is None:
            print(
                f"No GGUF shards matching variant '{variant}' in {repo_id}",
                file = sys.stderr,
            )
            sys.exit(1)
        targets = list(plan.target_filenames)
        expected_files = list(plan.expected_files)
        main_blob_hashes = plan.main_hashes
        companion_blob_hashes = plan.companion_hashes
        download_manifest.write_manifest(
            "model",
            repo_id,
            variant,
            expected_files,
            mode,
        )
    else:
        # Metadata unreachable (offline / gated / private). Resume the exact
        # shards the original attempt recorded so snapshot_download can range
        # over the surviving .incomplete blobs without a model_info call.
        manifest = download_manifest.read_manifest("model", repo_id, variant)
        if manifest is None or not manifest.expected_files:
            print(
                f"Metadata unavailable and no manifest to resume GGUF "
                f"variant '{variant}' for {repo_id}",
                file = sys.stderr,
            )
            sys.exit(1)
        plan = plan_from_expected_files(variant, manifest.expected_files)
        targets = list(plan.target_filenames)
        expected_files = list(plan.expected_files)
        download_manifest.write_manifest(
            "model",
            repo_id,
            variant,
            expected_files,
            mode,
        )
        main_blob_hashes = plan.main_hashes
        companion_blob_hashes = plan.companion_hashes
        print(
            f"Metadata unavailable; resuming GGUF variant '{variant}' for "
            f"{repo_id} from the existing manifest.",
            file = sys.stderr,
        )

    download_manifest.clear_cancel_marker("model", repo_id, variant)
    purge_blob_hashes = main_blob_hashes
    if not main_blob_hashes:
        if has_active_incomplete_blobs("model", repo_id):
            print(
                f"GGUF variant '{variant}' for {repo_id} has partial cache state "
                "but no resolvable blob hashes; delete the partial download or "
                "retry when metadata is available.",
                file = sys.stderr,
            )
            sys.exit(1)
        purge_blob_hashes = frozenset()
        print(
            f"GGUF variant '{variant}' for {repo_id} has no resolvable blob "
            "hashes; starting without partial cache reuse.",
            file = sys.stderr,
        )
    # Main quant blobs are owned by this variant (variant-scoped marker). The
    # shared vision companion (mmproj) is judged by a separate companion marker
    # and never purged while a concurrent peer is writing it.
    purged = prepare_cache_for_transport(
        "model",
        repo_id,
        mode,
        variant,
        only_blob_hashes = purge_blob_hashes,
        companion_blob_hashes = companion_blob_hashes,
        protected_blob_hashes = _protected_blob_hashes(),
    )
    if purged:
        print(
            f"Purged {purged} untrusted partial blob(s) for {repo_id} "
            f"before starting {mode} download.",
            file = sys.stderr,
        )
    _preflight_disk_space("model", repo_id, expected_files)
    snapshot_path = snapshot_download(
        repo_id = repo_id,
        token = _hf_token_arg(hf_token),
        allow_patterns = targets,
        max_workers = 1,
    )
    _verify_completed_download(
        "model",
        repo_id,
        variant,
        snapshot_path,
        metadata_unavailable = metadata_unavailable,
    )
    if plan is not None:
        try:
            from hub.services.models.deletion import reclaim_replaced_gguf_variant
            reclaim_replaced_gguf_variant(
                repo_id,
                variant,
                plan.main_hashes,
                hf_token,
                hub_cache = Path(snapshot_path).parents[2],
            )
        except Exception as e:
            print(
                f"Verified GGUF update for {repo_id} [{variant}], but stale-cache "
                f"reclaim failed ({type(e).__name__}: {e})",
                file = sys.stderr,
            )


def _download_scoped_snapshot(
    repo_id: str, scope: str, files: list[str], hf_token: str | None, mode: str
) -> None:
    """Fetch exactly ``files`` from ``repo_id``, keyed under ``scope``.

    For consumers that read a deliberate subset of a repo (the diffusion loader skips the
    packaged root single, transformer/ shards and fp16 twins). Keyed apart from the repo's
    full snapshot so neither manifest describes the other, and the repo is not later judged
    partial against expectations it was never meant to meet."""
    from huggingface_hub import HfApi, snapshot_download
    from hub.utils.download_registry import prepare_cache_for_transport
    from hub.utils import download_manifest
    from hub.utils.download_manifest import ExpectedFile

    wanted = set(files)
    try:
        info = _model_info_with_retry(repo_id, hf_token)
    except Exception as e:
        print(
            f"metadata unavailable for scoped download of {repo_id} " f"({type(e).__name__}: {e})",
            file = sys.stderr,
        )
        info = None

    expected_files: list[ExpectedFile] = []
    blob_hashes: frozenset[str] = frozenset()
    if info is not None:
        siblings = [s for s in info.siblings if getattr(s, "rfilename", None) in wanted]
        # Every requested file must resolve. Dropping an unmatched name would shrink the manifest and the
        # completion check to the survivors, and snapshot_download also succeeds when an allow pattern
        # matches nothing, so the job would report complete with a required file missing.
        missing = sorted(set(wanted) - {getattr(s, "rfilename", None) for s in siblings})
        if missing:
            print(
                f"Scoped download of {repo_id} cannot resolve "
                f"{len(missing)} requested file(s): {_format_path_list(missing)}",
                file = sys.stderr,
            )
            sys.exit(1)
        expected_files = [
            ExpectedFile(
                path = s.rfilename,
                size = int(getattr(s, "size", 0) or 0),
                sha256 = sibling_sha256(s),
            )
            for s in siblings
        ]
        from hub.utils.snapshot_filters import blob_hashes_for_siblings

        blob_hashes = blob_hashes_for_siblings(siblings)
        download_manifest.write_manifest("model", repo_id, scope, expected_files, mode)

    download_manifest.clear_cancel_marker("model", repo_id, scope)
    purged = prepare_cache_for_transport(
        "model",
        repo_id,
        mode,
        scope,
        only_blob_hashes = blob_hashes or None,
        protected_blob_hashes = _protected_blob_hashes(),
    )
    if purged:
        print(
            f"Purged {purged} untrusted partial blob(s) for {repo_id} [{scope}] "
            f"before starting {mode} download.",
            file = sys.stderr,
        )
    _preflight_disk_space("model", repo_id, expected_files)
    snapshot_path = snapshot_download(
        repo_id = repo_id,
        token = _hf_token_arg(hf_token),
        allow_patterns = files,
        max_workers = 1,
    )
    if info is None:
        # With no metadata there is no manifest, so _verify_completed_download below is a no-op -- and
        # snapshot_download RETURNS AN EXISTING SNAPSHOT FOLDER, fetching nothing, when its own repo_info
        # call also fails. A repo already on disk from a full snapshot job (which ignores *.gguf) would
        # therefore flip this job to complete having downloaded no weights. The requested list needs no
        # network, so check it against the disk directly.
        root = Path(snapshot_path)
        absent = tuple(f for f in files if not (root / f).exists())
        if absent:
            print(
                f"Could not reach Hugging Face for {repo_id} [{scope}] and the copy on disk "
                f"is incomplete ({len(absent)} file(s) missing): {_format_path_list(absent)}. "
                "Reconnect (or set a valid HF token) and resume the download.",
                file = sys.stderr,
            )
            sys.exit(1)
    _verify_completed_download(
        "model",
        repo_id,
        scope,
        snapshot_path,
        metadata_unavailable = info is None,
    )


def _download_dataset(repo_id: str, hf_token: str | None, mode: str) -> None:
    from huggingface_hub import snapshot_download
    from hub.utils.download_registry import prepare_cache_for_transport
    from hub.utils import download_manifest

    try:
        info = _dataset_info_with_retry(repo_id, hf_token)
    except Exception as e:
        print(
            f"dataset metadata unavailable, downloading full dataset for {repo_id} "
            f"({type(e).__name__}: {e})",
            file = sys.stderr,
        )
        info = None
    # Cancel-marker clear and manifest write run on every transport. See
    # _download_snapshot for why per-file size verification is valid under XET.
    download_manifest.clear_cancel_marker("dataset", repo_id, None)
    if info is not None:
        expected_files = _dataset_expected_files(info)
        download_manifest.write_manifest(
            "dataset",
            repo_id,
            None,
            expected_files,
            mode,
        )
    else:
        expected_files = []
    purged = prepare_cache_for_transport("dataset", repo_id, mode)
    if purged:
        print(
            f"Purged {purged} untrusted partial blob(s) for {repo_id} "
            f"before starting {mode} download.",
            file = sys.stderr,
        )
    _preflight_disk_space("dataset", repo_id, expected_files)
    snapshot_path = snapshot_download(
        repo_id = repo_id,
        token = _hf_token_arg(hf_token),
        repo_type = "dataset",
        max_workers = 1,
    )
    if info is None:
        _recover_manifest_after_download(
            "dataset",
            repo_id,
            snapshot_path,
            mode,
            fetch_info = lambda: _dataset_info_with_retry(repo_id, hf_token),
            expected_files_from_info = _dataset_expected_files,
            label = "dataset ",
        )
    _verify_completed_download(
        "dataset",
        repo_id,
        None,
        snapshot_path,
        metadata_unavailable = info is None,
    )


def main() -> None:
    parser = argparse.ArgumentParser(description = "HuggingFace Hub download worker")
    parser.add_argument("--repo-id", required = True)
    parser.add_argument("--variant", default = None)
    parser.add_argument("--dataset", action = "store_true")
    parser.add_argument("--transport", choices = ("http", "xet"), default = "http")
    parser.add_argument("--parent-pid", type = int, default = None)
    parser.add_argument(
        "--files-json",
        default = None,
        help = "Temp JSON file holding a scoped job's exact file list (deleted after reading).",
    )
    args = parser.parse_args()

    scoped_files: list[str] = []
    if args.files_json:
        import json
        try:
            with open(args.files_json, encoding = "utf-8") as handle:
                scoped_files = [str(f) for f in json.load(handle)]
        finally:
            try:
                os.unlink(args.files_json)
            except OSError:
                pass

    _install_signal_handlers()
    _install_parent_death_watchdog(args.parent_pid)

    hf_token = os.environ.get("HF_TOKEN") or None

    try:
        if args.dataset:
            _download_dataset(args.repo_id, hf_token, args.transport)
        elif scoped_files:
            _download_scoped_snapshot(
                args.repo_id, args.variant, scoped_files, hf_token, args.transport
            )
        elif args.variant:
            _download_gguf_variant(args.repo_id, args.variant, hf_token, args.transport)
        else:
            _download_snapshot(args.repo_id, hf_token, args.transport)
        sys.exit(0)
    except SystemExit:
        raise
    except Exception as e:
        # Surface a precise message so the UI doesn't show a generic "worker
        # exited with code 1". huggingface_hub's consistency check recommends
        # force_download=True to recover, which our "Restart" UI maps to a fresh
        # start by purging the partial via prepare_cache_for_transport.
        print(f"{type(e).__name__}: {e}", file = sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()