#!/usr/bin/env python3
"""
mm_link.py — a "perfect modem" between THOR and a Micromate, with a readable
log and deliberate fault injection.

Why
---
THOR gives almost no visibility into a connection: a refresh button, two poll
intervals, and no way to see whether a check succeeded, timed out, or was never
sent.  When a unit "won't stay connected" there is nothing to look at.

This sits where the cellular modem would sit and answers the question directly:

  * **What is THOR actually doing?**  Every frame is decoded and timestamped —
    `POLL`, `MONITOR_STATUS`, `SETUP_NAME_READ` — not a hex dump.
  * **Is it even trying?**  Silence is visible: the log shows gaps.
  * **How does it behave when the link misbehaves?**  Faults can be injected on
    demand, which a real cell link will not do on cue.

Point THOR at this host and port exactly as if it were a modem (Communication:
TCP, IP: <this host>, Port: <--listen>).

Fault injection
---------------
Write a mode into the control file (default `mm_link.ctl`) and it takes effect
on the next byte:

    echo pass      > mm_link.ctl    # normal relay
    echo blackhole > mm_link.ctl    # TCP stays up, bytes are swallowed
    echo drop      > mm_link.ctl    # close the connection abruptly (RST-ish)
    echo delay:2.0 > mm_link.ctl    # forward, but 2 s late in both directions
    echo onewaydev > mm_link.ctl    # THOR->unit passes, unit->THOR is swallowed

**`blackhole` is the one that matters.**  It reproduces the classic cellular
failure: the socket is still open as far as both ends are concerned, but nothing
crosses.  A client that relies on TCP to tell it the peer is gone will sit there
until the OS keepalive fires — which by default is about two hours.

Usage
-----
    python3 bridges/mm_link.py --serial /dev/ttyACM0 --baud 115200 \\
        --listen 12345 --logdir ~/mm-captures

Writes, per session:
    <logdir>/mmlink_<ts>/session.log   decoded, timestamped, human-readable
    <logdir>/mmlink_<ts>/raw_bw.bin    THOR -> unit, raw
    <logdir>/mmlink_<ts>/raw_s3.bin    unit -> THOR, raw

The raw pair loads straight into `scratch/mm_frame_parse.py`.
"""

from __future__ import annotations

import argparse
import datetime
import os
import socket
import sys
import threading
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "scratch"))
try:
    from mm_frame_parse import SUBNAME, destuff  # noqa: F401
except Exception:                                 # pragma: no cover
    SUBNAME = {}

import errno
import select
import termios

DLE, STX, ETX, ACK = 0x10, 0x02, 0x03, 0x41

_BAUD = {9600: termios.B9600, 19200: termios.B19200, 38400: termios.B38400,
         57600: termios.B57600, 115200: termios.B115200}


class SerialPort:
    """Minimal raw serial port on stdlib termios — no pyserial dependency.

    The bench hosts are whatever is to hand; requiring a pip install on someone
    else's machine is a poor trade for the ~30 lines this saves.
    """

    def __init__(self, path: str, baud: int):
        if baud not in _BAUD:
            raise ValueError(f"unsupported baud {baud}; pick one of {sorted(_BAUD)}")
        self.fd = os.open(path, os.O_RDWR | os.O_NOCTTY | os.O_NONBLOCK)
        a = termios.tcgetattr(self.fd)
        a[0] = 0                                             # iflag: no translation
        a[1] = 0                                             # oflag: raw
        a[2] = termios.CS8 | termios.CREAD | termios.CLOCAL  # cflag: 8N1, ignore modem lines
        a[3] = 0                                             # lflag: non-canonical, no echo
        a[4] = a[5] = _BAUD[baud]
        a[6] = list(a[6])
        a[6][termios.VMIN] = 0
        a[6][termios.VTIME] = 0
        termios.tcsetattr(self.fd, termios.TCSANOW, a)
        termios.tcflush(self.fd, termios.TCIOFLUSH)

    def read(self, n: int) -> bytes:
        r, _, _ = select.select([self.fd], [], [], 0.2)
        if not r:
            return b""
        try:
            return os.read(self.fd, n)
        except OSError as e:
            if e.errno in (errno.EAGAIN, errno.EWOULDBLOCK):
                return b""
            raise

    def write(self, data: bytes) -> None:
        while data:
            try:
                data = data[os.write(self.fd, data):]
            except OSError as e:
                if e.errno in (errno.EAGAIN, errno.EWOULDBLOCK):
                    select.select([], [self.fd], [], 0.2)
                    continue
                raise

    def close(self) -> None:
        try:
            os.close(self.fd)
        except OSError:
            pass


def name_of(sub: int, is_request: bool) -> str:
    if is_request:
        return SUBNAME.get(sub, f"SUB_{sub:02X}")
    return "rsp " + SUBNAME.get(0xFF - sub, f"SUB_{0xFF - sub:02X}")


class FrameSniffer:
    """Accumulate bytes and report complete frames, without altering the stream."""

    def __init__(self, is_request: bool):
        self.is_request = is_request
        self.buf = bytearray()

    def feed(self, data: bytes):
        """Yield (sub, payload_len) for each complete frame seen."""
        self.buf.extend(data)
        while True:
            start = -1
            for i, b in enumerate(self.buf):
                if self.is_request and b == ACK and i + 1 < len(self.buf) and self.buf[i + 1] == STX:
                    start = i
                    break
                if not self.is_request and b == STX:
                    start = i
                    break
            if start < 0:
                if len(self.buf) > 8192:
                    del self.buf[:-16]
                return
            j = start + (2 if self.is_request else 1)
            end = -1
            while j < len(self.buf):
                if self.buf[j] == DLE and j + 1 < len(self.buf):
                    j += 2
                    continue
                if self.buf[j] == ETX:
                    end = j
                    break
                j += 1
            if end < 0:
                return                      # wait for more bytes
            body = self.buf[start:end + 1]
            del self.buf[:end + 1]
            # SUB sits at a fixed spot past the leading framing -- but it is
            # DLE-escaped when its own value is 0x02/0x03/0x04/0x10, so a raw
            # read reports 0x10 for those.  SUB 0x02 was being logged as
            # "SUB_10" until this was handled.
            off = 5 if self.is_request else 3
            if len(body) > off:
                sub = body[off]
                if sub == DLE and len(body) > off + 1:
                    sub = body[off + 1]
                yield sub, len(body)


class Link:
    def __init__(self, args):
        self.args = args
        self.mode = "pass"
        self.delay = 0.0
        self.ctl = Path(args.control)
        self.session: Path | None = None
        self.log_fh = None
        self.raw = {}
        self.t0 = time.time()
        self.counts = {}

    # ── logging ────────────────────────────────────────────────────────────
    def open_session(self):
        ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        self.session = Path(self.args.logdir) / f"mmlink_{ts}"
        self.session.mkdir(parents=True, exist_ok=True)
        self.log_fh = open(self.session / "session.log", "a", buffering=1)
        self.raw = {
            "bw": open(self.session / "raw_bw.bin", "ab"),
            "s3": open(self.session / "raw_s3.bin", "ab"),
        }
        self.say(f"=== session {ts} — serial {self.args.serial} @ {self.args.baud} ===")

    def say(self, text: str):
        line = f"{datetime.datetime.now().strftime('%H:%M:%S.%f')[:-3]}  {text}"
        print(line, flush=True)
        if self.log_fh:
            self.log_fh.write(line + "\n")

    # ── control file ───────────────────────────────────────────────────────
    def poll_control(self):
        while True:
            try:
                if self.ctl.exists():
                    want = self.ctl.read_text().strip().lower()
                    if want.startswith("delay:"):
                        d = float(want.split(":", 1)[1])
                        if ("delay", d) != (self.mode, self.delay):
                            self.mode, self.delay = "delay", d
                            self.say(f"*** MODE -> delay {d}s ***")
                    elif want and want != self.mode:
                        self.mode, self.delay = want, 0.0
                        self.say(f"*** MODE -> {want} ***")
            except Exception:
                pass
            time.sleep(0.25)

    # ── the relay ──────────────────────────────────────────────────────────
    def pump(self, src, dst, tag: str, is_request: bool, stop: threading.Event):
        sniff = FrameSniffer(is_request)
        arrow = "THOR->unit" if is_request else "unit->THOR"
        last = time.time()
        while not stop.is_set():
            timed_out = False
            try:
                data = src.recv(4096) if isinstance(src, socket.socket) else src.read(4096)
            except TimeoutError:
                timed_out = True
                # socket.timeout subclasses OSError, so it MUST be caught first.
                # Treating it as a dead socket closes the connection after 200 ms
                # of quiet -- which is exactly what `blackhole` produces, so the
                # relay killed the link it was supposed to be faking a fault on.
                data = b""
            except OSError:
                break
            if isinstance(src, socket.socket) and data == b"" and not timed_out:
                self.say(f"{arrow}: peer closed the connection")
                break
            if not data:
                if time.time() - last > self.args.quiet_after and self.counts:
                    self.say(f"--- {self.args.quiet_after:.0f}s with no traffic ---")
                    last = time.time()
                continue
            last = time.time()

            self.raw[tag].write(data)
            self.raw[tag].flush()
            for sub, ln in sniff.feed(data):
                label = name_of(sub, is_request)
                self.counts[label] = self.counts.get(label, 0) + 1
                self.say(f"{arrow}  {label:<20} ({ln} B)"
                         + ("" if self.mode == "pass" else f"   [mode={self.mode}]"))

            mode = self.mode
            if mode == "drop":
                self.say(f"{arrow}: DROPPING the connection (fault injection)")
                stop.set()
                break
            if mode == "blackhole":
                continue                      # swallow, keep the socket open
            if mode == "onewaydev" and not is_request:
                continue                      # unit's replies never reach THOR
            if mode == "delay" and self.delay:
                time.sleep(self.delay)
            try:
                if isinstance(dst, socket.socket):
                    dst.sendall(data)
                else:
                    dst.write(data)
            except OSError:
                break
        stop.set()

    def serve(self):
        srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        srv.bind(("0.0.0.0", self.args.listen))
        srv.listen(5)
        self.open_session()
        self.say(f"listening on 0.0.0.0:{self.args.listen}   control file: {self.ctl}")
        self.say("point THOR at this host/port as Communication=TCP")
        threading.Thread(target=self.poll_control, daemon=True).start()

        while True:
            conn, addr = srv.accept()
            conn.settimeout(0.2)
            self.say(f"+++ THOR connected from {addr[0]}:{addr[1]} +++")
            try:
                ser = SerialPort(self.args.serial, self.args.baud)
            except OSError as e:
                self.say(f"!!! cannot open {self.args.serial}: {e}")
                conn.close()
                continue
            stop = threading.Event()
            ts = [
                threading.Thread(target=self.pump, args=(conn, ser, "bw", True, stop), daemon=True),
                threading.Thread(target=self.pump, args=(ser, conn, "s3", False, stop), daemon=True),
            ]
            for t in ts:
                t.start()
            for t in ts:
                t.join()
            conn.close()
            ser.close()
            summary = ", ".join(f"{k}x{v}" for k, v in sorted(self.counts.items()))
            self.say(f"--- connection closed.  frames this session: {summary or 'none'} ---")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--serial", default="/dev/ttyACM0")
    ap.add_argument("--baud", type=int, default=115200)
    ap.add_argument("--listen", type=int, default=12345)
    ap.add_argument("--logdir", default=os.path.expanduser("~/mm-captures"))
    ap.add_argument("--control", default="mm_link.ctl")
    ap.add_argument("--quiet-after", type=float, default=30.0,
                    help="log a marker after this many seconds of silence")
    Link(ap.parse_args()).serve()


if __name__ == "__main__":
    main()