#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: 2026 Martino Ferrari
"""A minimal UDPStreamer (UDPS) producer, for exercising udpscope without MARTe2.

Unicast only: binds `--port`, waits for a CONNECT datagram, then streams CONFIG
once and DATA packets to whoever connected. Publish mode is Accumulate, so each
packet carries `--batch` consecutive samples per (scalar) signal -- the shape the
viewer's per-sample path expects.

Signals: `time` (the stream's own clock, seconds), `sine`, `square`, `ramp`.
Everything is float64 and unquantised.

    ./tools/mock_udpstreamer.py --port 8000 --rate 10000 --batch 100
"""

import argparse
import math
import random
import socket
import struct
import time

import numpy as np

MAGIC = 0x53504455
PKT_DATA, PKT_CONFIG, PKT_ACK, PKT_CONNECT, PKT_DISCONNECT = 0, 1, 2, 3, 4
T_FLOAT64 = 9
QUANT_NONE = 0
TIME_PACKET, TIME_FULL_ARRAY, TIME_FIRST_SAMPLE = 0, 1, 2
PUBLISH_ACCUMULATE = 1
NO_TIME_SIGNAL = 0xFFFFFFFF


def header(pkt_type, counter, payload_bytes):
    return struct.pack("<IBIHHI", MAGIC, pkt_type, counter, 0, 1, payload_bytes)


def signal_desc(name, *, rate, time_mode, time_idx, unit, rows=1, cols=1, dims=0):
    """One 136-byte signal descriptor."""
    d = bytearray(136)
    d[0:64] = name.encode()[:64].ljust(64, b"\0")
    d[64] = T_FLOAT64
    d[65] = QUANT_NONE
    d[66] = dims
    struct.pack_into("<II", d, 67, rows, cols)
    struct.pack_into("<dd", d, 75, -10.0, 10.0)
    d[91] = time_mode
    struct.pack_into("<d", d, 92, rate)
    struct.pack_into("<I", d, 100, time_idx)
    d[104:136] = unit.encode()[:32].ljust(32, b"\0")
    return bytes(d)


def build_config(rate, time_mode=TIME_FULL_ARRAY, unit="V"):
    sigs = [
        # The stream's clock: referenced by every data signal, so udpscope treats
        # it as the time base rather than a plottable channel. Its *own* time_mode
        # is what udpscope switches on, so `--time-mode first` has to land here.
        signal_desc(
            "time", rate=0.0, time_mode=time_mode, time_idx=NO_TIME_SIGNAL, unit="s"
        ),
        signal_desc("sine", rate=rate, time_mode=time_mode, time_idx=0, unit=unit),
        signal_desc("square", rate=rate, time_mode=time_mode, time_idx=0, unit=unit),
        signal_desc("ramp", rate=rate, time_mode=time_mode, time_idx=0, unit=unit),
        signal_desc("noise", rate=rate, time_mode=time_mode, time_idx=0, unit=unit),
    ]
    payload = struct.pack("<I", len(sigs)) + b"".join(sigs)
    payload += bytes([PUBLISH_ACCUMULATE])
    return header(PKT_CONFIG, 0, len(payload)) + payload


def build_data(
    counter,
    t0,
    dt,
    batch,
    freq,
    time_mode=TIME_FULL_ARRAY,
    noise=0.0,
    time_spike=0.0,
    glitch_at=-1.0,
    glitch_width=0.0,
    quiet_at=-1.0,
):
    stamps = [t0 + i * dt for i in range(batch)]
    # `time_spike` models a device whose clock glitches for one packet: the stamps jump
    # far ahead and then the next packet carries on from where it would have been. The
    # waveforms are still computed from the *true* stamps, so only the time signal is
    # wrong -- which is what a glitching clock looks like on the wire.
    wire_stamps = [t + time_spike for t in stamps] if time_spike > 0.0 else stamps
    # The interferer is a tone at 100x the fundamental with a randomly scaled amplitude, which
    # is the shape that is hardest on anything counting crossings: near the sine's own mid level
    # the fundamental's slope is 2*pi*f while the interferer's is 100x that, so the signal
    # recrosses the level several times per true crossing. `--noise 0` is the ideal waveform, and
    # it is the default because the documented probe assertions are *exact* (`peak to peak=2`,
    # `top=0.636620`) and any interferer moves the extremes.
    if noise > 0.0:
        sine = [
            math.sin(2.0 * math.pi * freq * t)
            + random.random() * noise * math.sin(2.0 * math.pi * 100.0 * freq * t)
            for t in stamps
        ]
    else:
        sine = [math.sin(2.0 * math.pi * freq * t) for t in stamps]
    square = [
        1.0 if math.sin(2.0 * math.pi * freq * 0.5 * t) >= 0.0 else -1.0 for t in stamps
    ]
    ramp = [(freq * t) % 1.0 * 2.0 - 1.0 for t in stamps]
    noisy = np.sin(100 * math.pi * np.array(stamps) + 0.5) + np.random.normal(
        0, 0.5, batch
    )
    # A glitch is an *inversion* of the square for `glitch_width` seconds, repeated once per
    # square period (2/freq -- the period of sin(2*pi*freq*0.5*t)). Inverting rather than
    # forcing a level makes it a narrow pulse whatever the square's polarity is, and repeating
    # it lets a probe assert the *spacing* of the hits a width condition finds instead of an
    # absolute count, which depends on when the probe happened to start looking.
    if glitch_at >= 0.0 and glitch_width > 0.0:
        sq_period = 2.0 / freq
        for i, t in enumerate(stamps):
            if t < glitch_at:
                continue
            k = math.floor((t - glitch_at) / sq_period)
            start = glitch_at + k * sq_period
            if t < start + glitch_width:
                square[i] = -square[i]
    # `--quiet-at` is for the dropout condition: the sine stops at a level on the *opposite*
    # side of zero from where it was, so `quiet_at` itself is its last zero crossing and a
    # dropout at level 0 fires at exactly `quiet_at + timeout`. Holding at the same sign would
    # leave the last crossing at whichever earlier zero happened to precede it, which is not a
    # figure a probe can assert. Choose a `quiet_at` away from the sine's own zeros.
    if quiet_at >= 0.0:
        held = -0.5 if math.sin(2.0 * math.pi * freq * quiet_at) >= 0.0 else 0.5
        for i, t in enumerate(stamps):
            if t >= quiet_at:
                sine[i] = held
    # In FIRST_SAMPLE mode only element 0 of the time signal is meaningful: the producer
    # stamps the first sample of the batch and says nothing about the rest. A scalar
    # signal in Accumulate mode still occupies `batch` slots on the wire (the client
    # reads `num_samples` values for any 1-element signal), so the anchor is repeated --
    # which is exactly the shape that made udpscope derive a step of zero.
    wire_time = (
        [wire_stamps[0]] * batch if time_mode == TIME_FIRST_SAMPLE else wire_stamps
    )
    payload = struct.pack("<QI", int(t0 * 1e9), batch)
    for arr in (wire_time, sine, square, ramp, noisy):
        payload += struct.pack("<%dd" % batch, *arr)
    return header(PKT_DATA, counter, len(payload)) + payload


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8000)
    ap.add_argument("--rate", type=float, default=10000.0, help="samples per second")
    ap.add_argument("--batch", type=int, default=100, help="samples per packet")
    ap.add_argument("--freq", type=float, default=5.0, help="waveform frequency, Hz")
    ap.add_argument("--duration", type=float, default=0.0, help="0 = forever")
    ap.add_argument(
        "--stream-scale",
        type=float,
        default=1.0,
        help="stream seconds per wall second: 0.1 mimics a producer whose own "
        "clock advances ten times slower than the wire delivers samples",
    )
    ap.add_argument(
        "--declared-rate",
        type=float,
        default=-1.0,
        help="sampling_rate to put in CONFIG; -1 = the real rate, 0 = 'unknown'",
    )
    ap.add_argument(
        "--time-mode",
        choices=("full", "first"),
        default="full",
        help="how the stream carries its clock: 'full' stamps every sample "
        "(UDPS_TIME_FULL_ARRAY), 'first' stamps only the batch's first "
        "sample (UDPS_TIME_FIRST_SAMPLE) and leaves the viewer to space the "
        "rest by the declared cadence -- the shape a real MARTe2-style "
        "producer with a single u64 clock has",
    )
    ap.add_argument(
        "--unit",
        default="V",
        help="the unit the data signals declare in CONFIG; '' declares none, which "
        "is an ordinary stream (and what every math channel looks like) and is "
        "the only way a probe can see the unitless readout formatting",
    )
    ap.add_argument(
        "--noise",
        type=float,
        default=0.0,
        help="amplitude of an interferer added to `sine`: a tone at 100x the fundamental, "
        "scaled by a fresh uniform random factor per sample. 0 = the ideal sine, which is "
        "what the exact probe assertions (peak to peak = 2, top = 2/pi) are read against; "
        "0.1 is the figure the measurement hysteresis is sized against",
    )
    ap.add_argument(
        "--clock-jitter",
        type=float,
        default=0.0,
        help="per-packet stamp jitter, as a fraction of one packet's span: a device "
        "that stamps each batch from its own clock does not land on a perfect "
        "grid, so the stream's timestamps wobble about one. Non-cumulative "
        "(the mean rate is unchanged); 0 = the ideal grid",
    )
    ap.add_argument(
        "--time-glitch-at",
        type=float,
        default=-1.0,
        help="stream time (s) at which ONE packet's time signal jumps ahead by "
        "--time-glitch-by and then carries on normally: a device whose clock "
        "glitches. -1 = never",
    )
    ap.add_argument(
        "--time-glitch-by",
        type=float,
        default=10.0,
        help="how far ahead the glitching packet's stamps jump, in stream seconds",
    )
    ap.add_argument(
        "--glitch-at",
        type=float,
        default=-1.0,
        help="stream time (s) of the first narrow pulse injected into `square`: the square "
        "is inverted for --glitch-width seconds, once per square period (2/freq), so a "
        "pulse-width or glitch trigger has a repeating event to find. -1 = never",
    )
    ap.add_argument(
        "--glitch-width",
        type=float,
        default=0.002,
        help="width of each injected pulse, in stream seconds (default 0.002 s; a zero width "
        "silently disables the glitch even when --glitch-at is set)",
    )
    ap.add_argument(
        "--quiet-at",
        type=float,
        default=-1.0,
        help="stream time (s) after which `sine` is held at a constant level on the "
        "opposite side of zero, so --quiet-at is its last zero crossing: the shape a "
        "dropout trigger is read against. -1 = never",
    )
    ap.add_argument(
        "--drop-every",
        type=int,
        default=0,
        help="withhold every Nth DATA packet from the wire -- packet loss. Deterministic "
        "rather than a probability for the reason --glitch-at gives: a probe asserts the "
        "spacing of the events it finds, and a random count is not a figure. The counter "
        "and the stream clock advance over a withheld packet, which is the whole point: "
        "the receiver orders the stream by the counter alone, so that is the only thing "
        "that tells it a packet existed. 0 = never",
    )
    ap.add_argument(
        "--drop-burst",
        type=int,
        default=1,
        help="consecutive packets each --drop-every event withholds: one models a lost "
        "datagram, a long run models a link that went away. The run length is what the "
        "ingest's own glitch guard is sized against, since it bridges a packet whose "
        "stamps jump more than 256 packet spans ahead. Measured at 1000 packets a second, "
        "a burst of 200 is counted as loss and nothing else; one of 300 is read as a clock "
        "glitch instead, and because each bridge advances the anchor by one packet span "
        "while the true timeline is 300 ahead, the four following packets fail too and the "
        "stream falls back to arrival timing. Keep a burst under the 2 s silence timeout "
        "or the client reconnects and the mock restarts its clock, which is a different "
        "test",
    )
    ap.add_argument(
        "--reorder-every",
        type=int,
        default=0,
        help="hold every Nth packet back and send it after the next one that goes out -- "
        "an inversion on the wire. What the receiver does depends on how late the held "
        "packet arrives: inside the 20 ms reorder window it is resequenced with no hole; "
        "past it a hole is committed first and the packet backfills it. "
        "--reorder-delay-ms controls which path is taken (see its help). "
        "A --drop-every that swallows the packet after a held one deepens the inversion "
        "rather than cancelling it, since the held packet waits for the next packet "
        "actually sent. 0 = never",
    )
    ap.add_argument(
        "--reorder-delay-ms",
        type=float,
        default=0.0,
        metavar="MS",
        help="hold a withheld packet this many milliseconds instead of releasing it "
        "immediately behind the packet that overtook it. Under the receiver's reorder "
        "window the packet is resequenced and no hole is ever seen; past it the hole is "
        "committed first and the packet has to backfill it -- which is the only way to "
        "reach the backfill path at all, since a 0 ms delay always wins the race.",
    )
    args = ap.parse_args()

    dt = args.stream_scale / args.rate  # stream seconds per sample
    period = args.batch / args.rate  # wall seconds per packet
    declared = args.rate if args.declared_rate < 0.0 else args.declared_rate
    time_mode = TIME_FIRST_SAMPLE if args.time_mode == "first" else TIME_FULL_ARRAY

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(("0.0.0.0", args.port))
    sock.settimeout(0.001)
    print(
        "mock udpstreamer: waiting for CONNECT on :%d (%.0f S/s wire, %d per packet, "
        "stream clock x%.3g, declared %.0f S/s, time-mode %s)"
        % (
            args.port,
            args.rate,
            args.batch,
            args.stream_scale,
            declared,
            args.time_mode,
        ),
        flush=True,
    )

    peer = None
    counter = 1
    glitched = False
    t_stream = 0.0
    started = None
    next_send = time.monotonic()
    held = None  # (counter, datagram, release_deadline) built but withheld, for --reorder-every
    drop_left = 0
    dropped = 0
    inverted = 0
    noted_drop = False
    noted_reorder = False

    while True:
        try:
            buf, addr = sock.recvfrom(2048)
            if len(buf) >= 17:
                magic, ptype = struct.unpack_from("<IB", buf, 0)
                if magic == MAGIC and ptype == PKT_CONNECT:
                    peer = addr
                    counter = 1
                    glitched = False
                    t_stream = 0.0
                    started = time.monotonic()
                    next_send = started
                    held = None
                    drop_left = 0
                    sock.sendto(build_config(declared, time_mode, args.unit), peer)
                    print(
                        "mock udpstreamer: CONNECT from %s:%d, CONFIG sent" % peer,
                        flush=True,
                    )
                elif magic == MAGIC and ptype == PKT_DISCONNECT:
                    print("mock udpstreamer: DISCONNECT", flush=True)
                    peer = None
        except socket.timeout:
            pass

        if peer is None:
            continue

        now = time.monotonic()
        if now < next_send:
            time.sleep(min(next_send - now, 0.002))
            continue

        # The jitter varies each *packet's* own spacing rather than offsetting its start, because
        # the stream's timestamps have to stay strictly increasing -- a jittered t0 can overlap the
        # packet before it, and the buffer, the trigger and every binary search here assume a
        # monotonic clock. The mean spacing is still `dt`, so the declared rate stays honest; what
        # moves is the *local* spacing, which is what a uniform grid over the window cannot model.
        pdt = dt
        if args.clock_jitter > 0.0:
            pdt = dt * (1.0 + random.uniform(-args.clock_jitter, args.clock_jitter))
        spike = 0.0
        if (
            args.time_glitch_at >= 0.0
            and not glitched
            and t_stream >= args.time_glitch_at
        ):
            spike = args.time_glitch_by
            glitched = True
            print(
                "mock udpstreamer: time-signal glitch at t=%.3f, +%.3f s for one packet"
                % (t_stream, spike),
                flush=True,
            )
        pkt = build_data(
            counter,
            t_stream,
            pdt,
            args.batch,
            args.freq,
            time_mode,
            args.noise,
            spike,
            args.glitch_at,
            args.glitch_width,
            args.quiet_at,
        )
        burst = max(1, args.drop_burst)
        if args.drop_every > 0 and drop_left == 0 and counter % args.drop_every == 0:
            drop_left = burst
        if drop_left > 0:
            drop_left -= 1
            dropped += 1
            if not noted_drop:
                noted_drop = True
                print(
                    "mock udpstreamer: dropping packet %d and the %d after it"
                    % (counter, burst - 1),
                    flush=True,
                )
        elif (
            args.reorder_every > 0
            and held is None
            and counter % args.reorder_every == 0
        ):
            held = (counter, pkt, now + args.reorder_delay_ms / 1000.0)
        else:
            sock.sendto(pkt, peer)
            if held is not None and now >= held[2]:
                # Within the receiver's reorder window, this packet is put back in order and
                # costs nothing. Past the window, the receiver has already committed a hole
                # and this packet is what fills it.
                sock.sendto(held[1], peer)
                inverted += 1
                if not noted_reorder:
                    noted_reorder = True
                    print(
                        "mock udpstreamer: sent packet %d before %d, releasing it %.1f ms late"
                        % (counter, held[0], args.reorder_delay_ms),
                        flush=True,
                    )
                held = None
        # The counter and the stream clock advance whatever became of the datagram. That
        # is what makes a withheld packet a *lost* one and not a slower producer: nothing
        # in a payload says which packet it is, only the header's counter does.
        counter += 1
        t_stream += args.batch * pdt
        next_send += period
        if (
            args.duration > 0.0
            and started is not None
            and now - started > args.duration
        ):
            print(
                "mock udpstreamer: done (%d packets, %d dropped, %d inverted)"
                % (counter, dropped, inverted),
                flush=True,
            )
            return


if __name__ == "__main__":
    main()