#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: 2026 Martino Ferrari
"""A protocol mock of the Measurement Computing E-1608 Ethernet DAQ, for exercising udpscope
without a board. It is the *contract* the `e1608` source is verified against until a board is
attached, so its layouts are only as good as the two references they were taken from:

  * MARTe2 `E1608DataSource` (~/Projects/marte2_projects/MCC_E1608_DataSource), the primary, and
  * the Go `uopi/internal/datasource/mcc` implementation, the cross-check.

    ./tools/mock_e1608.py --port 54211 [--code N] [--drop-after S] [--overrun-after S] [--log F]

Serves: UDP `D` discovery and `C` connection-code handshake on --port; the TCP command channel on
--port; the TCP scan stream on --port+1 (it starts when the one filler byte follows the scan-start
command). Frame: 0xDB, cmd, frame id, status, count LE16, payload, checksum (the whole frame sums
to 0xFF mod 256). Reply: cmd | 0x80, same id.

UNVERIFIED -- every item below rests on a single reference, or the references disagree. They are
the first places to look when a real board misbehaves:

  * 0x31 counter reset: Go only (no payload, no reply payload).
  * Pacer: MARTe2 sends round(80e6 / (scan_rate * channels) - 1), the *aggregate* rate; Go sends the
    passed frequency. This mock follows MARTe2.
  * The scan-start options byte (trigger): none 0, rising 4, falling 8, high 12, low 16 -- MARTe2 only.
  * Handshake status 2 ("code ignored") is accepted by MARTe2 and refused by Go. The mock never sends it.
  * The first frame id: MARTe2 pre-increments (first frame is 1), Go starts at 0. The mock echoes it.
  * The maximum frame size (MARTe2: 1024 in total; Go: 1024 of payload). The mock enforces neither.
  * Agreed by both references (not in doubt): the filler byte after scan start, the scan port being
    command port + 1, DIFF wire channel = channel + 8.

Switches: --drop-after S closes both connections once, S seconds after the first command connection
(logs DROP); --overrun-after S raises the overrun bit and stops the stream S seconds after a scan
starts, once (logs OVERRUN); --code N is the connection code the board accepts (anything else gets
status 1).
"""

import argparse
import socket
import struct
import sys
import threading
import time

import numpy as np

CLOCK = 80e6
PRODUCT_ID = 0x012F
FS = [10.0, 5.0, 2.0, 1.0]  # full scale per range code BIP10V, BIP5V, BIP2V, BIP1V


def f32(x):
    return struct.unpack("<f", struct.pack("<f", x))[0]


DF_CAL = [(f32(1.01 - 0.004 * r), f32(150.0 - 30.0 * r)) for r in range(4)]
SE_CAL = [(f32(1.02 - 0.005 * r), f32(300.0 - 40.0 * r)) for r in range(4)]
AO_CAL = [(f32(0.99 + 0.01 * c), f32(200.0 - 100.0 * c)) for c in range(2)]
CAL_MEM = b"".join(struct.pack("<ff", s, i) for s, i in DF_CAL + SE_CAL + AO_CAL)  # 80 bytes


def close_quietly(sock):
    if sock is None:
        return
    try:
        sock.shutdown(socket.SHUT_RDWR)
    except OSError:
        pass
    try:
        sock.close()
    except OSError:
        pass


def recv_exact(conn, n):
    buf = b""
    while len(buf) < n:
        chunk = conn.recv(n - len(buf))
        if not chunk:
            return None
        buf += chunk
    return buf


class Board:
    def __init__(self, args):
        self.args = args
        self.mac = bytes(int(x, 16) for x in args.mac.split(":"))
        self.lock = threading.RLock()
        self.cmd_conn = None
        self.scan_conn = None
        self.dout = 0
        self.dconf = 0xFF
        self.ao = [32768, 32768]
        self.queue = []
        self.pacer = 0
        self.armed = False
        self.streaming = False
        self.stream_gen = 0
        self.overrun = False
        self.overrun_done = False
        self.t_boot = time.monotonic()
        self.counter_t0 = self.t_boot
        self.first_conn = None
        self.logfile = open(args.log, "a", buffering=1) if args.log else None

    def log(self, text):
        print("[mock] " + text, flush=True)
        if self.logfile:
            self.logfile.write(text + "\n")

    # --- the pretend hardware -------------------------------------------------------------
    def din(self):
        pins = int((time.monotonic() - self.t_boot) * 4) & 0xFF
        return (pins & self.dconf) | (self.dout & ~self.dconf & 0xFF)

    def counter(self):
        return int((time.monotonic() - self.counter_t0) * 5) & 0xFFFFFFFF

    def raw(self, wire, rng, t):
        diff = wire >= 8
        ch = wire - 8 if diff else wire
        fs = FS[rng]
        dc = min(0.5 * (ch + 1), 0.8 * fs)
        v = np.clip(dc + 0.1 * np.sin(2.0 * np.pi * 5.0 * t), -0.98 * fs, 0.98 * fs)
        code = v * 32768.0 / fs + 32768.0
        slope, icpt = (DF_CAL if diff else SE_CAL)[rng]
        return np.clip(np.rint((code - icpt) / slope), 0, 65535)

    # --- commands ---------------------------------------------------------------------------
    def handle(self, cmd, p):
        """Returns (status, reply payload). Status 1 protocol, 2 parameter."""
        if cmd == 0x00:
            return 0, bytes([self.din()])
        if cmd == 0x02:
            return 0, bytes([self.dout])
        if cmd == 0x03:
            if len(p) < 1:
                return 2, b""
            self.dout = p[0]
            self.log("DOUT 0x%02x" % p[0])
            return 0, b""
        if cmd == 0x05:
            if len(p) < 1:
                return 2, b""
            self.dconf = p[0]
            self.log("DCONF 0x%02x" % p[0])
            return 0, b""
        if cmd == 0x10:
            return 0, struct.pack("<H", 32768)
        if cmd == 0x11:
            if len(p) < 9 or not self.queue:
                return 2, b""
            _count, pacer, opts = struct.unpack("<IIB", p[:9])
            with self.lock:
                self.pacer = pacer
                self.armed = True
                self.overrun = False
            self.log("SCAN_START pacer=%d options=%d channels=%d" % (pacer, opts, len(self.queue)))
            return 0, b""
        if cmd == 0x13:
            self.stop_stream(close_scan=bool(p and p[0]))
            self.log("SCAN_STOP")
            return 0, b""
        if cmd == 0x15:
            if not p:
                return 2, b""
            n = p[0]
            if n < 1 or n > 8 or len(p) != 1 + 2 * n:
                return 2, b""
            q = [(p[1 + 2 * i], p[2 + 2 * i]) for i in range(n)]
            if any(r > 3 for _, r in q):
                return 2, b""
            self.queue = q
            self.log("QUEUE " + " ".join("%d:%d" % e for e in q))
            return 0, b""
        if cmd == 0x20:
            return 0, struct.pack("<HH", *self.ao)
        if cmd == 0x21:
            if len(p) < 3 or p[0] > 1:
                return 2, b""
            ch, code = struct.unpack("<BH", p[:3])
            self.ao[ch] = code
            slope, icpt = AO_CAL[ch]
            volts = ((code - icpt) / slope - 32768.0) * 10.0 / 32768.0
            self.log("AOUT ch=%d code=%d volts=%.4f" % (ch, code, volts))
            return 0, b""
        if cmd == 0x30:
            return 0, struct.pack("<I", self.counter())
        if cmd == 0x31:
            self.counter_t0 = time.monotonic()
            self.log("COUNTER_RESET")
            return 0, b""
        if cmd == 0x40:
            if len(p) < 4:
                return 2, b""
            addr, count = struct.unpack("<HH", p[:4])
            data = CAL_MEM[addr:addr + count]
            if len(data) != count:
                return 2, b""
            return 0, data
        if cmd == 0x50:
            self.log("BLINK")
            return 0, b""
        if cmd == 0x52:
            st = 0
            if self.scan_conn is not None:
                st |= 1
            if self.streaming:
                st |= 2
            if self.overrun:
                st |= 4
            return 0, struct.pack("<H", st)
        return 1, b""

    # --- streaming --------------------------------------------------------------------------
    def stop_stream(self, close_scan):
        with self.lock:
            self.stream_gen += 1
            self.streaming = False
            self.armed = False
            conn = None
            if close_scan:
                conn = self.scan_conn
                self.scan_conn = None
        close_quietly(conn)

    def start_stream(self):
        for _ in range(100):  # the scan socket's accept may trail the filler by a moment
            with self.lock:
                if self.scan_conn is not None:
                    break
            time.sleep(0.01)
        with self.lock:
            conn = self.scan_conn
            if conn is None or not self.armed:
                return
            self.armed = False
            self.streaming = True
            self.stream_gen += 1
            gen = self.stream_gen
            queue = list(self.queue)
            pacer = self.pacer
        threading.Thread(target=self.stream, args=(gen, conn, queue, pacer), daemon=True).start()

    def stream(self, gen, conn, queue, pacer):
        nq = len(queue)
        agg = CLOCK / (pacer + 1) if pacer > 0 else 1000.0 * nq  # external clock: 1 kS/s each
        rate = agg / nq
        t0 = time.monotonic()
        sent = 0
        ov_at = None
        if self.args.overrun_after > 0 and not self.overrun_done:
            ov_at = t0 + self.args.overrun_after
        self.log("STREAM rate=%.1f channels=%d" % (rate, nq))
        try:
            while self.stream_gen == gen:
                now = time.monotonic()
                if ov_at is not None and now >= ov_at:
                    with self.lock:
                        if self.stream_gen == gen:
                            self.overrun = True
                            self.overrun_done = True
                            self.streaming = False
                    self.log("OVERRUN")
                    return
                due = min(int((now - t0) * rate) - sent, 65536)
                if due > 0:
                    t = np.arange(sent, sent + due) / rate
                    cols = [self.raw(w, r, t) for (w, r) in queue]
                    conn.sendall(np.stack(cols, axis=1).astype("<u2").tobytes())
                    sent += due
                time.sleep(0.005)
        except (OSError, ValueError):
            pass
        finally:
            with self.lock:
                if self.stream_gen == gen:
                    self.streaming = False

    def drop(self):
        self.log("DROP")
        with self.lock:
            self.stream_gen += 1
            self.streaming = False
            self.armed = False
            conns = [self.cmd_conn, self.scan_conn]
            self.cmd_conn = None
            self.scan_conn = None
        for c in conns:
            close_quietly(c)

    # --- servers ----------------------------------------------------------------------------
    def serve_udp(self, port):
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(("", port))
        while True:
            data, addr = sock.recvfrom(64)
            if data[:1] == b"D":
                reply = b"D" + self.mac + struct.pack("<HH", PRODUCT_ID, 0x0100)
                reply += b"E-1608".ljust(16, b"\0") + struct.pack("<H", port)
                sock.sendto(reply.ljust(64, b"\0"), addr)
            elif data[:1] == b"C" and len(data) >= 5:
                (code,) = struct.unpack("<I", data[1:5])
                status = 0 if code == self.args.code else 1
                self.log("HANDSHAKE code=%d status=%d" % (code, status))
                sock.sendto(b"C" + bytes([status]), addr)

    def serve_scan_port(self, port):
        srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        srv.bind(("", port))
        srv.listen(4)
        while True:
            conn, _ = srv.accept()
            with self.lock:
                old = self.scan_conn
                self.scan_conn = conn
            close_quietly(old)

    def serve_cmd(self, conn):
        with self.lock:
            old = self.cmd_conn
            self.cmd_conn = conn
            first = self.first_conn is None
            if first:
                self.first_conn = time.monotonic()
        close_quietly(old)
        if first and self.args.drop_after > 0:
            threading.Timer(self.args.drop_after, self.drop).start()
        try:
            while True:
                b = conn.recv(1)
                if not b:
                    break
                if b[0] == 0x00:
                    self.start_stream()  # the filler byte after scan start
                    continue
                if b[0] != 0xDB:
                    continue
                rest = recv_exact(conn, 5)
                if rest is None:
                    break
                cmd, fid, _status, count = struct.unpack("<BBBH", rest)
                body = recv_exact(conn, count + 1)
                if body is None:
                    break
                frame = b"\xDB" + rest + body
                if (sum(frame) & 0xFF) == 0xFF:
                    status, reply = self.handle(cmd, body[:count])
                else:
                    status, reply = 1, b""
                out = struct.pack("<BBBBH", 0xDB, cmd | 0x80, fid, status, len(reply)) + reply
                out += bytes([(0xFF - sum(out)) & 0xFF])
                conn.sendall(out)
        except (OSError, ValueError):
            pass
        finally:
            with self.lock:
                mine = self.cmd_conn is conn
                if mine:
                    self.cmd_conn = None
            if mine:
                self.stop_stream(close_scan=True)
            close_quietly(conn)

    def serve_cmd_port(self, port):
        srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        srv.bind(("", port))
        srv.listen(4)
        while True:
            conn, _ = srv.accept()
            conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            threading.Thread(target=self.serve_cmd, args=(conn,), daemon=True).start()


def main():
    ap = argparse.ArgumentParser(description="E-1608 protocol mock")
    ap.add_argument("--port", type=int, default=54211)
    ap.add_argument("--mac", default="00:80:2f:00:16:08")
    ap.add_argument("--code", type=int, default=0, help="connection code the board accepts")
    ap.add_argument("--drop-after", type=float, default=0.0, help="close both links once, S s after the first command connection")
    ap.add_argument("--overrun-after", type=float, default=0.0, help="raise the overrun bit S s after a scan starts, once")
    ap.add_argument("--log", default="", help="append event lines to this file")
    args = ap.parse_args()
    board = Board(args)
    for target in (board.serve_udp, board.serve_scan_port, board.serve_cmd_port):
        port = args.port + 1 if target == board.serve_scan_port else args.port
        threading.Thread(target=target, args=(port,), daemon=True).start()
    print("mock_e1608: listening on port %d (scan %d)" % (args.port, args.port + 1), flush=True)
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        sys.exit(0)


if __name__ == "__main__":
    main()