"""Private root I/O progress observer for the C01 supervisor (wire version1).

The supervisor retains the accepted application endpoint and the direct worker's
pidfd/cgroup. This observer revokes that endpoint on deadline/protocol failure;
the caller must request retirement and retain the worker until actually reaped.
It never waits for filesystem completion or creates a replacement worker.
"""
import array
import os
import socket
import struct
import time

PACKET = struct.Struct('!IIQQ')
VERSION, READY, STARTED, COMPLETED = 1, 1, 2, 3
IO_LIMIT_NS = 2_000_000_000
SETUP_LIMIT_NS = 10_000_000_000


def now():
    return time.clock_gettime_ns(time.CLOCK_BOOTTIME)


def require(condition, message):
    if not condition:
        raise ValueError(message)


class HostWatch:
    def __init__(self, control, application, worker_pid, setup_deadline):
        require(os.geteuid() == 0 and type(worker_pid) is int and worker_pid > 1,
                'host observer requires root and a direct worker')
        require(control.family == socket.AF_UNIX and
                control.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) == socket.SOCK_SEQPACKET and
                control.getsockopt(socket.SOL_SOCKET, socket.SO_PASSCRED) == 1,
                'enable credential delivery before worker fork')
        require(application.family == socket.AF_UNIX and
                application.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) == socket.SOCK_SEQPACKET,
                'host observer requires retained application revocation endpoint')
        stamp = now()
        require(type(setup_deadline) is int and stamp < setup_deadline <= stamp + SETUP_LIMIT_NS,
                'invalid host setup deadline')
        control.setblocking(False)
        control.set_inheritable(False)
        application.set_inheritable(False)
        self.control, self.application = control, application
        self.worker_pid = worker_pid
        self.state = 'starting'
        self.sequence = 0
        self.deadline = setup_deadline
        self.failure = None

    def revoke(self, reason):
        if self.state == 'failed':
            return
        self.state, self.failure = 'failed', reason
        # shutdown affects the worker's duplicate endpoint too. Merely closing
        # this alias would not revoke a worker blocked in a filesystem syscall.
        try:
            self.application.shutdown(socket.SHUT_RDWR)
        except OSError:
            pass
        self.control.close()

    def check_deadline(self):
        if self.state == 'failed':
            return False
        if self.deadline is not None and now() >= self.deadline:
            self.revoke('host worker deadline expired')
            return False
        return True

    def receive(self):
        if not self.check_deadline():
            return False
        try:
            raw, ancillary, flags, _ = self.control.recvmsg(
                PACKET.size + 1, socket.CMSG_SPACE(12) + socket.CMSG_SPACE(16 * array.array('i').itemsize),
                socket.MSG_DONTWAIT | socket.MSG_CMSG_CLOEXEC)
        except BlockingIOError:
            return self.check_deadline()
        except OSError:
            self.revoke('host worker control disconnected')
            return False
        credentials, unexpected = [], False
        # Close any installed rights even on truncation or a malformed message.
        for level, kind, data in ancillary:
            if (level, kind) == (socket.SOL_SOCKET, socket.SCM_RIGHTS):
                rights = array.array('i')
                rights.frombytes(data[:len(data) - len(data) % rights.itemsize])
                for fd in rights:
                    os.close(fd)
                unexpected = True
            elif (level, kind) == (socket.SOL_SOCKET, socket.SCM_CREDENTIALS) and len(data) == 12:
                credentials.append(struct.unpack('3i', data))
            else:
                unexpected = True
        try:
            require(not unexpected and flags & ~socket.MSG_CMSG_CLOEXEC == 0 and
                    len(raw) == PACKET.size and credentials == [(self.worker_pid, 0, 0)],
                    'host progress sender/framing refused')
            require(self.check_deadline(), 'host progress arrived after deadline')
            version, event, sequence, deadline = PACKET.unpack(raw)
            require(version == VERSION, 'host progress version refused')
            if event == READY:
                require(self.state == 'starting' and sequence == 0 and deadline == 0,
                        'host readiness replay refused')
                self.state, self.deadline = 'idle', None
            elif event == STARTED:
                stamp = now()
                require(self.state == 'idle' and sequence == self.sequence + 1 and
                        stamp < deadline <= stamp + IO_LIMIT_NS, 'host start/replay/deadline refused')
                self.state, self.deadline = 'busy', deadline
            elif event == COMPLETED:
                require(self.state == 'busy' and sequence == self.sequence + 1 and deadline == self.deadline,
                        'host completion/replay/deadline refused')
                self.sequence, self.state, self.deadline = sequence, 'idle', None
            else:
                raise ValueError('unknown host progress event')
        except ValueError:
            self.revoke('host worker progress refused')
        return self.state != 'failed'

    def worker_exited(self):
        self.revoke('host worker exited')

    def release_reaped(self):
        """Caller invokes only after waitpid confirms the owned worker is reaped."""
        self.revoke('host worker retired')
        self.application.close()