"""The kernel's half of the durable session layer.

A detached kernel listens on a unix socket in its run directory. Whoever
attaches (a bridge the daemon runs, locally today and over ssh later) presents
the kernel's secret token; the newest attach wins and cuts the previous
connection off. Every frame after the handshake travels as
``{"seq": n, "ack": m, "frame": {...}}``: the kernel keeps what it sent until
the daemon acknowledges it and resends the rest on the next attach, and it
drops any daemon frame it has already seen, so a reconnect neither loses nor
repeats a message. A kernel run over plain stdio (the remote plugin's targets)
keeps the raw frames and ends with its pipe.
"""

from __future__ import annotations

import collections
import hmac
import json
import os
import socket
import struct
import threading
import time
from collections.abc import Callable
from typing import Any, BinaryIO

PROTOCOL = 1  # hello, snapshot and shutdown frames never change shape
MAX_FRAME = 8 * 1024 * 1024
SOCKET = "kernel.sock"
ATTACH_TIMEOUT = 10.0  # seconds a fresh connection has to present its token
OUTBOX_FRAMES = 4096
OUTBOX_BYTES = 64 * 1024 * 1024
# Frames whose newest copy supersedes every older unacknowledged one, by the
# field naming what they describe. Everything else must be delivered.
COALESCE = {"mirror": "handle", "trace": "id", "jobs": None}


def encode(value: object) -> bytes:
    return json.dumps(value, ensure_ascii=True, separators=(",", ":")).encode()


def envelope(seq: int, ack: int, frame: bytes) -> bytes:
    return b'{"seq":%d,"ack":%d,"frame":%s}' % (seq, ack, frame)


def write_all(write: Callable[[Any], int | None], data: bytes) -> None:
    remaining = memoryview(data)
    while remaining:
        written = write(remaining)
        remaining = remaining[len(remaining) if written is None else written :]


def write_frame(write: Callable[[Any], int | None], data: bytes) -> None:
    write_all(write, struct.pack(">I", len(data)) + data)


def read_exact(read: Callable[[int], bytes], size: int) -> bytes:
    data = bytearray()
    while len(data) < size:
        chunk = read(size - len(data))
        if not chunk:
            raise EOFError()
        data.extend(chunk)
    return bytes(data)


def read_frame(read: Callable[[int], bytes]) -> object:
    size = struct.unpack(">I", read_exact(read, 4))[0]
    if size > MAX_FRAME:
        raise ValueError("control frame too large")
    return json.loads(read_exact(read, size))


class Outbox:
    """Frames sent and not yet acknowledged, oldest first.

    Coalescing frames replace their older unacknowledged copy, so a flood of
    output mirrors costs one entry. Past the bound, the oldest coalescing
    frames go first; a must-deliver frame is dropped only when nothing else is
    left, and `dropped` counts every such loss for the next hello.
    """

    def __init__(self, frames: int = OUTBOX_FRAMES, size: int = OUTBOX_BYTES):
        self.limit_frames = frames
        self.limit_bytes = size
        self.seq = 0
        self.size = 0
        self.dropped = 0
        self.entries: collections.OrderedDict[int, tuple[bytes, object]] = (
            collections.OrderedDict()
        )
        self.keys: dict[object, int] = {}

    def add(self, frame: dict[str, object]) -> tuple[int, bytes]:
        data = encode(frame)
        self.seq += 1
        kind = frame.get("type")
        key: object = None
        if isinstance(kind, str) and kind in COALESCE:
            field = COALESCE[kind]
            key = (kind, frame.get(field) if field else None)
            older = self.keys.pop(key, None)
            if older is not None:
                self._remove(older)
        self.entries[self.seq] = (data, key)
        self.size += len(data)
        if key is not None:
            self.keys[key] = self.seq
        self._bound()
        return self.seq, data

    def ack(self, upto: int) -> None:
        while self.entries:
            seq = next(iter(self.entries))
            if seq > upto:
                return
            self._remove(seq)

    def pending(self) -> list[tuple[int, bytes]]:
        return [(seq, data) for seq, (data, _) in self.entries.items()]

    def _remove(self, seq: int) -> None:
        entry = self.entries.pop(seq, None)
        if entry is None:
            return
        data, key = entry
        self.size -= len(data)
        if key is not None and self.keys.get(key) == seq:
            del self.keys[key]

    def _bound(self) -> None:
        def over() -> bool:
            return len(self.entries) > self.limit_frames or self.size > self.limit_bytes

        if not over():
            return
        for seq in [seq for seq, (_, key) in self.entries.items() if key is not None]:
            if not over():
                return
            self._remove(seq)
        while over() and len(self.entries) > 1:
            self._remove(next(iter(self.entries)))
            self.dropped += 1


class Inbound:
    """The newest daemon frame applied; anything at or below it is a replay."""

    def __init__(self) -> None:
        self.last = 0

    def accept(self, seq: int) -> bool:
        if seq <= self.last:
            return False
        self.last = seq
        return True


class JobBook:
    """Job groups this kernel still owns, read from the frames it sends.

    The same bookkeeping the daemon keeps: a started job's group is owned
    until a job frame proves it gone, and so is the heavy-job slot it was
    granted. A reattaching daemon learns both from the hello, and the kernel
    stays alive past its grace while any job is live.
    """

    def __init__(self) -> None:
        self.started: dict[str, dict[str, object]] = {}
        self.held: set[str] = set()
        self.external = 0

    def hold(self, id: str) -> None:
        """The daemon granted this job a slot."""
        self.held.add(id)

    def observe(self, frame: dict[str, object]) -> None:
        kind, id = frame.get("type"), frame.get("id")
        if kind == "job_start" and isinstance(id, str):
            pgid = frame.get("pgid")
            if isinstance(pgid, int) and not isinstance(pgid, bool) and pgid > 1:
                self.started[id] = frame
        elif kind == "job" and isinstance(id, str):
            cleanup = frame.get("cleanup")
            if isinstance(cleanup, dict) and cleanup.get("gone") is True:
                self.started.pop(id, None)
                self.held.discard(id)
        elif kind == "jobs":
            live = frame.get("live")
            if isinstance(live, int) and not isinstance(live, bool) and live >= 0:
                self.external = live

    def live(self) -> int:
        return len(self.started) + self.external


class StdioLink:
    """Raw frames over the inherited stdin/stdout; the pipe closing ends us."""

    def __init__(self) -> None:
        self.control_in: BinaryIO = os.fdopen(os.dup(0), "rb", buffering=0)
        self.control_out: BinaryIO = os.fdopen(os.dup(1), "wb", buffering=0)
        self.lock = threading.Lock()
        self.jobs = JobBook()

    def send(self, frame: dict[str, object]) -> None:
        data = encode(frame)
        with self.lock:
            self.jobs.observe(frame)
            write_frame(self.control_out.write, data)

    def serve(self, dispatch: Callable[[object], None]) -> None:
        """Deliver frames until the pipe closes; raises when it does."""
        while True:
            dispatch(read_frame(self.control_in.read))

    def attached(self) -> bool:
        return True


class SocketLink:
    """The session layer over a unix socket in the kernel's run directory."""

    def __init__(
        self,
        run_dir: str,
        token: str,
        hello: Callable[[], dict[str, object]],
        *,
        bundle: str,
        grace: float,
    ) -> None:
        self.run_dir = run_dir
        self.token = token
        self.hello = hello
        self.bundle = bundle
        self.grace = grace
        self.lock = threading.RLock()
        self.outbox = Outbox()
        self.inbound = Inbound()
        self.jobs = JobBook()
        self.epoch = 0
        self.connection: socket.socket | None = None
        self.detached_at = time.monotonic()
        self.listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        # Bound relative to the run directory: sun_path holds ~104 bytes and a
        # temporary ALBEDO_HOME alone can use most of that.
        previous = os.getcwd()
        os.chdir(run_dir)
        try:
            if os.path.exists(SOCKET):
                os.unlink(SOCKET)
            self.listener.bind(SOCKET)
            os.chmod(SOCKET, 0o600)
        finally:
            os.chdir(previous)
        self.listener.listen(4)

    def send(self, frame: dict[str, object]) -> None:
        with self.lock:
            self.jobs.observe(frame)
            seq, data = self.outbox.add(frame)
            self._write(envelope(seq, self.inbound.last, data))

    def attached(self) -> bool:
        return self.connection is not None

    def idle_for(self) -> float:
        """Seconds with nothing attached; zero while attached."""
        if self.connection is not None:
            return 0.0
        return time.monotonic() - self.detached_at

    def serve(self, dispatch: Callable[[object], None]) -> None:
        """Accept attaches forever; each connection reads on its own thread."""
        while True:
            connection, _ = self.listener.accept()
            threading.Thread(
                target=self._session, args=(connection, dispatch), daemon=True
            ).start()

    def _write(self, data: bytes) -> None:
        connection = self.connection
        if connection is None:
            return
        try:
            write_frame(connection.send, data)
        except OSError:
            self._detach(connection)

    def _detach(self, connection: socket.socket) -> None:
        with self.lock:
            if self.connection is connection:
                self.connection = None
                self.detached_at = time.monotonic()
        try:
            connection.close()
        except OSError:
            pass

    def _attach(self, connection: socket.socket) -> bool:
        connection.settimeout(ATTACH_TIMEOUT)
        message = read_frame(connection.recv)
        attach = message.get("attach") if isinstance(message, dict) else None
        token = attach.get("token") if isinstance(attach, dict) else None
        if not isinstance(attach, dict) or not isinstance(token, str):
            write_frame(connection.send, encode({"refused": "expected attach"}))
            return False
        if not hmac.compare_digest(token.encode(), self.token.encode()):
            write_frame(connection.send, encode({"refused": "wrong token"}))
            return False
        connection.settimeout(None)
        grace = attach.get("grace")
        ack = attach.get("ack")
        with self.lock:
            if isinstance(grace, int | float) and not isinstance(grace, bool):
                self.grace = max(float(grace), 0.0)
            previous, self.connection = self.connection, connection
            self.epoch += 1
            if previous is not None:
                try:
                    previous.shutdown(socket.SHUT_RDWR)
                except OSError:
                    pass
            if isinstance(ack, int) and not isinstance(ack, bool):
                self.outbox.ack(ack)
            hello = {
                **self.hello(),
                "protocol": PROTOCOL,
                "bundle": self.bundle,
                "epoch": self.epoch,
                "ack": self.inbound.last,
                "jobs": list(self.jobs.started.values()),
                "held": sorted(self.jobs.held),
                "external": self.jobs.external,
                "dropped": self.outbox.dropped,
            }
            write_frame(connection.send, encode({"hello": hello}))
            for seq, data in self.outbox.pending():
                write_frame(connection.send, envelope(seq, self.inbound.last, data))
        return True

    def _session(
        self, connection: socket.socket, dispatch: Callable[[object], None]
    ) -> None:
        try:
            if not self._attach(connection):
                return
            while True:
                message = read_frame(connection.recv)
                if not isinstance(message, dict):
                    raise ValueError("expected an envelope")
                ack = message.get("ack")
                seq = message.get("seq")
                with self.lock:
                    if isinstance(ack, int) and not isinstance(ack, bool):
                        self.outbox.ack(ack)
                    fresh = isinstance(seq, int) and self.inbound.accept(seq)
                if fresh:
                    dispatch(message.get("frame"))
                if isinstance(seq, int):
                    with self.lock:
                        if self.connection is connection:
                            self._write(encode({"ack": self.inbound.last}))
        except (OSError, EOFError, ValueError):
            pass
        finally:
            self._detach(connection)