#!/usr/bin/env python3
"""Private monotonic runtime-generation state shared by runtime hook processes."""

from __future__ import annotations

import fcntl
import os
import stat
from collections.abc import Mapping
from pathlib import Path
from typing import Final

MAX_GENERATION: Final = (1 << 63) - 1
MAX_GENERATION_BYTES: Final = 32


def parse_generation(value: object) -> int:
    if not isinstance(value, str) or not value.isascii() or not value.isdigit():
        raise ValueError("invalid runtime generation")
    generation = int(value)
    if generation < 0 or generation > MAX_GENERATION:
        raise ValueError("invalid runtime generation")
    return generation


def _validate_descriptor(descriptor: int) -> None:
    metadata = os.fstat(descriptor)
    if not stat.S_ISREG(metadata.st_mode):
        raise ValueError("runtime generation state is not a regular file")
    if metadata.st_uid != os.geteuid():
        raise ValueError("runtime generation state has the wrong owner")
    if stat.S_IMODE(metadata.st_mode) & 0o077:
        raise ValueError("runtime generation state permissions are not private")
    if metadata.st_size > MAX_GENERATION_BYTES:
        raise ValueError("runtime generation state is oversized")


def _read_descriptor(descriptor: int) -> int:
    os.lseek(descriptor, 0, os.SEEK_SET)
    raw = os.read(descriptor, MAX_GENERATION_BYTES + 1)
    if len(raw) > MAX_GENERATION_BYTES:
        raise ValueError("runtime generation state is oversized")
    try:
        return parse_generation(raw.decode("ascii").strip())
    except UnicodeDecodeError as exc:
        raise ValueError("invalid runtime generation state") from exc


def _write_descriptor(descriptor: int, generation: int) -> None:
    payload = f"{generation}\n".encode("ascii")
    os.lseek(descriptor, 0, os.SEEK_SET)
    os.ftruncate(descriptor, 0)
    remaining = memoryview(payload)
    while remaining:
        written = os.write(descriptor, remaining)
        if written <= 0:
            raise OSError("runtime generation write made no progress")
        remaining = remaining[written:]
    os.fsync(descriptor)


def initialize_runtime_generation(path: Path, generation: int) -> None:
    if generation < 0 or generation > MAX_GENERATION:
        raise ValueError("invalid runtime generation")
    descriptor = os.open(
        path,
        os.O_RDWR | os.O_CREAT | os.O_TRUNC | os.O_CLOEXEC | os.O_NOFOLLOW,
        0o600,
    )
    try:
        os.fchmod(descriptor, 0o600)
        _validate_descriptor(descriptor)
        _write_descriptor(descriptor, generation)
    finally:
        os.close(descriptor)


def read_runtime_generation(environ: Mapping[str, str]) -> int:
    state_path = environ.get("MOSAIC_LEASE_GENERATION_FILE")
    if not state_path:
        return parse_generation(environ["MOSAIC_RUNTIME_GENERATION"])
    descriptor = os.open(state_path, os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW)
    try:
        _validate_descriptor(descriptor)
        return _read_descriptor(descriptor)
    finally:
        os.close(descriptor)


def bump_runtime_generation(environ: Mapping[str, str]) -> int:
    state_path = environ.get("MOSAIC_LEASE_GENERATION_FILE")
    if not state_path:
        raise ValueError("runtime generation file is required for same-PID rollover")
    descriptor = os.open(state_path, os.O_RDWR | os.O_CLOEXEC | os.O_NOFOLLOW)
    try:
        fcntl.flock(descriptor, fcntl.LOCK_EX)
        _validate_descriptor(descriptor)
        current = _read_descriptor(descriptor)
        if current >= MAX_GENERATION:
            raise ValueError("runtime generation exhausted")
        generation = current + 1
        _write_descriptor(descriptor, generation)
        return generation
    finally:
        os.close(descriptor)