"""Fork fresh spares from one interpreter that already imported the kernel (Linux).

    albedo_zygote.py <socket-path>

Imports what a spare imports, freezes the heap so children share its pages,
prints `{"ready": bundle}`, then forks one child per connection to the unix
socket. The child serves albedo_spare's stdio protocol over that connection;
once the bundle has changed, the connection gets `{"stale": bundle}` instead.
The zygote stays single-threaded, so forking is safe; it exits on stdin EOF
(its owner closed the port) or once the bundle changes under it.
"""

from __future__ import annotations

import asyncio
import contextlib
import gc
import json
import os
import select
import signal
import socket
import sys
import traceback
from typing import Any

import albedo_bundle
import albedo_cells
import albedo_proc
import albedo_spare


def child(connection: socket.socket, bundle: str) -> None:
    albedo_proc.set_name("spare")
    os.dup2(connection.fileno(), 0)
    os.dup2(connection.fileno(), 1)
    connection.close()
    signal.signal(signal.SIGCHLD, signal.SIG_DFL)
    # Every child needs its own selector. Closing the inherited loop would
    # unregister its self-pipe from the epoll set the zygote shares, so only
    # this process's copies of its descriptors are dropped (macOS does not
    # inherit a kqueue at all).
    inherited: Any = albedo_cells.LOOP
    for resource in (inherited._ssock, inherited._csock, inherited._selector._selector):
        with contextlib.suppress(OSError, ValueError):
            resource.close()
    albedo_cells.LOOP = asyncio.new_event_loop()
    asyncio.set_event_loop(albedo_cells.LOOP)
    albedo_proc.reap_stopped_safely(albedo_cells.LOOP)
    albedo_spare.main(bundle)


def main(path: str) -> None:
    if os.getpgrp() != os.getpid():
        os.setsid()
    bundle = albedo_spare.preimport()
    listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    directory, name = os.path.split(path)
    os.chdir(directory)
    if os.path.exists(name):
        os.unlink(name)
    listener.bind(name)
    os.chmod(name, 0o600)
    listener.listen(64)
    # Children are reaped by the kernel; each resets this before running code.
    signal.signal(signal.SIGCHLD, signal.SIG_IGN)
    gc.collect()
    gc.freeze()
    print(json.dumps({"ready": bundle, "pid": os.getpid()}), flush=True)
    while True:
        readable, _, _ = select.select([listener, 0], [], [])
        if 0 in readable and not os.read(0, 4096):
            return
        if listener not in readable:
            continue
        connection, _ = listener.accept()
        if albedo_bundle.digest() != bundle:
            # Not a failure: the caller starts this spare directly.
            with contextlib.suppress(OSError):
                connection.sendall(json.dumps({"stale": bundle}).encode() + b"\n")
            connection.close()
            return
        if os.fork() == 0:
            listener.close()
            try:
                child(connection, bundle)
            except BaseException:
                traceback.print_exc()
                os._exit(1)
            os._exit(0)
        connection.close()


if __name__ == "__main__":
    main(sys.argv[1])