src/prepare.py

from __future__ import annotations

import json
import os
import shutil
import site
import stat
import subprocess
import tarfile
import tempfile
import zipfile
from collections import deque
from datetime import UTC, datetime
from pathlib import Path
from urllib.request import urlretrieve

from huggingface_hub import hf_hub_download

from src.config import settings
from src.errors import ApiError


_runtime_notes: list[str] = []
_events: deque[dict[str, object]] = deque(maxlen=200)


def log_event(event: str, **fields: object) -> None:
    record: dict[str, object] = {
        "ts": datetime.now(UTC).isoformat(timespec="seconds"),
        "event": event,
        **fields,
    }
    _events.append(record)
    print(json.dumps(record, ensure_ascii=False, default=str), flush=True)


def runtime_events(limit: int = 100) -> list[dict[str, object]]:
    limit = max(1, min(limit, 200))
    return list(_events)[-limit:]


def nvidia_library_paths() -> list[str]:
    paths: list[str] = []
    seen: set[str] = set()
    roots = [*site.getsitepackages(), site.getusersitepackages()]
    for root in roots:
        if not root:
            continue
        base = Path(root) / "nvidia"
        if not base.exists():
            continue
        for candidate in base.glob("*/lib"):
            if candidate.is_dir():
                value = str(candidate)
                if value not in seen:
                    seen.add(value)
                    paths.append(value)
    return paths


def ld_library_path_for(binary_path: Path | None = None) -> str:
    paths: list[str] = []
    if binary_path:
        paths.append(str(binary_path.resolve().parent))
    paths.extend(nvidia_library_paths())
    existing = os.getenv("LD_LIBRARY_PATH", "")
    if existing:
        paths.append(existing)
    return ":".join(paths)


def runtime_status() -> dict[str, object]:
    return {
        "bin": str(settings.llama_diffusion_bin),
        "bin_exists": settings.llama_diffusion_bin.exists(),
        "model_cache_dir": str(settings.model_cache_dir),
        "model_file": str(settings.model_cache_dir / settings.gguf_filename),
        "model_file_exists": (settings.model_cache_dir / settings.gguf_filename).exists(),
        "nvidia_library_paths": nvidia_library_paths(),
        "ld_library_path": ld_library_path_for(settings.llama_diffusion_bin),
        "notes": list(_runtime_notes[-20:]),
    }


def _mark_executable(path: Path) -> None:
    mode = path.stat().st_mode
    path.chmod(mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)


def _find_binary(root: Path) -> Path | None:
    for candidate in root.rglob("llama-diffusion-cli"):
        if candidate.is_file():
            return candidate
    return None


def _download_prebuilt_binary(url: str) -> None:
    settings.bin_dir.mkdir(parents=True, exist_ok=True)
    log_event("runner.download.start", url=url)
    with tempfile.TemporaryDirectory() as td:
        tempdir = Path(td)
        archive_path = tempdir / "llama-diffusion-download"
        urlretrieve(url, archive_path)

        extracted_dir = tempdir / "extract"
        extracted_dir.mkdir(parents=True, exist_ok=True)

        if tarfile.is_tarfile(archive_path):
            with tarfile.open(archive_path) as tf:
                tf.extractall(extracted_dir)
            binary = _find_binary(extracted_dir)
        elif zipfile.is_zipfile(archive_path):
            with zipfile.ZipFile(archive_path) as zf:
                zf.extractall(extracted_dir)
            binary = _find_binary(extracted_dir)
        else:
            binary = archive_path

        if not binary or not binary.exists():
            raise ApiError("runtime_error", "Could not find llama-diffusion-cli in downloaded artifact", 500)

        # Keep sibling shared libraries with the binary. This mirrors the qwen36
        # Space strategy: a CUDA runner package is not just the executable.
        for child in binary.parent.iterdir():
            target = settings.bin_dir / child.name
            if child.is_file() or child.is_symlink():
                shutil.copy2(child, target)
            elif child.is_dir() and child.name in {"lib", "lib64"}:
                if target.exists():
                    shutil.rmtree(target)
                shutil.copytree(child, target, symlinks=True)
        _mark_executable(settings.llama_diffusion_bin)
        _runtime_notes.append(f"Installed prebuilt binary from {url}")
        log_event("runner.download.finish", bin=str(settings.llama_diffusion_bin))


def build_llama_diffusion() -> None:
    script = Path(__file__).resolve().parent.parent / "scripts" / "build_llama_diffusion.sh"
    if not script.exists():
        raise ApiError("runtime_error", f"Missing build script: {script}", 500)

    log_event("runner.build.start", script=str(script))
    env = dict(os.environ)
    env.setdefault("LLAMA_SRC_DIR", str(settings.llama_src_dir))
    env.setdefault("LLAMA_BIN_DIR", str(settings.bin_dir))
    env.setdefault("LLAMA_DIFFUSION_BIN", str(settings.llama_diffusion_bin))
    env.setdefault("LLAMA_BUILD_CUDA", "1" if settings.llama_build_cuda else "0")
    env.setdefault("LLAMA_CMAKE_EXTRA_ARGS", settings.llama_cmake_extra_args)
    env["LD_LIBRARY_PATH"] = ld_library_path_for(settings.llama_diffusion_bin)

    proc = subprocess.run(
        ["bash", str(script)],
        text=True,
        capture_output=True,
        env=env,
        timeout=int(os.getenv("LLAMA_BUILD_TIMEOUT_SECONDS", "1800")),
    )
    if proc.returncode != 0:
        log_event("runner.build.failed", stderr_tail=proc.stderr[-1200:])
        raise ApiError(
            "runtime_error",
            "Failed to build llama-diffusion-cli. stderr tail: " + proc.stderr[-4000:],
            500,
        )
    _runtime_notes.append("Built llama-diffusion-cli from llama.cpp DiffusionGemma PR")
    log_event("runner.build.finish", bin=str(settings.llama_diffusion_bin))


def ensure_runner_binary() -> Path:
    if settings.llama_diffusion_bin.exists():
        _mark_executable(settings.llama_diffusion_bin)
        log_event("runner.cache_hit", bin=str(settings.llama_diffusion_bin))
        return settings.llama_diffusion_bin

    if settings.llama_diffusion_bin_url:
        _download_prebuilt_binary(settings.llama_diffusion_bin_url)
        return settings.llama_diffusion_bin

    if settings.build_llama_diffusion:
        build_llama_diffusion()
        if settings.llama_diffusion_bin.exists():
            _mark_executable(settings.llama_diffusion_bin)
            return settings.llama_diffusion_bin

    raise ApiError(
        "runtime_error",
        "llama-diffusion-cli not found. Set LLAMA_DIFFUSION_BIN, LLAMA_DIFFUSION_BIN_URL, or BUILD_LLAMA_DIFFUSION=1.",
        500,
    )


def ensure_model_file(repo_id: str, filename: str) -> str:
    settings.model_cache_dir.mkdir(parents=True, exist_ok=True)
    log_event("model.download.start", repo_id=repo_id, filename=filename)
    path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        local_dir=str(settings.model_cache_dir),
        token=os.getenv("HF_TOKEN") or None,
    )
    _runtime_notes.append(f"Model ready: {repo_id}/{filename}")
    log_event("model.download.finish", path=str(path))
    return path


def prepare_runtime_if_requested() -> None:
    if settings.prepare_runtime_on_startup:
        ensure_runner_binary()
    if settings.download_model_on_startup:
        ensure_model_file(settings.gguf_repo_id, settings.gguf_filename)