SnapOtter/packages/ai/python/install_feature.py

"""Pre-built AI bundle installer for SnapOtter.

Downloads a pre-built tar.gz archive (or uses a local file), verifies its
SHA256 checksum, extracts site-packages and models, and writes installed.json.

Invoked by the Node.js backend as a subprocess.

Usage:
    python3 install_feature.py <bundleId> <manifestPath> <modelsDir>

Progress is reported via JSON lines on stderr (parsed by the Node bridge).
Final result is a JSON object on stdout.
"""

import csv
import errno
import glob
import hashlib
import importlib
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import tarfile
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone

DOWNLOAD_CHUNK_BYTES = 4 * 1024 * 1024
DOWNLOAD_META_BYTES = 64 * 1024 * 1024


# -- Helpers --

def emit_progress(percent: int, stage: str) -> None:
    """Emit a progress update via stderr JSON line."""
    sys.stderr.write(json.dumps({"progress": percent, "stage": stage}) + "\n")
    sys.stderr.flush()


def fail(message: str) -> None:
    """Print error to stderr and exit non-zero."""
    sys.stderr.write(json.dumps({"error": message}) + "\n")
    sys.stderr.flush()
    sys.exit(1)


# -- Architecture detection --

def detect_arch() -> str:
    """Return the bundle archive key for this host.

    Only two archive variants are currently published to the bundle repo:
    'amd64-gpu' and 'arm64-cpu' (see deepsafe/feature-bundles). There is no
    CPU-only amd64 variant yet, so amd64 hosts always resolve to 'amd64-gpu'
    even when no GPU is present: this downloads working CUDA-capable
    packages, just larger than a CPU-only host strictly needs. Do not change
    this to branch on GPU presence without first publishing an 'amd64-cpu'
    archive for every bundle; requesting a key that doesn't exist in the
    manifest fails the install outright (see the archives.get(arch) lookup
    below), which would be worse than the current oversized-but-working
    download.
    """
    machine = platform.machine().lower()
    if machine in ("aarch64", "arm64"):
        return "arm64-cpu"
    return "amd64-gpu"


# -- Disk space --

def _existing_ancestor(path: str) -> str:
    """Nearest existing ancestor of path (so disk_usage never raises on a
    not-yet-created dir like the venv)."""
    p = os.path.abspath(path)
    while p and not os.path.exists(p):
        parent = os.path.dirname(p)
        if parent == p:
            break
        p = parent
    return p or "/"


def check_disk_space(path: str, needed_bytes: int) -> None:
    """Fail if insufficient disk space on the filesystem holding path."""
    usage = shutil.disk_usage(_existing_ancestor(path))
    if usage.free < needed_bytes:
        free_gb = usage.free / (1024 ** 3)
        need_gb = needed_bytes / (1024 ** 3)
        fail(
            f"Insufficient disk space: need {need_gb:.1f} GB, "
            f"have {free_gb:.1f} GB free. "
            f"Free up space and retry."
        )


def estimate_extracted(compressed: int, extracted: int) -> int:
    """Extracted-size estimate for the disk preflight. When the manifest omits
    extractedSize (0), the budget would otherwise collapse to just the
    compressed size and under-reserve for the extracted payload; fall back to a
    conservative 3x of compressed (measured extracted/compressed ratios reach
    ~3x). This is only the early sanity bail; the accurate guard is the
    real-on-disk re-check just before the destructive venv write."""
    return extracted if extracted > 0 else compressed * 3


def dir_size(path: str) -> int:
    """Total size in bytes of all files under path (best-effort)."""
    total = 0
    for root, _dirs, files in os.walk(path):
        for f in files:
            try:
                total += os.path.getsize(os.path.join(root, f))
            except OSError:
                pass
    return total


def same_filesystem(a: str, b: str) -> bool:
    """True if paths a and b live on the same filesystem (so a rename between
    them is a cheap metadata op rather than a full copy)."""
    try:
        return os.stat(_existing_ancestor(a)).st_dev == os.stat(_existing_ancestor(b)).st_dev
    except OSError:
        return False


# -- Venv site-packages discovery --

def get_site_packages_dir(venv_path: str) -> str:
    """Find the site-packages directory inside a Python venv."""
    matches = glob.glob(os.path.join(venv_path, "lib", "python*", "site-packages"))
    if matches:
        return matches[0]
    return ""


# -- SHA256 verification --

def verify_sha256(filepath: str, expected: str) -> bool:
    """Stream-hash a file and compare to expected hex digest."""
    h = hashlib.sha256()
    with open(filepath, "rb") as f:
        while True:
            chunk = f.read(8192)
            if not chunk:
                break
            h.update(chunk)
    return h.hexdigest() == expected


# -- Download with resume --

def _set_env_temporarily(key: str, value: str):
    previous = os.environ.get(key)
    os.environ[key] = value
    return previous


def _restore_env(key: str, previous) -> None:
    if previous is None:
        os.environ.pop(key, None)
    else:
        os.environ[key] = previous


def _cleanup_hf_local_dir(local_dir: str, archive_file: str) -> None:
    if "/" in archive_file:
        top_level = archive_file.split("/", 1)[0]
        shutil.rmtree(os.path.join(local_dir, top_level), ignore_errors=True)
    shutil.rmtree(os.path.join(local_dir, ".cache"), ignore_errors=True)


def ensure_hf_hub(venv_path: str) -> None:
    """Guarantee the accelerated Hugging Face client is importable before the
    download so the multi-GB bundle transfer takes the fast Xet path.

    The installer runs under the on-disk venv (PYTHON_VENV_PATH, i.e.
    /data/ai/venv in Docker). That venv is normally seeded from the image's
    /opt/venv, which bakes huggingface-hub[hf_xet]. But an install whose venv
    predates the base package (an upgrade where the reseed stamp didn't move, a
    hand-copied or offline-imported venv) would import-fail in
    download_with_hf_hub and silently fall back to the slow single-stream urllib
    downloader. Self-heal by pip-installing the client into this same venv.

    A bundle install already requires network and lifts the offline guard (see
    main()), so this adds no new offline dependency; if the pip install fails we
    fall through to the resumable urllib downloader, the correct degraded path.
    """
    try:
        import huggingface_hub  # noqa: F401

        return
    except Exception:
        pass

    python_path = os.path.join(venv_path, "bin", "python3")
    if not os.path.exists(python_path):
        return

    emit_progress(1, "Preparing accelerated download client...")
    try:
        subprocess.run(
            [
                python_path, "-m", "pip", "install", "--quiet",
                "huggingface-hub[hf_xet,hf_transfer]==0.36.2",
            ],
            capture_output=True, text=True, timeout=300, check=True,
        )
        # The finder caches the venv's site-packages listing; drop it so the
        # just-installed package is visible to the import in download_with_hf_hub.
        importlib.invalidate_caches()
    except Exception as e:
        emit_progress(1, f"Accelerated client unavailable ({e}); using resumable download.")


def download_with_hf_hub(
    bundle_repo: str,
    archive_file: str,
    dest: str,
    expected_size: int,
    progress_start: int,
    progress_end: int,
    force_download: bool = False,
) -> bool:
    """Download through huggingface_hub when available.

    huggingface_hub 0.32+ can use hf_xet for faster large-file transfers and
    manages retries/resume internally. Return False when the client is missing
    or fails so callers can fall back to the manual urllib downloader.
    """
    try:
        from huggingface_hub import hf_hub_download
    except Exception:
        return False

    # Enable hf_transfer (Rust multi-connection downloader) ONLY when the
    # package is actually importable. For a Xet-backed repo hf_xet takes
    # precedence and this is a no-op, but if the Xet CAS endpoint is unreachable
    # (e.g. a firewall that allows huggingface.co but blocks transfer.xethub.hf.co)
    # hf_hub_download falls back to plain HTTP, and hf_transfer makes that
    # fallback multi-connection instead of single-stream. Gating on the import
    # avoids the "HF_HUB_ENABLE_HF_TRANSFER set but package missing" hard error
    # on a venv that only has hf_xet.
    try:
        import hf_transfer  # noqa: F401

        os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
    except Exception:
        pass

    local_dir = os.path.dirname(dest)
    os.makedirs(local_dir, exist_ok=True)
    emit_progress(progress_start, "Downloading with accelerated Hugging Face client...")

    previous_progress = _set_env_temporarily("HF_HUB_DISABLE_PROGRESS_BARS", "1")
    try:
        downloaded_path = hf_hub_download(
            repo_id=bundle_repo,
            filename=archive_file,
            repo_type="model",
            local_dir=local_dir,
            force_download=force_download,
        )
    except Exception as e:
        emit_progress(
            progress_start,
            f"Accelerated download unavailable, using resumable fallback: {e}",
        )
        # Reclaim any partial blob/metadata hf_hub_download staged under
        # local_dir/.cache so the urllib fallback starts clean and disk is freed.
        _cleanup_hf_local_dir(local_dir, archive_file)
        return False
    finally:
        _restore_env("HF_HUB_DISABLE_PROGRESS_BARS", previous_progress)

    try:
        if not os.path.exists(downloaded_path):
            emit_progress(
                progress_start,
                "Accelerated download did not produce an archive, using resumable fallback...",
            )
            return False

        if os.path.abspath(downloaded_path) != os.path.abspath(dest):
            if os.path.exists(dest):
                os.unlink(dest)
            os.replace(downloaded_path, dest)

        size = os.path.getsize(dest)
        if expected_size > 0:
            pct = min(size / expected_size, 1.0)
            progress = int(progress_start + pct * (progress_end - progress_start))
        else:
            progress = progress_end
        emit_progress(
            min(progress, progress_end),
            f"Downloaded with accelerated client ({size / (1024**3):.1f} GB)",
        )
        return True
    except Exception as e:
        emit_progress(
            progress_start,
            f"Accelerated download post-processing failed, using resumable fallback: {e}",
        )
        return False
    finally:
        # Always drop the hf staging tree (local_dir/<top>, local_dir/.cache).
        # On success the archive is already moved to dest; on any failure this
        # stops the transient hf cache copy from leaking across the fallback.
        _cleanup_hf_local_dir(local_dir, archive_file)


def download_with_resume(
    url: str,
    dest: str,
    expected_size: int,
    progress_start: int,
    progress_end: int,
) -> None:
    """Download a file with resume support via Range headers.

    Uses .partial and .meta sidecar files for crash recovery.
    """
    partial_path = dest + ".partial"
    meta_path = dest + ".meta"

    # Check for existing partial download
    bytes_downloaded = 0
    if os.path.exists(partial_path) and os.path.exists(meta_path):
        try:
            with open(meta_path, "r") as f:
                meta = json.load(f)
            bytes_downloaded = meta.get("bytesDownloaded", 0)
            if bytes_downloaded > 0:
                actual_size = os.path.getsize(partial_path)
                if actual_size != bytes_downloaded:
                    bytes_downloaded = 0  # Mismatch, restart
        except (json.JSONDecodeError, OSError):
            bytes_downloaded = 0

    if bytes_downloaded == 0 and os.path.exists(partial_path):
        os.unlink(partial_path)

    def _cleanup_partial() -> None:
        for p in (partial_path, meta_path):
            if os.path.exists(p):
                try:
                    os.unlink(p)
                except OSError:
                    pass

    max_retries = 5
    for attempt in range(max_retries):
        try:
            headers = {"User-Agent": "snapotter-installer/2.0"}
            if bytes_downloaded > 0:
                headers["Range"] = f"bytes={bytes_downloaded}-"
                emit_progress(
                    progress_start,
                    f"Resuming download from {bytes_downloaded / (1024**3):.1f} GB...",
                )

            req = urllib.request.Request(url, headers=headers)
            with urllib.request.urlopen(req, timeout=300) as resp:
                status = getattr(resp, "status", None) or resp.getcode()
                # If we asked to resume (sent a Range) but the server sent the
                # whole file back (200 instead of 206 Partial Content -- a proxy
                # or CDN that ignores Range), restart from byte 0. Appending a
                # full body onto the existing partial would corrupt the archive
                # and fail the checksum on every retry.
                if bytes_downloaded > 0 and status != 206:
                    bytes_downloaded = 0
                mode = "ab" if bytes_downloaded > 0 else "wb"
                next_meta_at = bytes_downloaded + DOWNLOAD_META_BYTES
                with open(partial_path, mode) as f:
                    while True:
                        chunk = resp.read(DOWNLOAD_CHUNK_BYTES)
                        if not chunk:
                            break
                        f.write(chunk)
                        bytes_downloaded += len(chunk)

                        # Update progress
                        if expected_size > 0:
                            pct = bytes_downloaded / expected_size
                            progress = int(
                                progress_start + pct * (progress_end - progress_start)
                            )
                            progress = min(progress, progress_end)
                            stage = f"Downloading... {bytes_downloaded / (1024**3):.1f} GB"
                            emit_progress(progress, stage)

                        # Write meta periodically so a crash can resume.
                        if bytes_downloaded >= next_meta_at:
                            with open(meta_path, "w") as mf:
                                json.dump({"bytesDownloaded": bytes_downloaded}, mf)
                            next_meta_at = bytes_downloaded + DOWNLOAD_META_BYTES

            # Guard against a truncated body or an error page served as the
            # archive: the completed size must match what the manifest expects.
            # A mismatch is retryable (transient truncation / a stale CDN edge).
            if expected_size > 0 and bytes_downloaded != expected_size:
                raise RuntimeError(
                    f"incomplete download: got {bytes_downloaded} bytes, "
                    f"expected {expected_size} (truncated response or error page)"
                )

            # Download complete
            os.rename(partial_path, dest)
            if os.path.exists(meta_path):
                os.unlink(meta_path)
            return

        except urllib.error.HTTPError as e:
            # HTTPError subclasses OSError, so it MUST be caught before the
            # OSError clause below. 4xx (except 408 Timeout / 429 Too Many
            # Requests) won't fix on retry -- a wrong URL, a private repo, or a
            # removed archive -- so fail fast with the manual-download hint.
            if 400 <= e.code < 500 and e.code not in (408, 429):
                _cleanup_partial()
                raise RuntimeError(
                    f"Download failed with HTTP {e.code} ({e.reason}). The archive "
                    f"URL may be wrong or access-restricted."
                )
            _retry_or_raise(e, attempt, max_retries, bytes_downloaded, meta_path,
                            progress_start, _cleanup_partial)
        except OSError as e:
            # Disk full is not transient: retrying can't create space. Fail fast
            # with an actionable message instead of burning the backoff budget.
            # (URLError/connection errors also land here; their errno is None, so
            # they fall through to the retry path.)
            if getattr(e, "errno", None) == errno.ENOSPC:
                _cleanup_partial()
                raise RuntimeError(
                    "Ran out of disk space while downloading the bundle. "
                    "Free up space and retry."
                )
            _retry_or_raise(e, attempt, max_retries, bytes_downloaded, meta_path,
                            progress_start, _cleanup_partial)
        except Exception as e:
            _retry_or_raise(e, attempt, max_retries, bytes_downloaded, meta_path,
                            progress_start, _cleanup_partial)


def _retry_or_raise(err, attempt, max_retries, bytes_downloaded, meta_path,
                    progress_start, cleanup) -> None:
    """Shared transient-failure handler for download_with_resume: persist resume
    metadata and back off, or clean up and raise on the final attempt."""
    try:
        with open(meta_path, "w") as mf:
            json.dump({"bytesDownloaded": bytes_downloaded}, mf)
    except OSError:
        pass

    if attempt < max_retries - 1:
        delay = min(60, 5 * (2 ** attempt))
        emit_progress(
            progress_start,
            f"Download failed (attempt {attempt + 1}/{max_retries}), "
            f"retrying in {delay}s: {err}",
        )
        time.sleep(delay)
    else:
        cleanup()
        raise RuntimeError(f"Failed to download after {max_retries} attempts: {err}")


# -- Safe tar extraction --

def safe_extract(tar_path: str, staging_dir: str) -> None:
    """Extract a tar.gz with security guards."""
    os.makedirs(staging_dir, exist_ok=True)
    with tarfile.open(tar_path, "r:gz") as tf:
        for member in tf.getmembers():
            # Block symlinks, hardlinks, devices
            if not member.isfile() and not member.isdir():
                raise RuntimeError(f"Blocked unsafe tar entry type: {member.name}")
            # Block absolute paths and traversal
            if member.name.startswith("/") or ".." in member.name.split("/"):
                raise RuntimeError(f"Blocked unsafe tar path: {member.name}")
        # The filter= kwarg was added in Python 3.12; the manual guards above
        # already block unsafe entries on older interpreters (e.g. 3.11).
        if sys.version_info >= (3, 12):
            tf.extractall(staging_dir, filter="data")
        else:
            tf.extractall(staging_dir)


# -- File move --

def move_tree(src: str, dst: str) -> None:
    """Merge src into dst, replacing entries crash-atomically where possible.

    This writes into the SHARED /data/ai/venv site-packages, so a crash mid-move
    must never leave a package in a half-replaced state (that tears the venv and
    breaks every other AI tool). For a file replacing a file, os.replace swaps in
    place with NO delete-then-write window, so an interruption leaves either the
    old or the new file intact, never a missing one. Cross-filesystem copies go
    through a temp sibling then an atomic rename for the same reason. Renames
    (vs copytree) also avoid transiently doubling the payload on disk."""
    if not os.path.isdir(src):
        return
    os.makedirs(dst, exist_ok=True)
    for name in os.listdir(src):
        s = os.path.join(src, name)
        d = os.path.join(dst, name)
        if os.path.isdir(s) and os.path.isdir(d):
            # Both dirs exist: merge recursively rather than replace.
            move_tree(s, d)
            continue
        try:
            # A type mismatch (dir<->file) can't be atomically swapped by rename,
            # so clear the destination first. A file-over-file or new entry needs
            # no pre-delete: os.replace is atomic and leaves no torn window.
            if os.path.exists(d) and os.path.isdir(d) != os.path.isdir(s):
                if os.path.isdir(d):
                    shutil.rmtree(d)
                else:
                    os.remove(d)
            os.replace(s, d)
            continue
        except OSError as e:
            if getattr(e, "errno", None) != errno.EXDEV:
                raise
        # Cross-filesystem: rename isn't allowed. Copy to a temp sibling and then
        # atomically replace, so a mid-copy ENOSPC never leaves a truncated file
        # where a working one used to be.
        if os.path.isdir(s):
            if os.path.exists(d):
                shutil.rmtree(d)
            shutil.copytree(s, d)
        else:
            tmp = d + ".part"
            shutil.copy2(s, tmp)
            os.replace(tmp, d)
    # Remove whatever remains of src (emptied by renames, or copied originals).
    shutil.rmtree(src, ignore_errors=True)


# -- ONNX Runtime flavor reconciliation (GPU wins) --

def _onnx_dist_infos(sp_dir: str) -> tuple:
    """Return (cpu, gpu) lists of onnxruntime dist-info directory names.

    Wheel metadata dirs are `<normalized-name>-<version>.dist-info`, and name
    normalization turns every other onnxruntime distribution into
    `onnxruntime_<suffix>` (onnxruntime-gpu -> onnxruntime_gpu), so a directory
    starting with exactly "onnxruntime-" can only be the CPU build.
    """
    if not os.path.isdir(sp_dir):
        return ([], [])
    names = os.listdir(sp_dir)
    cpu = sorted(n for n in names if n.startswith("onnxruntime-") and n.endswith(".dist-info"))
    gpu = sorted(n for n in names if n.startswith("onnxruntime_gpu-") and n.endswith(".dist-info"))
    return (cpu, gpu)


def reconcile_onnxruntime(staging_sp: str, site_packages_dir: str) -> None:
    """Never let a bundle downgrade the venv's ONNX Runtime from GPU to CPU.

    `onnxruntime` (CPU) and `onnxruntime-gpu` unpack into the SAME package
    directory (`onnxruntime/`), so whichever bundle lands last wins
    file-by-file. A CPU build arriving after the GPU build (e.g. transcription's
    faster-whisper dependency, after background-removal put onnxruntime-gpu in
    place) silently strips CUDAExecutionProvider while the surviving
    onnxruntime_gpu dist-info keeps claiming it is installed (#490). The GPU
    build ships CPUExecutionProvider too, so it satisfies every consumer of the
    CPU build; the reverse is not true. Hence: GPU wins, in both install orders.
    """
    staging_cpu, staging_gpu = _onnx_dist_infos(staging_sp)
    venv_cpu, venv_gpu = _onnx_dist_infos(site_packages_dir)

    if staging_cpu and (venv_gpu or staging_gpu):
        # Incoming CPU flavor while a GPU build exists: drop the CPU metadata,
        # and when the staged package files ARE the CPU build (no GPU flavor in
        # this same bundle), drop them too so they cannot clobber the venv's
        # GPU libraries.
        for name in staging_cpu:
            shutil.rmtree(os.path.join(staging_sp, name), ignore_errors=True)
        if not staging_gpu:
            pkg_dir = os.path.join(staging_sp, "onnxruntime")
            if os.path.isdir(pkg_dir):
                shutil.rmtree(pkg_dir, ignore_errors=True)
        sys.stderr.write(
            "[install] kept the GPU build of onnxruntime: dropped this bundle's CPU "
            "onnxruntime so it cannot disable CUDA for other AI tools (#490)\n"
        )
        sys.stderr.flush()

    if staging_gpu and venv_cpu:
        # Incoming GPU flavor over a CPU install: the GPU files win the merge;
        # clear the venv's stale CPU metadata so pip reflects reality. This also
        # makes reinstalling any GPU bundle repair a venv clobbered before this
        # guard existed.
        for name in venv_cpu:
            shutil.rmtree(os.path.join(site_packages_dir, name), ignore_errors=True)
        sys.stderr.write(
            "[install] replacing the CPU build of onnxruntime with the GPU build; "
            "removed its stale metadata\n"
        )
        sys.stderr.flush()


# -- Distribution reconciliation (one version per distribution) --

DIST_INFO_SUFFIX = ".dist-info"


def canonical_dist_name(name: str) -> str:
    """Normalize a distribution name the way PEP 503 does, so the same project
    compares equal however its wheel spelled it (`hf-xet`, `hf_xet`, `HF.Xet`)."""
    return re.sub(r"[-_.]+", "-", name).lower()


def list_distributions(sp_dir: str) -> dict:
    """Map canonical distribution name -> [(version, dist-info dir name), ...].

    A healthy site-packages holds exactly one entry per name. A venv merged by
    an installer that never uninstalls can hold several, which is the corruption
    this section exists to prevent: `move_tree` copies file over file, so the
    two versions overlay and CPython can end up loading one version's compiled
    extension underneath the other's Python modules.
    """
    found = {}
    if not os.path.isdir(sp_dir):
        return found
    for entry in os.listdir(sp_dir):
        if not entry.endswith(DIST_INFO_SUFFIX):
            continue
        name, sep, version = entry[: -len(DIST_INFO_SUFFIX)].rpartition("-")
        if not sep or not name:
            continue
        found.setdefault(canonical_dist_name(name), []).append((version, entry))
    return found


def distribution_files(sp_dir: str, dist_info: str) -> list:
    """Relative paths a distribution owns, read from its RECORD.

    Entries that point outside site-packages (console scripts are recorded as
    `../../bin/<name>`) are skipped: they cannot shadow an import, and deleting
    outside the directory we were handed is not this function's business.
    """
    record = os.path.join(sp_dir, dist_info, "RECORD")
    if not os.path.exists(record):
        return []
    paths = []
    with open(record, newline="", errors="replace") as f:
        for row in csv.reader(f):
            if not row:
                continue
            rel = row[0].strip().replace("\\", "/")
            if not rel or rel.startswith("/") or ".." in rel.split("/"):
                continue
            paths.append(rel)
    return paths


def paths_claimed_by_others(sp_dir: str, exclude: set) -> set:
    """Every file some OTHER installed distribution says it owns.

    Distributions are supposed to own disjoint files, and almost all of them do.
    The exceptions are the ones that package the same import under different
    project names: `opencv-python`, `opencv-python-headless` and
    `opencv-contrib-python` all write `cv2/`, exactly as `onnxruntime` and
    `onnxruntime-gpu` both write `onnxruntime/`. Uninstalling one of those by its
    RECORD would take the shared import away from the siblings that are still
    installed, and the merge only puts back the files of the distribution it is
    placing, so `cv2` would simply vanish.

    Shared files are therefore left where they are and let the merge overwrite
    them, which is what pip does with these packages too. The version metadata
    still ends up single-valued, which is the invariant that matters.
    """
    claimed = set()
    if not os.path.isdir(sp_dir):
        return claimed
    for entry in os.listdir(sp_dir):
        if entry.endswith(DIST_INFO_SUFFIX) and entry not in exclude:
            claimed.update(distribution_files(sp_dir, entry))
    return claimed


def _is_inside(root: str, path: str) -> bool:
    return not os.path.relpath(os.path.abspath(path), os.path.abspath(root)).startswith("..")


def _relocate(src_root: str, dst_root: str, rel_path: str) -> bool:
    """Move one relative entry between two trees, keeping its relative position."""
    src = os.path.join(src_root, rel_path)
    if not os.path.exists(src) and not os.path.islink(src):
        return False
    dst = os.path.join(dst_root, rel_path)
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    try:
        os.replace(src, dst)
    except OSError as e:
        if getattr(e, "errno", None) != errno.EXDEV:
            raise
        shutil.move(src, dst)
    return True


def _prune_empty_parents(root: str, rel_path: str) -> None:
    """Drop directories emptied by a removal, stopping at root and at the first
    directory that still holds something (another distribution's files)."""
    directory = os.path.dirname(os.path.join(root, rel_path))
    while _is_inside(root, directory) and os.path.abspath(directory) != os.path.abspath(root):
        try:
            os.rmdir(directory)
        except OSError:
            return
        directory = os.path.dirname(directory)


def _discard_stale_bytecode(root: str, rel_path: str) -> None:
    """Remove the cached bytecode of a source file that just left the tree.

    A `.pyc` whose `.py` is gone is not importable, but it does keep the package
    directory non-empty, which would stop `_prune_empty_parents` from clearing
    the way for the incoming version.
    """
    if not rel_path.endswith(".py"):
        return
    directory, base = os.path.split(os.path.join(root, rel_path))
    cache = os.path.join(directory, "__pycache__")
    for stale in glob.glob(os.path.join(cache, base[:-3] + ".*.pyc")):
        try:
            os.unlink(stale)
        except OSError:
            pass
    try:
        os.rmdir(cache)
    except OSError:
        pass


def plan_reconciliation(staging_sp: str, sp_dir: str) -> list:
    """Work out, per distribution the bundle is about to place, what it displaces.

    One entry per staged distribution that will actually change the venv:

        {"name", "version", "dist_info", "files", "superseded": [(version, dir)]}

    `superseded` holds every copy already installed under a different version.
    A distribution already present at the staged version is left out entirely:
    placing it is a no-op, so there is nothing to remove and nothing to undo.
    """
    staged = list_distributions(staging_sp)
    installed = list_distributions(sp_dir)
    plan = []
    for name in sorted(staged):
        version, dist_info = sorted(staged[name])[-1]
        existing = installed.get(name, [])
        superseded = [copy for copy in existing if copy[0] != version]
        if existing and not superseded:
            continue
        # RECORD describes the wheel, not the archive. A bundle's copy of
        # setuptools lists `_distutils_hack/` and `distutils-precedence.pth`
        # while carrying neither, so trusting RECORD alone would clear ground
        # that nothing is going to cover. Keep only what is really staged.
        files = [
            rel
            for rel in distribution_files(staging_sp, dist_info)
            if os.path.exists(os.path.join(staging_sp, rel))
        ]
        plan.append(
            {
                "name": name,
                "version": version,
                "dist_info": dist_info,
                "files": files,
                "replaces": {rel.split("/")[0] for rel in files},
                "superseded": superseded,
            }
        )
    return plan


def supersede_distributions(sp_dir: str, plan: list, quarantine_dir: str) -> int:
    """Uninstall every superseded version, holding its files for a rollback.

    Files move to `quarantine_dir/<name>/<version>/` rather than being deleted
    outright, so a failed verification can put the venv back exactly as it was
    instead of leaving it half-way between two bundles.
    """
    removed = 0
    superseded_dist_infos = {info for item in plan for _v, info in item["superseded"]}
    shared = paths_claimed_by_others(sp_dir, superseded_dist_infos)
    for item in plan:
        for version, dist_info in item["superseded"]:
            destination = os.path.join(quarantine_dir, item["name"], version)
            # Only clear ground the incoming copy is going to cover. A bundle
            # archive is not always a complete wheel: upscale-enhance carries
            # setuptools 74.1.3 as `setuptools/` and its dist-info, without the
            # `_distutils_hack/` and `distutils-precedence.pth` that the version
            # it replaces owns. Removing everything the old RECORD listed took
            # the distutils shim with it and basicsr, realesrgan and gfpgan all
            # stopped importing, because nothing was going to write those files
            # back. Anything the incoming copy does not carry stays put.
            replaces = item["replaces"] | {dist_info}
            for rel in distribution_files(sp_dir, dist_info):
                if rel in shared or rel.split("/")[0] not in replaces:
                    continue
                if _relocate(sp_dir, destination, rel):
                    _discard_stale_bytecode(sp_dir, rel)
                    _prune_empty_parents(sp_dir, rel)
            # RECORD does not always list every file pip wrote into the metadata
            # directory (INSTALLER and REQUESTED are frequently absent), so sweep
            # whatever is left before dropping the directory itself.
            leftover = os.path.join(sp_dir, dist_info)
            if os.path.isdir(leftover):
                for root, _dirs, files in os.walk(leftover):
                    for name in files:
                        rel = os.path.relpath(os.path.join(root, name), sp_dir)
                        _relocate(sp_dir, destination, rel)
                shutil.rmtree(leftover, ignore_errors=True)
            removed += 1
            sys.stderr.write(
                f"[install] replacing {item['name']} {version} with {item['version']}: "
                f"removed the superseded version so only one stays in the venv\n"
            )
    if removed:
        sys.stderr.flush()
    return removed


def restore_superseded(sp_dir: str, quarantine_dir: str) -> None:
    """Put every quarantined version back where it came from."""
    if not os.path.isdir(quarantine_dir):
        return
    for name in os.listdir(quarantine_dir):
        holder = os.path.join(quarantine_dir, name)
        if not os.path.isdir(holder):
            continue
        for version in os.listdir(holder):
            version_root = os.path.join(holder, version)
            if not os.path.isdir(version_root):
                continue
            for root, _dirs, files in os.walk(version_root):
                for entry in files:
                    rel = os.path.relpath(os.path.join(root, entry), version_root)
                    _relocate(version_root, sp_dir, rel)
    shutil.rmtree(quarantine_dir, ignore_errors=True)


def discard_placed_distributions(sp_dir: str, plan: list) -> None:
    """Remove the files this bundle placed, for the distributions it changed.

    Only distributions the plan actually touched are removed. One already
    present at the staged version never entered the plan, so a rollback cannot
    delete a version that was there before this install and is still correct.
    """
    placed_dist_infos = {item["dist_info"] for item in plan}
    shared = paths_claimed_by_others(sp_dir, placed_dist_infos)
    for item in plan:
        for rel in item["files"]:
            if rel in shared:
                continue
            target = os.path.join(sp_dir, rel)
            try:
                if os.path.isdir(target) and not os.path.islink(target):
                    shutil.rmtree(target, ignore_errors=True)
                else:
                    os.unlink(target)
            except OSError:
                pass
            _prune_empty_parents(sp_dir, rel)


def rollback_reconciliation(sp_dir: str, plan: list, quarantine_dir: str) -> None:
    """Undo a merge: drop what this bundle placed, restore what it displaced."""
    discard_placed_distributions(sp_dir, plan)
    restore_superseded(sp_dir, quarantine_dir)


def mark_venv_writing(marker_path: str, bundle_id: str) -> None:
    """Breadcrumb a destructive write to the shared venv.

    If the process dies while the marker exists, recoverInterruptedInstalls
    reseeds the venv from the image base on next boot. That is the blunt
    fallback for a tear this installer could not undo itself.
    """
    with open(marker_path, "w") as f:
        json.dump({"bundleId": bundle_id, "startedAt": datetime.now(timezone.utc).isoformat()}, f)


def clear_venv_writing(marker_path: str) -> None:
    if os.path.exists(marker_path):
        os.unlink(marker_path)


def abandon_merge(sp_dir, plan, quarantine_dir, marker_path, bundle_id) -> None:
    """Put the venv back the way it was before this bundle was merged.

    A refused install must leave nothing behind: without this, a bundle that
    fails verification would keep the versions it displaced from the bundles
    that were already working, so one bad install would break tools that used
    to run. Runs under the crash breadcrumb, and deliberately leaves the
    breadcrumb in place if the rollback itself fails, so a venv this installer
    could not repair still gets reseeded on next boot.
    """
    if not sp_dir or (not plan and not os.path.isdir(quarantine_dir)):
        return
    try:
        mark_venv_writing(marker_path, bundle_id)
        rollback_reconciliation(sp_dir, plan, quarantine_dir)
        clear_venv_writing(marker_path)
        sys.stderr.write("[install] rolled the venv back to its state before this bundle\n")
    except OSError as e:
        sys.stderr.write(
            f"[install] could not roll the venv back ({e}); the AI environment will be "
            f"reseeded on next restart\n"
        )
    sys.stderr.flush()


# -- Fixups (NCCL wheel) --

def apply_fixups(staging_dir: str, venv_path: str) -> None:
    """Install any wheels from fixups/ directory (local only, no network)."""
    fixups_dir = os.path.join(staging_dir, "fixups")
    if not os.path.isdir(fixups_dir):
        return
    wheels = [f for f in os.listdir(fixups_dir) if f.endswith(".whl")]
    if not wheels:
        return
    python_path = os.path.join(venv_path, "bin", "python3")
    if not os.path.exists(python_path):
        return
    for wheel in wheels:
        pkg_name = wheel.split("-")[0]
        try:
            subprocess.run(
                [python_path, "-m", "pip", "install", "--no-index",
                 f"--find-links={fixups_dir}", pkg_name],
                capture_output=True, text=True, timeout=60,
            )
        except Exception:
            pass  # Non-fatal


# -- installed.json management --

def read_installed(ai_dir: str) -> dict:
    """Read the current installed.json, returning empty structure if missing."""
    path = os.path.join(ai_dir, "installed.json")
    if not os.path.exists(path):
        return {"bundles": {}}
    try:
        with open(path, "r") as f:
            return json.load(f)
    except (json.JSONDecodeError, OSError):
        return {"bundles": {}}


def write_installed_atomic(ai_dir: str, data: dict) -> None:
    """Write installed.json atomically (write .tmp then rename)."""
    path = os.path.join(ai_dir, "installed.json")
    tmp_path = path + ".tmp"
    with open(tmp_path, "w") as f:
        json.dump(data, f, indent=2)
        f.write("\n")
    os.rename(tmp_path, path)


# -- Main --

def main() -> None:
    """Run the install with runtime-download restrictions lifted.

    In strict offline mode (SNAPOTTER_ALLOW_MODEL_DOWNLOAD=0) the sidecar
    runs with HF_HUB_OFFLINE/TRANSFORMERS_OFFLINE=1; a bundle install is an
    explicitly user-initiated download, so those flags are lifted here
    regardless. The previous values are restored in the finally block because
    this script can run in-process inside the long-lived dispatcher, where
    os.environ changes would otherwise leak into every later request.
    """
    saved = {key: os.environ.get(key) for key in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")}
    os.environ["HF_HUB_OFFLINE"] = "0"
    os.environ["TRANSFORMERS_OFFLINE"] = "0"
    try:
        _install()
    finally:
        for key, value in saved.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value


def _install() -> None:
    if len(sys.argv) < 4:
        fail(
            f"Usage: {sys.argv[0]} <bundleId> <manifestPath> <modelsDir>\n"
            f"Got {len(sys.argv) - 1} argument(s)"
        )

    bundle_id = sys.argv[1]
    manifest_path = sys.argv[2]
    models_dir = sys.argv[3]
    ai_dir = os.path.dirname(models_dir)
    staging_base = os.path.join(ai_dir, "staging")
    venv_path = os.environ.get("PYTHON_VENV_PATH", os.path.join(ai_dir, "venv"))

    # -- Load manifest --
    emit_progress(0, "Reading manifest...")

    try:
        with open(manifest_path, "r") as f:
            manifest = json.load(f)
    except Exception as e:
        fail(f"Failed to read manifest: {e}")

    bundles = manifest.get("bundles", {})
    if bundle_id not in bundles:
        fail(f"Unknown bundle: {bundle_id}")

    bundle = bundles[bundle_id]
    archives = bundle.get("archives")
    if not archives:
        fail(f"Bundle '{bundle_id}' has no archives in manifest (v2 required)")

    # -- Detect architecture --
    arch = detect_arch()
    archive = archives.get(arch)
    if not archive:
        fail(f"No archive for architecture '{arch}' in bundle '{bundle_id}'")

    archive_file = archive["file"]
    expected_sha256 = archive["sha256"]
    compressed_size = archive.get("compressedSize", 0)
    extracted_size = archive.get("extractedSize", 0)

    # -- Check for local file override (testing / offline) --
    local_path = os.environ.get("SNAPOTTER_BUNDLE_LOCAL_PATH")

    if local_path:
        # Local mode: use the file directly, verify checksum
        emit_progress(5, "Using local bundle archive...")
        tar_path = local_path

        if not os.path.exists(tar_path):
            fail(f"Local bundle file not found: {tar_path}")

        # Verify checksum
        emit_progress(10, "Verifying checksum...")
        if not verify_sha256(tar_path, expected_sha256):
            fail(
                f"SHA256 checksum mismatch for local file.\n"
                f"Expected: {expected_sha256}\n"
                f"This usually means the manifest and archive are out of sync."
            )
    else:
        # Remote mode: download from HuggingFace
        bundle_repo = manifest.get("bundleRepo", "deepsafe/feature-bundles")
        url = f"https://huggingface.co/{bundle_repo}/resolve/main/{archive_file}"

        # Disk space check (early sanity bail before a multi-GB download).
        # estimate_extracted covers the extractedSize:0 case so the budget can't
        # collapse to just the compressed size; the accurate guard is the
        # real-on-disk re-check just before the destructive venv write below.
        needed = compressed_size + estimate_extracted(compressed_size, extracted_size)
        needed += 500 * 1024 * 1024  # 500 MB buffer
        if needed > 0:
            check_disk_space(ai_dir, needed)

        # Download
        os.makedirs(staging_base, exist_ok=True)
        tar_path = os.path.join(staging_base, f"{bundle_id}-{arch}.tar.gz")

        emit_progress(2, f"Downloading {bundle.get('name', bundle_id)} bundle...")

        # Make sure the accelerated Xet client is importable in this venv, so a
        # drifted/upgraded venv doesn't silently fall back to slow urllib.
        ensure_hf_hub(venv_path)

        try:
            if not download_with_hf_hub(
                bundle_repo,
                archive_file,
                tar_path,
                compressed_size,
                2,
                85,
            ):
                download_with_resume(url, tar_path, compressed_size, 2, 85)
        except RuntimeError as e:
            fail(
                f"{e}\n\n"
                f"You can download the bundle manually from:\n"
                f"  {url}\n"
                f"Then upload it via Settings > AI Features > Offline Import."
            )

        # Verify checksum
        emit_progress(86, "Verifying integrity...")
        if not verify_sha256(tar_path, expected_sha256):
            # Delete and retry once from scratch
            os.unlink(tar_path)
            emit_progress(86, "Checksum mismatch, retrying download...")
            try:
                if not download_with_hf_hub(
                    bundle_repo,
                    archive_file,
                    tar_path,
                    compressed_size,
                    2,
                    85,
                    force_download=True,
                ):
                    download_with_resume(url, tar_path, compressed_size, 2, 85)
            except RuntimeError as e:
                fail(str(e))

            if not verify_sha256(tar_path, expected_sha256):
                os.unlink(tar_path)
                fail(
                    f"SHA256 checksum mismatch after re-download.\n"
                    f"Expected: {expected_sha256}\n"
                    f"The archive may be corrupted. Try again later."
                )

    # -- Extract to staging --
    staging_dir = os.path.join(ai_dir, f"staging-{bundle_id}")
    emit_progress(88, "Extracting packages and models...")

    try:
        if os.path.exists(staging_dir):
            shutil.rmtree(staging_dir)
        safe_extract(tar_path, staging_dir)
    except Exception as e:
        if os.path.exists(staging_dir):
            shutil.rmtree(staging_dir, ignore_errors=True)
        if isinstance(e, OSError) and getattr(e, "errno", None) == errno.ENOSPC:
            fail(
                "Ran out of disk space while extracting the bundle. "
                "Free up space and retry."
            )
        fail(f"Failed to extract archive: {e}")

    # -- Read bundle.json from tar --
    bundle_json_path = os.path.join(staging_dir, "bundle.json")
    if not os.path.exists(bundle_json_path):
        shutil.rmtree(staging_dir, ignore_errors=True)
        fail("Archive is missing bundle.json")

    try:
        with open(bundle_json_path, "r") as f:
            bundle_meta = json.load(f)
    except Exception as e:
        shutil.rmtree(staging_dir, ignore_errors=True)
        fail(f"Invalid bundle.json: {e}")

    version = bundle_meta.get("version", manifest.get("imageVersion", "unknown"))
    model_ids = bundle_meta.get("models", [])

    # -- Disk re-check before the first destructive venv write --
    # The upfront check ran before the download and used an estimate; now the
    # payload is really on disk, so measure it and verify there's room to place
    # it before we start writing into the venv. Running here (after the
    # local/remote branches merge) also covers the offline-import path, which
    # skipped the upfront check entirely. Each budget is checked against the
    # filesystem the bytes actually land on: when the venv lives on a
    # different filesystem than staging, the site-packages payload is COPIED
    # onto the venv's disk, so that disk (not ai_dir's) must hold it. Models
    # stay under ai_dir either way, moving by rename.
    disk_floor = 1024 ** 3  # 1 GB for fixups / installed.json / slack
    staging_sp = os.path.join(staging_dir, "site-packages")
    if not same_filesystem(staging_dir, venv_path):
        sp_bytes = dir_size(staging_sp) if os.path.isdir(staging_sp) else 0
        check_disk_space(venv_path, sp_bytes + disk_floor)
    check_disk_space(ai_dir, disk_floor)

    # -- Move site-packages --
    emit_progress(92, "Installing packages...")
    site_packages_dir = get_site_packages_dir(venv_path)
    venv_writing_marker = os.path.join(ai_dir, "venv.writing")
    # Superseded versions wait here until the install is verified. It sits inside
    # the venv so every move is a same-filesystem rename rather than a copy of
    # several GB, and outside site-packages so a quarantined dist-info cannot be
    # picked up by importlib.metadata while it waits.
    quarantine_dir = os.path.join(venv_path, ".superseded")
    reconciliation_plan = []

    try:
        if os.path.isdir(staging_sp) and site_packages_dir:
            # Breadcrumb the destructive shared-venv write. If the process is
            # killed mid-move (OOM/SIGKILL/power loss), move_tree can leave the
            # venv torn, which breaks OTHER installed tools. The marker survives
            # the crash; on next boot recoverInterruptedInstalls sees it and
            # reseeds the venv back to a known-good base. We clear it the instant
            # the site-packages move completes, since the venv is consistent
            # again then (a later models-move failure can't tear the venv).
            mark_venv_writing(venv_writing_marker, bundle_id)
            shutil.rmtree(quarantine_dir, ignore_errors=True)
            # Flavor first: `onnxruntime` and `onnxruntime-gpu` are two different
            # distributions writing the same package directory, so which one wins
            # is not a version question and the general pass below cannot decide
            # it. Running it first also means that when the CPU build is dropped
            # from staging, the general pass never sees it and so never uninstalls
            # the GPU build the venv is keeping (#490).
            reconcile_onnxruntime(staging_sp, site_packages_dir)
            reconciliation_plan = plan_reconciliation(staging_sp, site_packages_dir)
            supersede_distributions(site_packages_dir, reconciliation_plan, quarantine_dir)
            move_tree(staging_sp, site_packages_dir)
            clear_venv_writing(venv_writing_marker)

        # -- Move models --
        emit_progress(95, "Installing models...")
        staging_models = os.path.join(staging_dir, "models")
        if os.path.isdir(staging_models):
            os.makedirs(models_dir, exist_ok=True)
            move_tree(staging_models, models_dir)
    except OSError as e:
        shutil.rmtree(staging_dir, ignore_errors=True)
        abandon_merge(
            site_packages_dir, reconciliation_plan, quarantine_dir,
            venv_writing_marker, bundle_id,
        )
        if getattr(e, "errno", None) == errno.ENOSPC:
            fail("Ran out of disk space while installing the bundle. Free up space and retry.")
        fail(f"Failed to install bundle files: {e}")

    # -- Apply fixups --
    emit_progress(97, "Finalizing...")
    apply_fixups(staging_dir, venv_path)

    # -- Verify the bundle actually imports --
    # File-copy completion does NOT prove the bundle works: an incomplete
    # extraction or an ABI mismatch (e.g. a numpy/torch/protobuf skew) can leave
    # every file present yet the module unimportable, so the tool "installs" but
    # fails at first use. Import the bundle's key native libraries in the venv
    # now; if that fails, refuse to mark the bundle installed so the user gets a
    # clear retry instead of a silently broken tool.
    smoke_imports = bundle.get("smokeImports") or []
    if smoke_imports and os.environ.get("SNAPOTTER_SKIP_INSTALL_SMOKE") != "1":
        emit_progress(99, "Verifying installation...")
        venv_python = os.path.join(venv_path, "bin", "python3")
        if os.path.exists(venv_python):
            import_stmt = "\n".join(f"import {mod}" for mod in smoke_imports)
            try:
                proc = subprocess.run(
                    [venv_python, "-c", import_stmt],
                    capture_output=True, text=True, timeout=300,
                )
            except subprocess.TimeoutExpired:
                shutil.rmtree(staging_dir, ignore_errors=True)
                abandon_merge(
                    site_packages_dir, reconciliation_plan, quarantine_dir,
                    venv_writing_marker, bundle_id,
                )
                fail("Installation verification timed out. Please retry the install.")
            if proc.returncode != 0:
                shutil.rmtree(staging_dir, ignore_errors=True)
                abandon_merge(
                    site_packages_dir, reconciliation_plan, quarantine_dir,
                    venv_writing_marker, bundle_id,
                )
                tail = "\n".join((proc.stderr or "").strip().splitlines()[-6:])
                fail(
                    "Installation verification failed: the bundle installed but its "
                    "libraries could not be loaded, so the tool would not work.\n"
                    f"{tail}\n\n"
                    "This usually means an interrupted or corrupted install. Retry the "
                    "install; if it keeps failing, use Settings > AI Features > Reset AI "
                    "Environment, then reinstall."
                )

    # -- Write installed.json --
    emit_progress(98, "Recording installation...")
    installed = read_installed(ai_dir)
    installed["bundles"][bundle_id] = {
        "version": version,
        "installedAt": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "models": model_ids,
    }
    write_installed_atomic(ai_dir, installed)

    # -- Cleanup --
    # The install is recorded, so the versions this bundle displaced are never
    # coming back and their quarantine copy is just disk.
    shutil.rmtree(quarantine_dir, ignore_errors=True)
    if os.path.exists(staging_dir):
        shutil.rmtree(staging_dir, ignore_errors=True)
    # Clean up downloaded tar (but not if local override)
    if not local_path and os.path.exists(tar_path):
        os.unlink(tar_path)

    # -- Done --
    emit_progress(100, "Complete")

    result = {
        "success": True,
        "bundleId": bundle_id,
        "version": version,
        "models": model_ids,
    }
    print(json.dumps(result))


if __name__ == "__main__":
    main()