ArchiveBox/archivebox/services/binary_service.py

from __future__ import annotations

import asyncio
import json
from collections.abc import Mapping
from pathlib import Path
from typing import Any

from asgiref.sync import sync_to_async
from django.utils import timezone

from abxpkg.binary_service import BinaryEvent, BinaryRequestEvent
from abxpkg.config import load_derived_cache
from abxbus import BaseEvent, EventBus
from abx_dl.services.base import BaseService


class ArchiveBoxBinaryService(BaseService):
    """Preserve ArchiveBox's legacy Binary Process rows around abxpkg requests."""

    LISTENS_TO = [BinaryRequestEvent, BinaryEvent]
    EMITS: list[type[BaseEvent]] = []

    def __init__(self, bus: EventBus):
        super().__init__(bus)
        self.process_ids_by_request_id: dict[str, str] = {}
        self._missing_finalize_tasks: set[asyncio.Task] = set()
        self.bus.on(BinaryRequestEvent, self.on_BinaryRequestEvent__project_process)
        self.bus.on(BinaryRequestEvent, self.on_BinaryRequestEvent__schedule_missing_finalize)
        self.bus.on(BinaryEvent, self.on_BinaryEvent__finalize_process)

    async def on_BinaryRequestEvent__project_process(self, request: BinaryRequestEvent) -> None:
        from archivebox.machine.models import Machine, Process, _canonical_binary_name

        machine = await sync_to_async(Machine.current, thread_sensitive=True)()
        binary_name = _canonical_binary_name(request.name)
        if not binary_name:
            return
        binary = await self._get_or_create_binary(machine, binary_name, request)
        started_at = timezone.now()
        output_dir = self._process_output_dir(binary, request)
        await sync_to_async(output_dir.mkdir, thread_sensitive=True)(parents=True, exist_ok=True)
        process = await Process.objects.acreate(
            machine=machine,
            iface=None,
            process_type=Process.TypeChoices.BINARY,
            worker_type="",
            pwd=str(output_dir),
            cmd=self._process_cmd(request),
            env={},
            timeout=int(request.event_timeout or request.install_timeout or 600),
            pid=None,
            url=None,
            started_at=started_at,
            ended_at=None,
            stdout="",
            stderr="",
            exit_code=None,
            status=Process.StatusChoices.RUNNING,
            retry_at=None,
            binary=binary,
        )
        self.process_ids_by_request_id[request.event_id] = str(process.id)

    async def on_BinaryEvent__finalize_process(self, event: BinaryEvent) -> None:
        from archivebox.machine.models import Binary, Process, _canonical_binary_name

        request = await self.bus.find(
            BinaryRequestEvent,
            past=True,
            future=False,
            where=lambda candidate: self.bus.event_is_child_of(event, candidate),
        )
        request = request if isinstance(request, BinaryRequestEvent) else None
        process_id = self.process_ids_by_request_id.pop(request.event_id, "") if request is not None else ""
        if not process_id:
            return
        process = await Process.objects.filter(id=process_id).select_related("binary").afirst()
        if process is None:
            return
        binary_name = _canonical_binary_name(event.name)
        binary = process.binary
        if binary is not None and binary_name:
            binary.abspath = event.abspath
            if event.version:
                binary.version = str(event.version)
            if event.sha256:
                binary.sha256 = str(event.sha256)
            binary.binproviders = event.binproviders or binary.binproviders
            binary.binprovider = event.binprovider or binary.binprovider
            binary.status = Binary.StatusChoices.INSTALLED
            binary.retry_at = None
            await binary.asave(
                update_fields=["abspath", "version", "sha256", "binproviders", "binprovider", "status", "retry_at", "modified_at"],
            )
        process.ended_at = timezone.now()
        process.stdout = json.dumps(self._binary_event_json(event, binary)) + "\n"
        process.stderr = ""
        process.exit_code = 0
        process.status = Process.StatusChoices.EXITED
        await process.asave(update_fields=["ended_at", "stdout", "stderr", "exit_code", "status", "modified_at"])
        if binary is not None:
            await sync_to_async(self._write_binary_index, thread_sensitive=True)(binary, process, Path(process.pwd))

    async def _get_or_create_binary(self, machine, binary_name: str, request: BinaryRequestEvent):
        from archivebox.machine.models import Binary

        binary_id = str(request.extra_context.get("binary_id") or "")
        if binary_id:
            binary = await Binary.objects.filter(id=binary_id).afirst()
            if binary is not None:
                return binary
        binary = await Binary.objects.filter(machine=machine, name=binary_name).order_by("-modified_at").afirst()
        if binary is not None:
            return binary
        return await Binary.objects.acreate(
            machine=machine,
            name=binary_name,
            binproviders=_binproviders_to_str(request.binproviders),
            overrides=_persisted_overrides_for_request(request),
            status=Binary.StatusChoices.QUEUED,
        )

    def _process_cmd(self, request: BinaryRequestEvent) -> list[str]:
        cmd = [
            "abxpkg",
            "install",
            f"--name={request.name}",
            f"--binproviders={_binproviders_to_str(request.binproviders)}",
        ]
        if request.overrides:
            cmd.append(f"--overrides={json.dumps(request.overrides, sort_keys=True)}")
        return cmd

    def _binary_event_json(self, event: BinaryEvent, binary) -> dict[str, Any]:
        if binary is not None:
            data = binary.to_json()
        else:
            data = {"type": "Binary", "name": event.name}
        data.update(
            {
                "type": "Binary",
                "name": event.name,
                "binproviders": event.binproviders,
                "binprovider": event.binprovider,
                "abspath": event.abspath,
                "version": str(event.version or ""),
                "sha256": event.sha256 or "",
                "status": "installed",
            },
        )
        return data

    async def _finalize_missing_process(self, request: BinaryRequestEvent) -> None:
        from archivebox.machine.models import Process

        process_id = self.process_ids_by_request_id.pop(request.event_id, "")
        if not process_id:
            return
        process = await Process.objects.filter(id=process_id).afirst()
        if process is None or process.status == Process.StatusChoices.EXITED:
            return
        process.ended_at = timezone.now()
        process.stderr = f"Binary request did not resolve: {request.name}"
        process.exit_code = 1
        process.status = Process.StatusChoices.EXITED
        await process.asave(update_fields=["ended_at", "stderr", "exit_code", "status", "modified_at"])

    async def _finalize_request_when_done(self, request: BinaryRequestEvent) -> None:
        try:
            await request.wait(timeout=request.event_timeout)
        except TimeoutError:
            await self._finalize_missing_process(request)
            return
        binary_event = await self.bus.find(
            BinaryEvent,
            child_of=request,
            past=True,
            future=False,
            name=request.name,
            where=lambda candidate: bool(candidate.abspath),
        )
        if not isinstance(binary_event, BinaryEvent):
            await self._finalize_missing_process(request)

    def _schedule_missing_finalize(self, request: BinaryRequestEvent) -> None:
        task = asyncio.create_task(self._finalize_request_when_done(request))
        self._missing_finalize_tasks.add(task)
        task.add_done_callback(lambda done: self._missing_finalize_tasks.discard(done) or (None if done.cancelled() else done.exception()))

    async def flush_missing_finalizers(self) -> None:
        if self._missing_finalize_tasks:
            await asyncio.gather(*tuple(self._missing_finalize_tasks), return_exceptions=False)

    async def on_BinaryRequestEvent__schedule_missing_finalize(self, request: BinaryRequestEvent) -> None:
        self._schedule_missing_finalize(request)

    def _process_output_dir(self, binary, request: BinaryRequestEvent) -> Path:
        raw_output_dir = str(request.extra_context.get("output_dir") or "").strip()
        if raw_output_dir:
            output_dir = Path(raw_output_dir).expanduser()
            if output_dir.name == str(binary.id):
                return output_dir.parent
            return output_dir
        return binary.output_dir.parent

    def _write_binary_index(self, binary, process, output_dir: Path) -> None:
        output_dir.mkdir(parents=True, exist_ok=True)
        index_path = output_dir / "index.jsonl"
        with index_path.open("w", encoding="utf-8") as f:
            f.write(json.dumps(binary.to_json()) + "\n")
            f.write(json.dumps(process.to_json()) + "\n")


def project_abxpkg_derived_cache_to_db(lib_dir: Path | str | None) -> None:
    """Mirror abxpkg subprocess-resolved binaries into ArchiveBox's DB.

    Hook shebangs resolve binaries through the abxpkg CLI in subprocesses, so
    those resolutions cannot emit in-process BinaryRequestEvent/BinaryEvent
    objects on the ArchiveBox runner bus. abxpkg's generic cross-process
    projection point is ``LIB_DIR/env/derived.env``; ArchiveBox imports those
    resolved records here after hook execution.
    """

    if lib_dir is None:
        return

    lib_path = Path(lib_dir).expanduser()
    derived_env_paths = sorted(lib_path.rglob("derived.env")) if lib_path.is_dir() else []
    if not derived_env_paths:
        return

    from archivebox.machine.models import Binary, Machine, Process, _canonical_binary_name

    machine = Machine.current()
    for derived_env_path in derived_env_paths:
        for record in load_derived_cache(derived_env_path).values():
            if not isinstance(record, Mapping):
                continue
            binary_name = _canonical_binary_name(str(record.get("bin_name") or ""))
            if not binary_name:
                continue
            abspath = str(record.get("abspath") or "").strip()
            if not abspath:
                continue
            binary_path = Path(abspath).expanduser().resolve(strict=False)
            if not binary_path.exists():
                continue

            version = str(record.get("loaded_version") or "")
            sha256 = str(record.get("loaded_sha256") or "")
            provider_name = str(record.get("provider_name") or "")
            resolved_provider_name = str(record.get("resolved_provider_name") or provider_name)
            installed_abspath = str(binary_path)

            binary, _created = Binary.objects.get_or_create(
                machine=machine,
                name=binary_name,
                defaults={
                    "status": Binary.StatusChoices.QUEUED,
                    "binproviders": provider_name or resolved_provider_name or "env",
                },
            )
            previous_projection = (
                binary.status,
                binary.abspath,
                binary.version,
                binary.sha256,
                binary.binprovider,
            )
            binary.abspath = installed_abspath
            binary.version = version
            binary.sha256 = sha256
            binary.binproviders = provider_name or resolved_provider_name or binary.binproviders or "env"
            binary.binprovider = resolved_provider_name or provider_name or binary.binprovider
            binary.status = Binary.StatusChoices.INSTALLED
            binary.retry_at = None
            binary.save(
                update_fields=[
                    "abspath",
                    "version",
                    "sha256",
                    "binproviders",
                    "binprovider",
                    "status",
                    "retry_at",
                    "modified_at",
                ],
            )

            current_projection = (
                binary.status,
                binary.abspath,
                binary.version,
                binary.sha256,
                binary.binprovider,
            )
            if current_projection == previous_projection:
                continue

            output_dir = binary.output_dir.parent
            output_dir.mkdir(parents=True, exist_ok=True)
            now = timezone.now()
            process = Process.objects.create(
                machine=machine,
                iface=None,
                process_type=Process.TypeChoices.BINARY,
                worker_type="",
                pwd=str(output_dir),
                cmd=[
                    "abxpkg",
                    "run",
                    "--script",
                    f"--name={binary_name}",
                    f"--binproviders={binary.binproviders}",
                ],
                env={},
                timeout=0,
                pid=None,
                url=None,
                started_at=now,
                ended_at=now,
                stdout=json.dumps(
                    {
                        "type": "Binary",
                        "name": binary_name,
                        "binproviders": binary.binproviders,
                        "binprovider": binary.binprovider,
                        "abspath": binary.abspath,
                        "version": binary.version,
                        "sha256": binary.sha256,
                        "status": "installed",
                    },
                )
                + "\n",
                stderr="",
                exit_code=0,
                status=Process.StatusChoices.EXITED,
                retry_at=None,
                binary=binary,
            )
            index_path = output_dir / "index.jsonl"
            with index_path.open("w", encoding="utf-8") as f:
                f.write(json.dumps(binary.to_json()) + "\n")
                f.write(json.dumps(process.to_json()) + "\n")


def _provider_names(binproviders: str | list[str] | None) -> list[str]:
    if isinstance(binproviders, str):
        raw_names = [part.strip() for part in binproviders.split(",")]
    elif binproviders:
        raw_names = [str(part).strip() for part in binproviders]
    else:
        raw_names = ["env"]
    names: list[str] = []
    for name in raw_names:
        if name and name not in names:
            names.append(name)
    return names or ["env"]


def _binproviders_to_str(binproviders: str | list[str] | None) -> str:
    return ",".join(_provider_names(binproviders))


def _persisted_overrides_for_request(request: BinaryRequestEvent | None) -> dict[str, Any]:
    if request is None:
        return {}
    return dict(request.overrides or {})