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

"""HF cache inventory scanner.

Read-only walks of the HuggingFace hub cache plus legacy/default
cache locations. Builds the foundation that Hub inventory endpoints
and the DownloadRegistry both consume.

The worker spawn / transport-marker preparation / DownloadRegistry
layers built on top of these primitives live in download_registry.py.
"""

from __future__ import annotations

import hashlib
import re
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Optional

from loggers import get_logger

logger = get_logger(__name__)

from hub.utils.gguf import (
    extract_quant_label,
    is_gguf_filename,
    is_mmproj_filename,
    is_mtp_drafter_path,
)
from hub.utils.state_dir import RepoType

from hub.utils.hf_cache_state import (
    INCOMPLETE_SUFFIX,
    has_incomplete_blobs,
    hf_cache_roots,
    iter_repo_cache_dirs,
    latest_snapshot_dir,
    repo_cache_dir_has_incomplete_blobs,
)

# Inventory is invalidated explicitly on every app-driven cache mutation, so
# this TTL only bounds staleness from out-of-band edits while skipping re-walks
# on rapid UI navigation.
_HF_CACHE_SCANS_TTL_SECONDS = 15.0
_GGUF_SPLIT_RE = re.compile(r"-(\d{3,})-of-(\d{3,})(?=\.gguf$)", re.IGNORECASE)
_hf_cache_scans_lock = threading.Lock()


@dataclass
class _HfCacheScanFlight:
    event: threading.Event
    epoch: int
    result: Optional[list] = None
    error: Optional[BaseException] = None


_hf_cache_scans_flight: Optional[_HfCacheScanFlight] = None
_hf_cache_scans_result: Optional[list] = None
_hf_cache_scans_cached_at: float = 0.0
# Bumped on every invalidation. A scan tags itself with the epoch it began
# under; an invalidation mid-scan changes the epoch so the in-flight result is
# neither cached nor served to callers that arrived after the mutation.
_hf_cache_scans_epoch: int = 0


def invalidate_hf_cache_scans() -> None:
    global _hf_cache_scans_result, _hf_cache_scans_cached_at, _hf_cache_scans_epoch
    with _hf_cache_scans_lock:
        _hf_cache_scans_result = None
        _hf_cache_scans_cached_at = 0.0
        _hf_cache_scans_epoch += 1


def all_hf_cache_scans() -> list:
    global _hf_cache_scans_flight, _hf_cache_scans_result, _hf_cache_scans_cached_at

    now = time.monotonic()
    with _hf_cache_scans_lock:
        if (
            _hf_cache_scans_result is not None
            and (now - _hf_cache_scans_cached_at) < _HF_CACHE_SCANS_TTL_SECONDS
        ):
            return list(_hf_cache_scans_result)
        start_epoch = _hf_cache_scans_epoch
        flight = _hf_cache_scans_flight
        # Only coalesce onto an in-flight scan from the current epoch; one that
        # began before an intervening invalidation is superseded so
        # post-mutation callers never receive pre-mutation data.
        if flight is None or flight.epoch != start_epoch:
            flight = _HfCacheScanFlight(event = threading.Event(), epoch = start_epoch)
            _hf_cache_scans_flight = flight
            owner = True
        else:
            owner = False

    if not owner:
        flight.event.wait()
        if flight.error is not None:
            raise flight.error
        return list(flight.result or [])

    try:
        scans = _compute_all_hf_cache_scans()
        with _hf_cache_scans_lock:
            flight.result = scans
            if _hf_cache_scans_epoch == flight.epoch:
                _hf_cache_scans_result = scans
                _hf_cache_scans_cached_at = time.monotonic()
        return scans
    except Exception as exc:
        with _hf_cache_scans_lock:
            if _hf_cache_scans_epoch == flight.epoch:
                _hf_cache_scans_result = None
                _hf_cache_scans_cached_at = 0.0
            flight.error = exc
        raise
    finally:
        with _hf_cache_scans_lock:
            if _hf_cache_scans_flight is flight:
                _hf_cache_scans_flight = None
            flight.event.set()


def _compute_all_hf_cache_scans() -> list:
    from huggingface_hub import scan_cache_dir

    scans: list = []
    for cache_root in hf_cache_roots():
        try:
            scans.append(scan_cache_dir(cache_dir = str(cache_root)))
        except Exception as exc:
            logger.warning("Could not scan HF cache %s: %s", cache_root, exc)
    return scans


def token_fingerprint(hf_token: Optional[str]) -> str:
    """16-char SHA256 prefix used as a cache-key qualifier for gated repos.

    Lets per-token size/snapshot caches refuse to serve a previously
    fetched value back to a different token (a private/gated repo's
    metadata is only valid for the credential that fetched it).
    """
    if not hf_token:
        return ""
    return hashlib.sha256(hf_token.encode()).hexdigest()[:16]


def resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]:
    """Pick the most useful on-disk path for a HF cache repo dir.

    Prefers the most-recent snapshot dir (what ``from_pretrained``
    actually points at). Falls back to the cache repo root. Returns the
    resolved realpath so symlinks under ``snapshots/`` are followed back
    to ``blobs/``.
    """
    try:
        latest = latest_snapshot_dir(repo_dir)
        if latest is not None:
            return str(latest.resolve())
        return str(repo_dir.resolve())
    except Exception:
        return None


def resolve_snapshot_dir_for_scan(
    repo_type: str,
    repo_id: str,
    repo_cache_dir: Optional[Path] = None,
) -> Optional[Path]:
    """Latest snapshot dir for a cache row, or the first populated HF cache root.

    Scanner-side counterpart to snapshot_download()'s return value (which the
    scanner cannot access). With a *repo_cache_dir*, returns its newest
    snapshot. Otherwise scans roots in priority order (active, legacy, default)
    and returns the newest snapshot in the first root that holds one; active is
    where snapshot_download writes, so it is authoritative. Within a root,
    picks by mtime (what from_pretrained resolves to) rather than refs/main,
    since the user may have downloaded a non-main commit.
    """
    if repo_cache_dir is not None:
        latest = latest_snapshot_dir(repo_cache_dir)
        if latest is None:
            return None
        try:
            return latest.resolve()
        except OSError:
            return None
    for repo_dir in iter_repo_cache_dirs(repo_type, repo_id):
        latest = latest_snapshot_dir(repo_dir)
        if latest is None:
            continue
        try:
            return latest.resolve()
        except OSError:
            continue
    return None


def _compose_partial(*signals: Callable[[], bool]) -> bool:
    return any(signal() for signal in signals)


def _hub_cache_for_repo_dir(repo_cache_dir: Optional[Path]) -> Optional[Path]:
    return repo_cache_dir.parent if repo_cache_dir is not None else None


def _legacy_partial(
    repo_type: str,
    repo_id: str,
    repo_cache_dir: Optional[Path] = None,
) -> bool:
    if repo_cache_dir is not None:
        return repo_cache_dir_has_incomplete_blobs(repo_cache_dir)
    return has_incomplete_blobs(repo_type, repo_id)


def _repo_cache_dir_incomplete_hashes(repo_cache_dir: Path) -> set[str]:
    blobs_dir = repo_cache_dir / "blobs"
    if not blobs_dir.is_dir():
        return set()
    hashes: set[str] = set()
    try:
        entries = list(blobs_dir.iterdir())
    except OSError:
        return hashes
    for blob in entries:
        try:
            if blob.is_file() and blob.name.endswith(INCOMPLETE_SUFFIX):
                hashes.add(blob.name[: -len(INCOMPLETE_SUFFIX)])
        except OSError:
            continue
    return hashes


def _repo_cache_dir_has_non_gguf_broken_snapshot_symlinks(repo_cache_dir: Path) -> bool:
    latest = latest_snapshot_dir(repo_cache_dir)
    if latest is None:
        return False
    try:
        entries = list(latest.rglob("*"))
    except OSError:
        return False
    for entry in entries:
        try:
            if not entry.is_symlink() or entry.exists():
                continue
            rel = entry.relative_to(latest).as_posix()
            if is_gguf_filename(rel):
                continue
            return True
        except OSError:
            continue
    return False


def _gguf_variant_manifest_blob_hashes(
    repo_id: str, repo_cache_dir: Optional[Path] = None
) -> frozenset[str]:
    from hub.utils import download_manifest

    hashes: set[str] = set()
    hub_cache = _hub_cache_for_repo_dir(repo_cache_dir)
    for variant, _path in download_manifest.iter_variant_manifests(
        "model",
        repo_id,
        hub_cache = hub_cache,
    ):
        manifest = download_manifest.read_manifest(
            "model",
            repo_id,
            variant,
            hub_cache = hub_cache,
        )
        if manifest is None:
            continue
        for expected in manifest.expected_files:
            if expected.sha256 and is_gguf_filename(expected.path):
                hashes.add(expected.sha256)
    return frozenset(hashes)


def _repo_cache_dir_has_snapshot_legacy_partial(
    repo_cache_dir: Path, *, ignored_blob_hashes: frozenset[str]
) -> bool:
    incomplete_hashes = _repo_cache_dir_incomplete_hashes(repo_cache_dir)
    if any(blob_hash not in ignored_blob_hashes for blob_hash in incomplete_hashes):
        return True
    return _repo_cache_dir_has_non_gguf_broken_snapshot_symlinks(repo_cache_dir)


def _snapshot_legacy_partial(
    repo_type: str,
    repo_id: str,
    repo_cache_dir: Optional[Path] = None,
) -> bool:
    if repo_type != "model":
        return _legacy_partial(repo_type, repo_id, repo_cache_dir)
    ignored_hashes = _gguf_variant_manifest_blob_hashes(repo_id, repo_cache_dir)
    if repo_cache_dir is not None:
        return _repo_cache_dir_has_snapshot_legacy_partial(
            repo_cache_dir,
            ignored_blob_hashes = ignored_hashes,
        )
    return any(
        _repo_cache_dir_has_snapshot_legacy_partial(
            entry,
            ignored_blob_hashes = ignored_hashes,
        )
        for entry in iter_repo_cache_dirs(repo_type, repo_id)
    )


def _completed_gguf_variants(snapshot_dir: Optional[Path]) -> set[str]:
    if snapshot_dir is None:
        return set()
    complete: set[str] = set()
    split_groups: dict[str, dict[int, set[int]]] = {}
    try:
        paths = list(snapshot_dir.rglob("*"))
    except OSError:
        return set()
    for path in paths:
        try:
            if not path.is_file() or path.stat().st_size <= 0:
                continue
        except OSError:
            continue
        rel = path.relative_to(snapshot_dir).as_posix()
        if not is_gguf_filename(rel) or is_mmproj_filename(rel) or is_mtp_drafter_path(rel):
            continue
        quant = extract_quant_label(rel)
        split = _GGUF_SPLIT_RE.search(path.name)
        if split is None:
            complete.add(quant)
            continue
        index = int(split.group(1))
        total = int(split.group(2))
        if index <= 0 or total <= 0 or index > total:
            continue
        split_groups.setdefault(quant, {}).setdefault(total, set()).add(index)
    for quant, groups in split_groups.items():
        for total, indices in groups.items():
            if indices == set(range(1, total + 1)):
                complete.add(quant)
                break
    return complete


def _manifest_partial(
    repo_type: RepoType,
    repo_id: str,
    variant: Optional[str] = None,
    snapshot_dir: Optional[Path] = None,
    repo_cache_dir: Optional[Path] = None,
) -> bool:
    from hub.utils import download_manifest

    manifest = download_manifest.read_manifest(
        repo_type,
        repo_id,
        variant,
        hub_cache = _hub_cache_for_repo_dir(repo_cache_dir),
    )
    if manifest is None:
        return False
    resolved = (
        snapshot_dir
        if snapshot_dir is not None
        else resolve_snapshot_dir_for_scan(repo_type, repo_id, repo_cache_dir)
    )
    if resolved is None:
        return True
    if repo_type == "model" and variant is not None:
        if download_manifest.verify_against_disk(manifest, resolved).ok:
            return False
        for candidate in _manifest_snapshot_dirs(repo_type, repo_id, repo_cache_dir):
            if candidate == resolved:
                continue
            if download_manifest.verify_against_disk(manifest, candidate).ok:
                return False
        return True
    return not download_manifest.verify_against_disk(manifest, resolved).ok


def _manifest_snapshot_dirs(
    repo_type: RepoType,
    repo_id: str,
    repo_cache_dir: Optional[Path] = None,
) -> list[Path]:
    repo_dirs = (
        [repo_cache_dir]
        if repo_cache_dir is not None
        else list(iter_repo_cache_dirs(repo_type, repo_id))
    )
    snapshots: list[Path] = []
    seen: set[str] = set()
    for repo_dir in repo_dirs:
        if repo_dir is None:
            continue
        snapshots_dir = repo_dir / "snapshots"
        try:
            if not snapshots_dir.is_dir():
                continue
            entries = list(snapshots_dir.iterdir())
        except OSError:
            continue
        for entry in entries:
            try:
                if not entry.is_dir():
                    continue
                resolved = entry.resolve()
            except OSError:
                continue
            key = str(resolved)
            if key in seen:
                continue
            seen.add(key)
            snapshots.append(resolved)
    return snapshots


def is_snapshot_partial(
    repo_type: RepoType,
    repo_id: str,
    repo_cache_dir: Optional[Path] = None,
) -> bool:
    """Repo-row partial flag for snapshot-style downloads (full-snapshot
    models — safetensors/adapter/checkpoint — and all datasets).

    Composes three signals, cheapest first:
      1. Cancel marker (single stat).
      2. Snapshot-attributed legacy .incomplete blob / broken-symlink check.
      3. Manifest walk (stat per expected file under the latest snapshot).

    A manifest without a resolvable snapshot is partial: the worker got
    far enough to record expectations but did not leave a usable snapshot."""
    from hub.utils import download_manifest
    return _compose_partial(
        lambda: download_manifest.has_cancel_marker(
            repo_type,
            repo_id,
            None,
            hub_cache = _hub_cache_for_repo_dir(repo_cache_dir),
        ),
        lambda: _snapshot_legacy_partial(repo_type, repo_id, repo_cache_dir),
        lambda: _manifest_partial(
            repo_type,
            repo_id,
            None,
            None,
            repo_cache_dir,
        ),
    )


def _current_revisions(repo_info):
    """The revisions to judge a cached repo by: just the one the loader will actually open.

    ``from_pretrained`` resolves the newest snapshot by mtime (:func:`latest_snapshot_dir`), so a
    repo cached twice -- an older complete snapshot plus a newer companion-only scoped one -- must
    be judged on the newer one. Scanning every revision let the old snapshot's denoiser satisfy the
    completeness check, so the row read as complete while the snapshot the loader picks has no
    transformer/unet: offline loads then fail and online loads silently pull the multi-GB weight.

    Falls back to the newest revision by ``last_modified``, then to every revision, so a cache
    layout this cannot resolve behaves as before rather than reporting nothing.
    """
    revisions = list(getattr(repo_info, "revisions", ()) or ())
    if len(revisions) <= 1:
        return revisions
    repo_path = getattr(repo_info, "repo_path", None)
    if repo_path is not None:
        latest = latest_snapshot_dir(Path(repo_path))
        if latest is not None:
            scoped = [
                rev
                for rev in revisions
                if getattr(rev, "snapshot_path", None) is not None
                and Path(rev.snapshot_path) == latest
            ]
            if scoped:
                return scoped
    dated = [rev for rev in revisions if getattr(rev, "last_modified", None) is not None]
    if dated:
        return [max(dated, key = lambda rev: rev.last_modified)]
    return revisions


def repo_has_pipeline_index(repo_info) -> bool:
    """Whether the cached snapshot carries a ROOT model_index.json, i.e. is loadable
    as a full diffusers pipeline (from_pretrained reads only the repo root). A nested
    subdir/model_index.json does not count: loading the repo root still fails, so the
    row must keep its single_file flag. CachedFileInfo.file_name is the basename, so
    a name match alone would also claim nested copies -- scope by file_path when the
    scan provides it."""
    try:
        for rev in _current_revisions(repo_info):
            snapshot = getattr(rev, "snapshot_path", None)
            for f in rev.files:
                name = str(getattr(f, "file_name", "") or "")
                path = getattr(f, "file_path", None)
                if path is not None and snapshot is not None:
                    p = Path(path)
                    if p.name == "model_index.json" and p.parent == Path(snapshot):
                        return True
                elif name == "model_index.json":
                    return True
    except Exception:
        pass
    return False


def repo_pipeline_missing_denoiser(repo_info) -> bool:
    """True for a diffusers-pipeline snapshot (root ``model_index.json``) whose denoiser component
    (``transformer/`` or ``unet/``) carries NO weight file -- the shape of a companion-only prefetch
    where a GGUF image load pulled the base repo's VAE / text-encoder / manifest but skipped the
    multi-GB transformer (the GGUF supplies it). :func:`is_snapshot_partial` misses this (every
    file the manifest expected did arrive), so callers OR the two signals together and mark such
    rows partial. Best-effort: any scan error reports not-missing so a glitch never hides a
    genuinely complete pipeline."""
    if not repo_has_pipeline_index(repo_info):
        return False
    _DENOISER_DIRS = ("transformer", "unet")
    _WEIGHT_SUFFIXES = (".safetensors", ".bin")
    try:
        for rev in _current_revisions(repo_info):
            snapshot = getattr(rev, "snapshot_path", None)
            for f in rev.files:
                name = str(getattr(f, "file_name", "") or "")
                path = getattr(f, "file_path", None)
                parts: tuple[str, ...] = ()
                if path is not None and snapshot is not None:
                    try:
                        parts = Path(path).relative_to(Path(snapshot)).parts
                    except ValueError:
                        parts = ()
                if not parts:
                    # No snapshot scoping: fall back to the recorded name, which may itself carry the component
                    # subdir (e.g. 'transformer/model...').
                    parts = Path(name).parts
                if (
                    len(parts) >= 2
                    and parts[0].lower() in _DENOISER_DIRS
                    and parts[-1].lower().endswith(_WEIGHT_SUFFIXES)
                ):
                    return False
        return True
    except Exception:
        return False


def is_variant_partial(
    repo_id: str,
    variant: str,
    snapshot_dir: Optional[Path] = None,
    *,
    incomplete_blob_hashes: Optional[set[str]] = None,
    variant_blob_hashes: Optional[frozenset[str]] = None,
    repo_cache_dir: Optional[Path] = None,
) -> bool:
    """Per-variant partial detection. Owns its manifest, owns its marker.
    Used by the GGUF variants endpoint to flag a specific quant as broken
    without contaminating other quants in the same repo.

    snapshot_dir is an optional hint to avoid re-walking the cache when a
    caller is checking many variants of the same repo (see
    is_gguf_repo_partial for that usage)."""
    from hub.utils import download_manifest
    return _compose_partial(
        lambda: download_manifest.has_cancel_marker(
            "model",
            repo_id,
            variant,
            hub_cache = _hub_cache_for_repo_dir(repo_cache_dir),
        ),
        lambda: bool(
            incomplete_blob_hashes
            and variant_blob_hashes
            and incomplete_blob_hashes.intersection(variant_blob_hashes)
        ),
        lambda: _manifest_partial(
            "model",
            repo_id,
            variant,
            snapshot_dir,
            repo_cache_dir,
        ),
    )


def is_gguf_repo_partial(repo_id: str, repo_cache_dir: Optional[Path] = None) -> bool:
    """Repo-row partial flag for a GGUF repo. The inventory shows ONE row per
    GGUF repo (requires_variant=True); per-variant detail lives in
    GET /api/models/gguf-variants and uses is_variant_partial.

    *** DO NOT simplify this to "any variant partial -> repo partial" ***

    Tripwire scenario: user downloads Q8_0 fully, then starts Q4_K_M and
    cancels. Both variants share ONE inventory row. If row.partial flips True,
    _capabilities_for_format flips can_chat=False, so the user can no longer
    chat with the perfectly-good Q8_0 because of an unrelated cancelled Q4_K_M.

    Correct semantics: partial=True only when at least one variant is broken
    AND no other variant is clean. "Simplifying" to the obvious "any broken"
    form re-introduces this Q8+Q4 mixed-state regression.

    Composes signals:
      1. Cheap legacy fast-path (.incomplete blobs / broken symlinks).
      2. Per-variant manifest + marker enumeration, gated on "all broken".
    """
    from hub.utils import download_manifest

    has_legacy_partial = _legacy_partial("model", repo_id, repo_cache_dir)
    snapshot_dir = resolve_snapshot_dir_for_scan(
        "model",
        repo_id,
        repo_cache_dir,
    )
    variants: set[str] = set(_completed_gguf_variants(snapshot_dir))
    hub_cache = _hub_cache_for_repo_dir(repo_cache_dir)
    for variant, _path in download_manifest.iter_variant_manifests(
        "model",
        repo_id,
        hub_cache = hub_cache,
    ):
        if (
            download_manifest.read_manifest(
                "model",
                repo_id,
                variant,
                hub_cache = hub_cache,
            )
            is not None
        ):
            variants.add(variant)
    for variant, _path in download_manifest.iter_variant_markers(
        "model",
        repo_id,
        hub_cache = hub_cache,
    ):
        if download_manifest.has_cancel_marker(
            "model",
            repo_id,
            variant,
            hub_cache = hub_cache,
        ):
            variants.add(variant)
    if not variants:
        return has_legacy_partial
    has_clean = False
    has_broken = has_legacy_partial
    for variant in variants:
        if is_variant_partial(
            repo_id,
            variant,
            snapshot_dir,
            repo_cache_dir = repo_cache_dir,
        ):
            has_broken = True
        else:
            has_clean = True
    return has_broken and not has_clean


def partial_transport_for(
    repo_type: RepoType,
    repo_id: str,
    variant: Optional[str] = None,
    repo_cache_dir: Optional[Path] = None,
) -> Optional[str]:
    """Transport to surface on a partial row's resume affordance.

    Prefers the cancel marker's transport, then the manifest's. The fallback
    matters for rows partial without a marker (an errored/interrupted download
    leaves the manifest but no marker) so the UI can still show HTTP-resume vs
    XET-redownload instead of the neutral retry label. ``None`` when neither is
    available."""
    from hub.utils import download_manifest

    hub_cache = _hub_cache_for_repo_dir(repo_cache_dir)
    marker_transport = download_manifest.read_cancel_marker_transport(
        repo_type,
        repo_id,
        variant,
        hub_cache = hub_cache,
    )
    if marker_transport is not None:
        return marker_transport
    manifest = download_manifest.read_manifest(
        repo_type,
        repo_id,
        variant,
        hub_cache = hub_cache,
    )
    return manifest.transport if manifest is not None else None