"""Copy a detached kernel's frames between this process's stdio and its socket.

    albedo_bridge.py start <run_dir> <modules-json>
    albedo_bridge.py attach <run_dir>
    albedo_bridge.py launch <run_dir> <modules-json>
    albedo_bridge.py --frame

`--frame` takes the same arguments from its first stdin frame instead,
`{"bridge": {"argv": [...], "cwd": dir}}`, and starts in `cwd` (exit status 4
when it is no directory): the remote form, so the daemon's ssh command line
never carries a path or the module list.

The daemon talks to this process as it once talked to the kernel itself, over
4-byte-length-framed stdio, and writes the attach frame first. `start` reads
that frame to learn the kernel's token, starts the albedo_kernel.py beside this
file in its own session (the token travels over the kernel's stdin, never
argv), and waits for its socket; `attach` connects to a kernel that is already running. Either way this
process announces the bundle it runs from, then copies bytes both ways until
one side closes. Killing it never touches the kernel.

`launch` only starts the kernel, its token read as one stdin line, and exits 0
once the socket answers: a local daemon connects to that socket itself, so no
bridge stays running beside a local kernel.

Exit status 3 means no kernel is there to attach to, so the daemon can tell a
kernel that is gone from a connection that merely dropped.
"""

from __future__ import annotations

import errno
import os
import socket
import subprocess
import sys
import threading
import time
from collections.abc import Callable

import albedo_bundle
import albedo_link

GONE = 3
NO_FOLDER = 4
START_TIMEOUT = 10.0  # seconds a starting kernel has to bind its socket


def connect(run_dir: str) -> socket.socket | None:
    """The kernel's socket, or None when nothing listens there."""
    connection = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        os.chdir(run_dir)
        connection.connect(albedo_link.SOCKET)
        return connection
    except OSError as error:
        connection.close()
        if error.errno in (errno.ENOENT, errno.ECONNREFUSED, errno.ENOTDIR):
            return None
        raise


KERNEL = os.path.join(os.path.dirname(os.path.abspath(__file__)), "albedo_kernel.py")


def start(run_dir: str, modules: str, token: str) -> socket.socket:
    os.makedirs(run_dir, mode=0o700, exist_ok=True)
    with open(os.path.join(run_dir, "kernel.log"), "ab") as log:
        kernel = subprocess.Popen(
            [sys.executable, "-u", KERNEL, "--run", run_dir, modules],
            stdin=subprocess.PIPE,
            stdout=subprocess.DEVNULL,
            stderr=log,
            start_new_session=True,
            close_fds=True,
        )
    assert kernel.stdin is not None
    kernel.stdin.write(token.encode() + b"\n")
    kernel.stdin.close()
    deadline = time.monotonic() + START_TIMEOUT
    while time.monotonic() < deadline:
        connection = connect(run_dir)
        if connection is not None:
            return connection
        if kernel.poll() is not None:
            break
        time.sleep(0.02)
    sys.stderr.write(
        f"albedo_bridge: kernel did not start (see {run_dir}/kernel.log)\n"
    )
    sys.exit(GONE)


def stdin(size: int) -> bytes:
    return os.read(0, size)


def stdout(data: bytes) -> int:
    return os.write(1, data)


def pump(source: Callable[[int], bytes], sink: Callable[[bytes], int | None]) -> None:
    try:
        while data := source(65536):
            albedo_link.write_all(sink, data)
    except OSError:
        pass
    os._exit(0)


def framed() -> list[str]:
    """The arguments a remote daemon sent as the first frame."""
    first = albedo_link.read_frame(stdin)
    spec = first.get("bridge") if isinstance(first, dict) else None
    argv = spec.get("argv") if isinstance(spec, dict) else None
    if not isinstance(argv, list) or not all(isinstance(a, str) for a in argv):
        sys.exit(2)
    cwd = spec.get("cwd") if isinstance(spec, dict) else None
    if isinstance(cwd, str):
        try:
            os.chdir(cwd)
        except OSError:
            sys.exit(NO_FOLDER)
    return [sys.argv[0], *argv]


def main(argv: list[str]) -> None:
    if argv[1:] == ["--frame"]:
        argv = framed()
    mode, run_dir = argv[1], argv[2]
    if mode == "launch":
        start(run_dir, argv[3], sys.stdin.readline().strip()).close()
        return
    if mode == "start":
        attach = albedo_link.read_frame(stdin)
        if not isinstance(attach, dict) or not isinstance(attach.get("attach"), dict):
            sys.exit(2)
        token = attach["attach"].get("token")
        if not isinstance(token, str):
            sys.exit(2)
        connection = start(run_dir, argv[3], token)
        first = albedo_link.encode(attach)
    else:
        found = connect(run_dir)
        if found is None:
            sys.exit(GONE)
        connection, first = found, None
    albedo_link.write_frame(
        stdout, albedo_link.encode({"bridge": {"bundle": albedo_bundle.digest()}})
    )
    if first is not None:
        albedo_link.write_frame(connection.send, first)
    threading.Thread(target=pump, args=(stdin, connection.send), daemon=True).start()
    pump(connection.recv, stdout)


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