"""Benchmark: ChaCha20 encrypt via update() (returns bytes, extra memcpy from
the Rust temp buffer into the PyBytes) vs update_into() writing directly into
a caller-provided buffer.

Two scenarios:
  one-shot : a fresh Cipher per call (mirrors chab3._session, one per message)
  stream   : a reused encryptor   (mirrors chab3.random CSPRNG); isolates the
             allocation+memcpy difference from cipher setup cost.

The bytes-hack row from test1 is absent: update_into rejects readonly buffers.
"fresh bytearray" has the same allocate-zeroed-buffer profile a hypothetical
fresh bytes target would have.
"""

import sys
import timeit

from cryptography.hazmat.primitives.ciphers import Cipher
from cryptography.hazmat.primitives.ciphers.algorithms import ChaCha20

KEY = bytes(range(32))
NONCE = bytes(16)
SIZES = [16, 64, 256, 1024, 16384, 65536, 1048576]
REPEAT = 5


def iterations(size):
    if size <= 1024:
        return 5000
    if size <= 16384:
        return 1000
    if size <= 65536:
        return 300
    if size <= 1 << 20:
        return 20
    return 1


def make_enc():
    return Cipher(ChaCha20(KEY, NONCE), mode=None).encryptor()


def time_call(fn, number):
    return min(timeit.repeat(fn, number=number, repeat=REPEAT)) / number


def oneshot(size):
    data = bytes(size)
    out = bytearray(size)
    mv = memoryview(out)
    n = iterations(size)
    return {
        "update -> bytes": time_call(lambda: make_enc().update(data), n),
        "into fresh bytearray": time_call(
            lambda: make_enc().update_into(data, bytearray(size)), n
        ),
        "into reused bytearray": time_call(
            lambda: make_enc().update_into(data, out), n
        ),
        "into reused mv-slice": time_call(
            lambda: make_enc().update_into(data, mv[:size]), n
        ),
    }


def stream(size):
    data = bytes(size)
    out = bytearray(size)
    mv = memoryview(out)
    encs = [make_enc() for _ in range(4)]
    n = iterations(size)
    return {
        "update -> bytes": time_call(lambda: encs[0].update(data), n),
        "into fresh bytearray": time_call(
            lambda: encs[1].update_into(data, bytearray(size)), n
        ),
        "into reused bytearray": time_call(lambda: encs[2].update_into(data, out), n),
        "into reused mv-slice": time_call(
            lambda: encs[3].update_into(data, mv[:size]), n
        ),
    }


def report(title, scenario):
    print(f"\n=== {title} ===")
    header = f"{'size':>9} | {'variant':<22} | {'ns/op':>12} | {'MB/s':>9} | vs update"
    print(header)
    print("-" * len(header))
    for size in SIZES:
        results = scenario(size)
        base = results["update -> bytes"]
        for name, t in results.items():
            print(
                f"{size:>9} | {name:<22} | {t * 1e9:>12.0f} | "
                f"{size / t / 1e6:>9.1f} | {t / base:>8.3f}"
            )


def main():
    global SIZES
    SIZES = [int(s, 0) for s in sys.argv[1:]] or SIZES
    report("one-shot: fresh Cipher per call (chab3 message path)", oneshot)
    report("stream: reused encryptor (chab3 CSPRNG path)", stream)


if __name__ == "__main__":
    main()