#!/usr/bin/env python3
# SPDX-License-Identifier: CC0-1.0
"""Test two consumers of tee(2)-duplicated pipe buffers, not the tee(1) command."""
import argparse
import contextlib
import ctypes
import datetime
import importlib.util
import json
import os
from pathlib import Path
import platform
import subprocess
import tempfile

HERE = Path(__file__).resolve().parent
spec = importlib.util.spec_from_file_location("lifetime", HERE / "zero-copy-lifetime.py")
base = importlib.util.module_from_spec(spec)
spec.loader.exec_module(base)
libc = ctypes.CDLL(None, use_errno=True)
libc.tee.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_size_t, ctypes.c_uint]
libc.tee.restype = ctypes.c_ssize_t


def duplicate(source, destination, count):
    result = libc.tee(source, destination, count, 0)
    if result < 0:
        error = ctypes.get_errno()
        raise OSError(error, os.strerror(error))
    base.require(result == count, "short tee; no observation published")
    return result


def probe(mode, parent=None):
    if mode not in ("sendfile", "buffered"):
        raise ValueError("unknown mode")
    size = os.sysconf("SC_PAGESIZE")
    with tempfile.TemporaryDirectory(prefix="hesela-tee-", dir=parent) as directory:
        with contextlib.ExitStack() as cleanup:
            def own(fd):
                cleanup.callback(os.close, fd)
                return fd
            path = Path(directory) / "source"
            fd = own(os.open(path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600))
            base.write_all(fd, b"A" * size)
            os.fsync(fd)
            r1, w1 = os.pipe2(os.O_NONBLOCK | os.O_CLOEXEC)
            r2, w2 = os.pipe2(os.O_NONBLOCK | os.O_CLOEXEC)
            for pipe_fd in (r1, w1, r2, w2):
                own(pipe_fd)
            for pipe_fd in (w1, w2):
                base.require(base.fcntl.fcntl(pipe_fd, base.fcntl.F_GETPIPE_SZ) >= size, "small pipe")
            before = os.pread(fd, size, 0)
            base.require(before == b"A" * size, "unexpected initial payload")
            if mode == "sendfile":
                count = os.sendfile(w1, fd, 0, size)
            else:
                count = os.write(w1, before)
            base.require(count == size, "short enqueue")
            copied = duplicate(r1, w2, size)
            first = base.read_exact(r1, size)
            first_summary = base.payload_summary(first)
            base.write_all(fd, b"B" * size)
            second = base.read_exact(r2, size)
            second_summary = base.payload_summary(second)
            return {
                "mode": mode, "page_size_bytes": size,
                "enqueue_return_bytes": count, "tee_return_bytes": copied,
                "events": ["enqueue", "tee", "read-first", "hash-first", "overwrite-source", "read-second", "hash-second"],
                "first_consumer": first_summary,
                "second_consumer": second_summary,
                "consumer_hashes_equal": first_summary["sha256"] == second_summary["sha256"],
            }


def run(repetitions=5, parent=None):
    if type(repetitions) is not int or not 1 <= repetitions <= 100:
        raise ValueError("repetitions must be an integer from 1 to 100")
    rows = []
    with tempfile.TemporaryDirectory(prefix="hesela-tee-env-", dir=parent) as directory:
        filesystem = subprocess.run(["stat", "-f", "-c", "%T", directory],
                                    check=True, capture_output=True, text=True, timeout=5).stdout.strip()
        for repetition in range(1, repetitions + 1):
            for mode in ("sendfile", "buffered"):
                rows.append({"repetition": repetition, **probe(mode, directory)})
    return {
        "schemaVersion": 1, "license": "CC0-1.0",
        "generatedAt": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "method": "Duplicate a one-page pipe with tee(2), consume and hash the first branch, overwrite the file, then consume and hash the other branch.",
        "limits": ["This uses the Linux tee system call, not a test of GNU tee.",
                   "Mutation deliberately violates the source-lifetime requirement.",
                   "One host; no TCP, checksum-algorithm failure, performance or crash measurement."],
        "environment": {"kernel": platform.release(), "python": platform.python_version(),
                        "machine": platform.machine(), "filesystem_magic_label": filesystem},
        "repetitions": repetitions, "casesPerRepetition": 2,
        "sourceSha256": base.digest(Path(__file__).read_bytes()),
        "helperSha256": base.digest((HERE / "zero-copy-lifetime.py").read_bytes()),
        "observations": rows,
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repetitions", type=int, default=5)
    parser.add_argument("--temp-dir")
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    result = json.dumps(run(args.repetitions, args.temp_dir), indent=2) + "\n"
    if args.output:
        args.output.write_text(result, encoding="utf-8")
    else:
        print(result, end="")