#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
"""Bytes-only fixed worker result record and generation identity contract."""

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import struct


MAGIC = b"CHW10AQ1"
VERSION = 1
RECORD_SIZE = 128
HASHED_SIZE = 96
MAX_COPY_SIZE = 128 * 1024 * 1024
STATUS_SUCCESS = 0
STATUS_COPY_ERROR = 1
STATUS_RESTORE_ERROR = 2
STATUS_DEADLINE = 3
VALID_STATUSES = {STATUS_SUCCESS, STATUS_COPY_ERROR,
                  STATUS_RESTORE_ERROR, STATUS_DEADLINE}
PREFIX = struct.Struct("<8sHHIIIQQQII16s16s8s")


class WorkerRecordError(ValueError):
    """The fixed worker record is malformed or not precommitted."""


def _token(value: bytes, name: str) -> bytes:
    if type(value) is not bytes or len(value) != 16 or not any(value):
        raise WorkerRecordError(f"{name} must be a nonzero 16-byte token")
    return value


@dataclass(frozen=True)
class WorkerPrecommit:
    attempt_id: bytes
    nonce: bytes
    worker_pid: int
    child_pid: int
    generation: int
    requested: int

    def __post_init__(self) -> None:
        _token(self.attempt_id, "attempt ID")
        _token(self.nonce, "worker nonce")
        for value in (self.worker_pid, self.child_pid, self.generation,
                      self.requested):
            if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
                raise WorkerRecordError("precommit integer is invalid")
        if self.worker_pid == self.child_pid or self.worker_pid > 0xffffffff \
                or self.child_pid > 0xffffffff or self.generation > 0xffffffffffffffff \
                or self.requested > MAX_COPY_SIZE:
            raise WorkerRecordError("precommit identity or size is invalid")


@dataclass(frozen=True)
class WorkerRecord:
    status: int
    worker_pid: int
    child_pid: int
    generation: int
    requested: int
    copied: int
    restore_failure_bits: int
    flags: int
    attempt_id: bytes
    nonce: bytes

    @property
    def successful(self) -> bool:
        return (self.status == STATUS_SUCCESS and self.copied == self.requested
                and self.restore_failure_bits == 0 and self.flags == 0)


def encode_record(precommit: WorkerPrecommit, status: int, copied: int,
                  restore_failure_bits: int = 0, flags: int = 0) -> bytes:
    """Encode synthetic bytes; this function has no transport capability."""
    if type(precommit) is not WorkerPrecommit or status not in VALID_STATUSES:
        raise WorkerRecordError("record input is invalid")
    for value in (copied, restore_failure_bits, flags):
        if not isinstance(value, int) or isinstance(value, bool) or value < 0:
            raise WorkerRecordError("record numeric field is invalid")
    if copied > precommit.requested or restore_failure_bits > 3 \
            or flags > 0xffffffff:
        raise WorkerRecordError("record progress, restore bits or flags are invalid")
    if status == STATUS_SUCCESS and (copied != precommit.requested
                                     or restore_failure_bits or flags):
        raise WorkerRecordError("success record is not exact")
    prefix = PREFIX.pack(MAGIC, VERSION, RECORD_SIZE, status,
                         precommit.worker_pid, precommit.child_pid,
                         precommit.generation, precommit.requested, copied,
                         restore_failure_bits, flags, precommit.attempt_id,
                         precommit.nonce, bytes(8))
    if len(prefix) != HASHED_SIZE:
        raise WorkerRecordError("internal record layout differs")
    return prefix + hashlib.sha256(prefix).digest()


def parse_record(raw: bytes, expected: WorkerPrecommit) -> WorkerRecord:
    """Parse exactly one complete record and enforce the full precommit."""
    if type(raw) is not bytes or len(raw) != RECORD_SIZE \
            or type(expected) is not WorkerPrecommit:
        raise WorkerRecordError("record boundary is invalid")
    prefix, digest = raw[:HASHED_SIZE], raw[HASHED_SIZE:]
    if not hashlib.sha256(prefix).digest() == digest:
        raise WorkerRecordError("record digest differs")
    (magic, version, size, status, worker_pid, child_pid, generation,
     requested, copied, restore_bits, flags, attempt_id, nonce,
     reserved) = PREFIX.unpack(prefix)
    if magic != MAGIC or version != VERSION or size != RECORD_SIZE \
            or status not in VALID_STATUSES or any(reserved):
        raise WorkerRecordError("record header is invalid")
    if (worker_pid, child_pid, generation, requested, attempt_id, nonce) != (
            expected.worker_pid, expected.child_pid, expected.generation,
            expected.requested, expected.attempt_id, expected.nonce):
        raise WorkerRecordError("record identity differs from precommit")
    if copied > requested or restore_bits > 3 or flags != 0:
        raise WorkerRecordError("record result fields are invalid")
    record = WorkerRecord(status, worker_pid, child_pid, generation, requested,
                          copied, restore_bits, flags, attempt_id, nonce)
    if status == STATUS_SUCCESS and not record.successful:
        raise WorkerRecordError("success record is incomplete")
    if status == STATUS_RESTORE_ERROR and restore_bits == 0:
        raise WorkerRecordError("restore error has no failure bits")
    return record