#!/usr/bin/env python3
# Reusable bench probe (post-abort recovery timeline; lives in MDK-ARM/bench_diag/).
# RECOVERY TIMELINE after an aborted transfer: how long is the device impaired, and is
# the impairment specific to the FTP application (lftpd) or common to the net/chip layer?
#
# Trigger (t=0):
#   --target ftp : raw control login -> PASV -> data conn -> STOR -> 8192B burst -> RST the
#                  DATA channel; control channel teardown is either
#                    clean = 0.5s settle, read the reply, QUIT   (what real clients do)
#                    fast  = QUIT ~20ms after the abort           (events coincide)
#   --target sink: :7003 connect -> INIT -> LCDMEM -> N keepalive frames (RX in flight,
#                  unread) -> RST          (NO lftpd code involved at all)
#
# Then every --step ms a full short session is attempted until it succeeds --first-ok
# marks the moment the device is usable again-- and every attempt's payload integrity is
# recorded (FTP: NOOP reply; sink: 3209B LCDMEM frame length + CRC, and the echo count of a
# keepalive burst = RX direction integrity).
import argparse, re, socket, struct, sys, time

TAG_CLIENT = 0xAA
CMD_KEEP = 0
CMD_INIT = 1
CMD_LCDMEM = 3
MEMSIZE = 3200
LCDMEM_TOTAL = 3209
ACK_LEN = 5


def log(m):
    sys.stdout.write(m + "\n")
    sys.stdout.flush()


def xor_bytes(b):
    c = 0
    for x in b:
        c ^= x
    return c


def req(cmd, data=b""):
    body = bytes(data)
    return (bytes([TAG_CLIENT, cmd, (len(body) >> 8) & 0xFF, len(body) & 0xFF]) + body
            + bytes([xor_bytes(body)]))


def read_exact(s, want, deadline):
    buf = bytearray()
    while len(buf) < want:
        left = deadline - time.monotonic()
        if left <= 0:
            break
        s.settimeout(left)
        try:
            d = s.recv(want - len(buf))
        except (socket.timeout, OSError):
            break
        if not d:
            break
        buf.extend(d)
    return bytes(buf)


def read_frame(s, deadline):
    """Read exactly one RemoDispBus frame -> (got, total, crc_ok)."""
    buf = bytearray()
    total = None
    while True:
        want = (4 - len(buf)) if total is None else (total - len(buf))
        if want <= 0:
            break
        left = deadline - time.monotonic()
        if left <= 0:
            break
        s.settimeout(left)
        try:
            d = s.recv(min(want, 4096))
        except (socket.timeout, OSError):
            break
        if not d:
            break
        buf.extend(d)
        if total is None and len(buf) >= 4:
            total = 4 + (((buf[2] << 8) | buf[3])) + 1
    if total is None or len(buf) < total:
        return len(buf), (total or 0), False
    body = bytes(buf[4:total - 1])
    return len(buf), total, (buf[total - 1] == xor_bytes(body))


def rst(s):
    try:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
    except OSError:
        pass
    try:
        s.close()
    except OSError:
        pass


# ------------------------------- FTP target -------------------------------

def ftp_recv(s, timeout=2.0):
    s.settimeout(timeout)
    try:
        return s.recv(4096).decode("latin1", "replace")
    except (socket.timeout, OSError):
        return ""


def ftp_trigger(ip, port, teardown, name, send=8192):
    ctl = socket.create_connection((ip, port), timeout=5.0)
    ftp_recv(ctl)
    ctl.sendall(b"USER anonymous\r\n"); ftp_recv(ctl)
    ctl.sendall(b"TYPE I\r\n"); ftp_recv(ctl)
    ctl.sendall(b"PASV\r\n")
    r = ftp_recv(ctl)
    m = re.search(r"\((\d+),(\d+),(\d+),(\d+),(\d+),(\d+)\)", r)
    if not m:
        raise RuntimeError("PASV failed: %r" % r[:60])
    dhost = ".".join(m.group(i) for i in (1, 2, 3, 4))
    dport = (int(m.group(5)) << 8) | int(m.group(6))
    d = socket.create_connection((dhost, dport), timeout=5.0)
    ctl.sendall(("STOR " + name + "\r\n").encode())
    ftp_recv(ctl, 1.0)
    d.sendall(bytes((i * 7 + 3) & 0xFF for i in range(send)))
    t = time.monotonic()
    rst(d)                                   # abort mid-upload
    if teardown == "fast":
        time.sleep(0.02)
        ftp_recv(ctl, 0.2)
    else:
        time.sleep(0.5)
        ftp_recv(ctl, 1.0)
    try:
        ctl.sendall(b"QUIT\r\n")
        ftp_recv(ctl, 0.5)
    except OSError:
        pass
    try:
        ctl.close()
    except OSError:
        pass
    return t


def ftp_probe(ip, port, budget=0.4):
    """One short real session: connect -> NOOP reply integrity -> QUIT."""
    t = time.monotonic()
    try:
        s = socket.create_connection((ip, port), timeout=budget)
    except Exception as e:
        return False, "REFUSED" if isinstance(e, ConnectionRefusedError) else type(e).__name__
    banner = ftp_recv(s, budget)
    if not banner.startswith("220"):
        rst(s)
        return False, "nobanner"
    s.sendall(b"NOOP\r\n")
    r = ftp_recv(s, budget)
    ok = r.startswith("2")
    try:
        s.sendall(b"QUIT\r\n"); ftp_recv(s, 0.3)
    except OSError:
        pass
    try:
        s.close()
    except OSError:
        pass
    return True, ("ok" if ok else "noop=%r" % r[:16])


# ------------------------------- sink target -------------------------------

def sink_trigger(ip, port, burst):
    s = socket.create_connection((ip, port), timeout=5.0)
    s.sendall(req(CMD_INIT))
    read_frame(s, time.monotonic() + 5.0)
    if burst:
        s.sendall(req(CMD_KEEP) * burst)
    s.sendall(req(CMD_LCDMEM, bytes(4)))
    t = time.monotonic()
    rst(s)
    return t


def sink_probe(ip, port, burst=0, budget=2.0):
    """One real session: INIT + keepalive burst echo count + LCDMEM frame integrity."""
    try:
        s = socket.create_connection((ip, port), timeout=budget)
    except Exception as e:
        return False, "REFUSED" if isinstance(e, ConnectionRefusedError) else type(e).__name__
    dl = time.monotonic() + budget
    s.sendall(req(CMD_INIT))
    ini_got, _, ini_crc = read_frame(s, dl)
    ack = -1
    if burst:
        s.sendall(req(CMD_KEEP) * burst)
        ack = len(read_exact(s, burst * ACK_LEN, time.monotonic() + budget))
    s.sendall(req(CMD_LCDMEM, bytes(4)))
    got, total, crc = read_frame(s, time.monotonic() + budget)
    try:
        s.sendall(b"")            # no-op: close gracefully (FIN) so we never re-trigger
        s.shutdown(socket.SHUT_RDWR)
    except OSError:
        pass
    try:
        s.close()
    except OSError:
        pass
    ok = (ini_got == 13 and got == LCDMEM_TOTAL and crc and (not burst or ack == burst * ACK_LEN))
    return True, ("ok" if ok else "init=%d lcd=%d/%d crc=%s ack=%s"
                  % (ini_got, got, LCDMEM_TOTAL, "OK" if crc else "BAD",
                     ack if burst else "-"))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--target", choices=("ftp", "sink"), default="ftp")
    ap.add_argument("--ip", default="192.168.1.100")
    ap.add_argument("--port", type=int, default=0, help="0 = 21 for ftp, 7003 for sink")
    ap.add_argument("--rounds", type=int, default=3)
    ap.add_argument("--span", type=int, default=4000, help="timeline length in ms")
    ap.add_argument("--step", type=int, default=100)
    ap.add_argument("--teardown", choices=("clean", "fast"), default="clean")
    ap.add_argument("--burst", type=int, default=400, help="sink: keepalive frames in flight")
    a = ap.parse_args()
    port = a.port or (21 if a.target == "ftp" else 7003)

    log("== recovery timeline target=%s port=%d teardown=%s burst=%d step=%dms span=%dms"
        % (a.target, port, a.teardown, a.burst, a.step, a.span))
    for r in range(1, a.rounds + 1):
        if a.target == "ftp":
            t0 = ftp_trigger(a.ip, port, a.teardown, "/r%04d.bin" % r)
        else:
            t0 = sink_trigger(a.ip, port, a.burst)
        seq = []
        first_ok = None
        t = 0
        while t < a.span:
            time.sleep(a.step / 1000.0)
            t = int((time.monotonic() - t0) * 1e3)
            if a.target == "ftp":
                up, det = ftp_probe(a.ip, port)
            else:
                up, det = sink_probe(a.ip, port, a.burst)
            seq.append("%d:%s" % (t, det))
            if up:
                if first_ok is None:
                    first_ok = t
                if det == "ok" and first_ok is not None and t > first_ok + 400:
                    break
        log("  round %d first_usable=%s ms" % (r, first_ok if first_ok is not None else ">%d" % a.span))
        log("    " + " | ".join(seq))
        time.sleep(6.0)                       # let the device fully settle before the next trigger
    log("SUMMARY target=%s teardown=%s rounds=%d done" % (a.target, a.teardown, a.rounds))


main()