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

from __future__ import annotations

import os
import signal
import subprocess
import sys
import threading
from pathlib import Path
from typing import Callable, Optional

from fastapi import HTTPException

from hub.schemas.downloads import ActiveDownload, DownloadJobState
from hub.utils import download_manifest
from hub.utils import download_registry
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.hf_cache_state import EXIT_CANCELLED
from hub.utils.state_dir import RepoType


def backend_dir() -> Path:
    return Path(__file__).resolve().parent.parent.parent


def resolve_transport(use_xet: bool) -> str:
    transport = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
    unavailable_reason = download_registry.download_transport_unavailable_reason(transport)
    if unavailable_reason is not None:
        raise HTTPException(status_code = 400, detail = unavailable_reason)
    return transport


def spawn_worker(
    args: list[str],
    hf_token: Optional[str],
    *,
    use_xet: bool,
    protected_blob_hashes: Optional[frozenset[str]] = None,
) -> subprocess.Popen:
    """Spawn the download worker.

    XET and ``hf_transfer`` write chunks out of order, so their partials can't
    resume under a sequential writer; the HTTP path stays sequential so
    SIGKILL -> resume is byte-identical. ``protected_blob_hashes`` are blobs a
    concurrent same-repo peer is writing, excluded from the cache-prep purge so a
    shared ``.incomplete`` (e.g. bundled mmproj) is never deleted.
    """
    cwd = backend_dir()
    mode = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
    env = os.environ.copy()
    if protected_blob_hashes:
        env["UNSLOTH_PROTECTED_BLOB_HASHES"] = ",".join(sorted(protected_blob_hashes))
    else:
        env.pop("UNSLOTH_PROTECTED_BLOB_HASHES", None)
    env["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
    env["HF_HUB_DISABLE_TELEMETRY"] = "1"
    env["HF_HUB_DISABLE_XET"] = "0" if use_xet else "1"
    env["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "0" if hf_token else "1"
    # hf_transfer's parallel Range chunks can leave sparse partials even in
    # "http" mode; disable so the worker's writer is always sequential.
    env["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
    for token_key in (
        "HF_TOKEN",
        "HF_HUB_TOKEN",
        "HUGGING_FACE_HUB_TOKEN",
        "HUGGINGFACE_HUB_TOKEN",
        "HUGGINGFACEHUB_API_TOKEN",
    ):
        env.pop(token_key, None)
    if hf_token:
        env["HF_TOKEN"] = hf_token
    existing_path = env.get("PYTHONPATH", "")
    env["PYTHONPATH"] = f"{cwd}{os.pathsep}{existing_path}" if existing_path else str(cwd)
    return subprocess.Popen(
        [
            sys.executable,
            "-m",
            "hub.workers.hf_download",
            *args,
            "--parent-pid",
            str(os.getpid()),
            "--transport",
            mode,
        ],
        env = env,
        cwd = str(cwd),
        stdout = subprocess.DEVNULL,
        stderr = subprocess.PIPE,
        start_new_session = sys.platform != "win32",
    )


def drain_stderr_excerpt(stream, edge_bytes: int = 500) -> bytes:
    """Drain a worker's stderr to EOF, retaining the first and last bytes.

    Incremental reads keep the pipe from filling while bounding memory; long
    messages keep both ends since stderr prefixes often name the failing repo."""
    if stream is None:
        return b""
    edge_bytes = max(1, edge_bytes)
    max_bytes = edge_bytes * 2
    full = bytearray()
    head = bytearray()
    tail = bytearray()
    truncated = False
    for chunk in iter(lambda: stream.read(4096), b""):
        if not truncated:
            full.extend(chunk)
            if len(full) <= max_bytes:
                continue
            truncated = True
            head.extend(full[:edge_bytes])
            tail.extend(full[-edge_bytes:])
            full.clear()
            continue
        tail.extend(chunk)
        if len(tail) > edge_bytes:
            del tail[:-edge_bytes]
    if not truncated:
        return bytes(full)
    return bytes(head + b"\n...[stderr truncated]...\n" + tail)


def _cancellation_return_codes() -> frozenset[int]:
    """Returncodes for intentional cancellation only (SIGKILL/SIGTERM/SIGINT); crash signals stay errors, and ``getattr`` tolerates Windows where these signals are absent."""
    codes: set[int] = set()
    for name in ("SIGKILL", "SIGTERM", "SIGINT"):
        sig = getattr(signal, name, None)
        if sig is not None:
            codes.add(-int(sig))
    return frozenset(codes)


_CANCELLATION_RETURN_CODES = _cancellation_return_codes()


def _sigpipe_return_codes() -> frozenset[int]:
    sig = getattr(signal, "SIGPIPE", None)
    if sig is None:
        return frozenset()
    value = int(sig)
    return frozenset({-value, 128 + value})


_SIGPIPE_RETURN_CODES = _sigpipe_return_codes()


def classify_exit(rc: int, *, cancel_requested: bool = False) -> str:
    """Map a worker process exit code to a job state.

    - rc == 0: clean completion.
    - rc == EXIT_CANCELLED (130): the worker trapped a stop signal and exited
      cleanly with a resumable partial. In-app cancel uses untrappable SIGKILL
      and the OOM killer never produces 130, so 130 is always a resumable cancel.
    - rc killed by SIGKILL/SIGTERM/SIGINT: a cancel only when *we* asked for it.
      The OOM killer also sends SIGKILL, so an unrequested kill surfaces as error.
    - rc killed by SIGPIPE (or 128+SIGPIPE): parent pipe is gone; treated as
      cancelled.
    - any other non-zero rc (incl. crash signals): worker errored out.

    Windows has no POSIX signal exit encoding, so a user cancel can't be told from
    an error by code alone; there ``cancel_requested`` decides.
    """
    if rc == 0:
        return "complete"
    if rc == EXIT_CANCELLED:
        return "cancelled"
    if rc in _SIGPIPE_RETURN_CODES:
        return "cancelled"
    if rc in _CANCELLATION_RETURN_CODES:
        return "cancelled" if cancel_requested else "error"
    if cancel_requested and sys.platform == "win32":
        return "cancelled"
    return "error"


def finalize_worker_exit(
    registry: download_registry.DownloadRegistry,
    key: str,
    proc: subprocess.Popen,
    *,
    hf_token: Optional[str],
    label: str,
    log_prefix: str,
    logger,
    repo_type: Optional[RepoType] = None,
    repo_id: Optional[str] = None,
    transport: Optional[str] = None,
) -> None:
    """Block until *proc* exits, then record the job's terminal state in
    *registry*. Drains and scrubs stderr first, then classifies the exit code.
    A no-op when the process was already dropped (e.g. superseded).

    No stall watchdog: huggingface_hub already times out chunk reads and raises
    a resumable error on a dead connection, so the worker's exit code is the
    single source of truth."""
    stderr_data = drain_stderr_excerpt(proc.stderr)
    rc = proc.wait()
    cancel_requested = registry.cancel_requested(key)
    if not registry.drop_process(key, proc):
        return
    stderr_text = download_registry.scrub_secrets(
        (stderr_data or b"").decode("utf-8", "replace").strip(),
        hf_token = hf_token,
    )
    state = classify_exit(rc, cancel_requested = cancel_requested)
    if state == "complete":
        registry.set_job(key, "complete")
        if stderr_text:
            if download_manifest.MANIFEST_DEGRADED_MARKER in stderr_text:
                logger.warning(
                    f"{log_prefix} complete with degraded diagnostics for "
                    f"{label}: {stderr_text}"
                )
            else:
                logger.info(f"{log_prefix} worker diagnostics for {label}: {stderr_text}")
        logger.info(f"{log_prefix} complete: {label}")
        # Defensive cleanup: the canonical clear is at download-start; this
        # catches the rare case where that failed but the download succeeded.
        if repo_type and repo_id:
            try:
                download_manifest.clear_cancel_marker(
                    repo_type,
                    repo_id,
                    download_registry.variant_from_key(key),
                )
            except Exception as exc:
                logger.debug(f"clear_cancel_marker failed for {repo_id} (rc=0): {exc}")
    elif state == "cancelled":
        # Read metadata before the terminal set_job so a concurrent eviction
        # can't drop it; the job key is the fallback variant label.
        metadata = registry.get_job_metadata(key)
        registry.set_job(key, "cancelled")
        logger.info(f"{log_prefix} cancelled: {label} (rc={rc})")
        download_registry.persist_cancel_marker(
            repo_type,
            repo_id,
            metadata.variant
            if metadata is not None and metadata.variant
            else download_registry.variant_from_key(key),
            transport,
            logger = logger,
        )
    else:
        registry.set_job(
            key,
            "error",
            stderr_text or f"worker exited with code {rc}",
        )
        logger.error(
            f"{log_prefix} failed for {label} (rc={rc}): {stderr_text}",
        )


def kill_and_reap_process(
    proc: subprocess.Popen,
    *,
    label: str,
    logger,
    timeout: float = 10.0,
) -> None:
    try:
        proc.kill()
    except ProcessLookupError:
        pass
    except Exception as exc:
        logger.warning(f"Cancel SIGKILL for {label} failed: {exc}")
    try:
        proc.wait(timeout = timeout)
    except subprocess.TimeoutExpired:
        logger.warning(f"Cancelled worker for {label} did not exit after SIGKILL")
    except Exception:
        pass


def register_worker(
    registry: download_registry.DownloadRegistry,
    key: str,
    proc: subprocess.Popen,
    *,
    hf_token: Optional[str],
    label: str,
    log_prefix: str,
    logger,
    repo_type: RepoType,
    repo_id: str,
    transport: str,
    watch_name: str,
) -> bool:
    if not registry.register_process(key, proc):
        kill_and_reap_process(proc, label = label, logger = logger)
        return False

    worker_token = hf_token

    def _watch() -> None:
        finalize_worker_exit(
            registry,
            key,
            proc,
            hf_token = worker_token,
            label = label,
            log_prefix = log_prefix,
            logger = logger,
            repo_type = repo_type,
            repo_id = repo_id,
            transport = transport,
        )
        if registry.get_job(key).state in ("error", "cancelled"):
            download_registry.purge_empty_marker_dir(
                repo_type,
                repo_id,
                download_registry.variant_from_key(key),
            )
        hf_cache_scan.invalidate_hf_cache_scans()

    threading.Thread(target = _watch, name = watch_name, daemon = True).start()
    return True


def launch_worker(
    registry: download_registry.DownloadRegistry,
    key: str,
    *,
    spawn: Callable[[], subprocess.Popen],
    hf_token: Optional[str],
    label: str,
    log_prefix: str,
    logger,
    repo_type: RepoType,
    repo_id: str,
    transport: str,
    watch_name: str,
) -> str:
    try:
        proc = spawn()
    except Exception as e:
        scrubbed = download_registry.scrub_secrets(str(e), hf_token = hf_token)
        logger.error(
            f"Failed to spawn {log_prefix.lower()} worker for {label}: {scrubbed}",
            exc_info = True,
        )
        registry.set_job(key, "error", scrubbed)
        raise HTTPException(
            status_code = 500,
            detail = f"Failed to start {log_prefix.lower()}: {scrubbed}",
        ) from e
    register_worker(
        registry,
        key,
        proc,
        hf_token = hf_token,
        label = label,
        log_prefix = log_prefix,
        logger = logger,
        repo_type = repo_type,
        repo_id = repo_id,
        transport = transport,
        watch_name = watch_name,
    )
    return registry.get_job(key).state


def cancel_worker(
    registry: download_registry.DownloadRegistry,
    key: str,
    *,
    generation: Optional[int],
    label: str,
    logger,
) -> str:
    proc = registry.get_process(key)
    # No worker process yet: arm a pending cancel so register_process kills it on
    # arrival during the claim-to-register window.
    if proc is None:
        if registry.mark_pending_cancel(key, generation):
            return "cancelling"
        return registry.get_job(key).state
    # Worker already exited; let its watcher classify the real return code.
    # Arming a pending cancel here could mislabel a genuine failure as a cancel.
    if proc.poll() is not None:
        return registry.get_job(key).state

    if not registry.request_cancel(key, proc, generation):
        return registry.get_job(key).state
    # No eager marker: finalize_worker_exit writes it on a "cancelled" exit.
    # Persisting before the kill races a clean completion and strands a stale marker.
    try:
        proc.kill()
    except ProcessLookupError:
        pass
    except Exception as e:
        logger.warning(f"Cancel SIGKILL for {label} failed: {e}")

    return "cancelling"


def idle_status(
    registry: download_registry.DownloadRegistry,
    key: str,
    *,
    repo_type: RepoType,
    repo_id: Optional[str],
    variant: Optional[str],
) -> tuple[DownloadJobState, Optional[str], int]:
    state = registry.get_job(key)
    generation = registry.current_generation(key)
    if (
        state.state == "idle"
        and repo_id
        and download_manifest.has_cancel_marker(
            repo_type,
            repo_id,
            variant,
        )
    ):
        return ("cancelled", None, generation)
    return (state.state, state.error, generation)


def active_download_refs(
    registry: download_registry.DownloadRegistry, repo_id: Optional[str], *, with_variant: bool
) -> list[ActiveDownload]:
    downloads: list[ActiveDownload] = []
    for ref in registry.active_job_refs(repo_id):
        metadata = ref.metadata
        if with_variant:
            ref_repo_id = metadata.repo_id if metadata is not None else ref.key.split("::", 1)[0]
            if metadata is not None:
                variant = metadata.variant
            else:
                _repo, sep, raw_variant = ref.key.partition("::")
                variant = raw_variant if sep and raw_variant else None
        else:
            ref_repo_id = metadata.repo_id if metadata is not None else ref.key
            variant = None
        downloads.append(
            ActiveDownload(
                repo_id = ref_repo_id,
                variant = variant,
                transport = metadata.transport if metadata is not None else None,
                state = ref.state,
                generation = ref.generation,
            )
        )
    return downloads