#!/usr/bin/env python3
"""
Inkpilot Bridge: lets the Inkpilot iPhone app drive a USB pen plotter over Wi-Fi.

iPhones can't talk to USB serial devices, so this tiny program runs on any
computer the plotter is plugged into (a Mac, a Raspberry Pi, an old laptop).
It forwards bytes between the app and the plotter and announces itself on
the network so the app finds it automatically. The app does all the thinking;
this is just a wire.

Works with AxiDraw (EiBotBoard) and GRBL plotters. Needs only Python 3.8+,
nothing to install.

    python3 inkpilot_bridge.py                # find the plotter and start
    python3 inkpilot_bridge.py --port /dev/cu.usbmodem1101
    python3 inkpilot_bridge.py --list         # show serial ports
"""
import argparse
import glob
import json
import os
import platform
import select
import shutil
import signal
import socket
import subprocess
import sys
import termios
import time
import tty

VERSION = "1.1"
DEFAULT_TCP_PORT = 7878
SERVICE_TYPE = "_inkpilot._tcp"


def log(msg):
    print(time.strftime("%H:%M:%S"), msg, flush=True)


# ---------------------------------------------------------------- serial

def candidate_ports():
    pats = ["/dev/cu.usbmodem*", "/dev/cu.usbserial*", "/dev/cu.wchusbserial*", "/dev/cu.SLAB_USBtoUART*",
            "/dev/ttyACM*", "/dev/ttyUSB*"]
    out = []
    for p in pats:
        out.extend(sorted(glob.glob(p)))
    return out


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


class Serial:
    def __init__(self, path, baud=115200):
        self.path = path
        self.fd = os.open(path, os.O_RDWR | os.O_NOCTTY | os.O_NONBLOCK)
        attrs = termios.tcgetattr(self.fd)
        tty.setraw(self.fd)
        attrs = termios.tcgetattr(self.fd)
        attrs[2] |= termios.CLOCAL | termios.CREAD  # cflag
        speed = BAUD.get(baud, termios.B115200)
        attrs[4] = speed
        attrs[5] = speed
        termios.tcsetattr(self.fd, termios.TCSANOW, attrs)
        termios.tcflush(self.fd, termios.TCIOFLUSH)

    def write(self, data):
        view = memoryview(data)
        while view:
            try:
                n = os.write(self.fd, view)
                view = view[n:]
            except BlockingIOError:
                select.select([], [self.fd], [], 1.0)

    def read(self):
        try:
            return os.read(self.fd, 4096)
        except BlockingIOError:
            return b""

    def read_for(self, seconds):
        end = time.time() + seconds
        buf = b""
        while time.time() < end:
            r, _, _ = select.select([self.fd], [], [], max(0, end - time.time()))
            if r:
                chunk = self.read()
                if not chunk:
                    break
                buf += chunk
        return buf

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


def identify(path):
    """Returns ('ebb'|'grbl', banner) or (None, '')."""
    try:
        s = Serial(path)
    except OSError as e:
        return None, str(e)
    try:
        s.write(b"V\r")
        reply = s.read_for(0.8).decode("ascii", "replace")
        if "EBB" in reply or "EB Firmware" in reply:
            return "ebb", reply.strip()
        # GRBL prints a banner after reset and answers $I
        s.write(b"\r\n$I\r\n")
        reply += s.read_for(1.5).decode("ascii", "replace")
        if "Grbl" in reply or "GRBL" in reply or "FluidNC" in reply or "[VER" in reply:
            return "grbl", reply.strip().splitlines()[0] if reply.strip() else "GRBL"
        return None, reply.strip()
    finally:
        s.close()


def find_plotter(explicit=None):
    ports = [explicit] if explicit else candidate_ports()
    for p in ports:
        kind, banner = identify(p)
        if kind:
            return p, kind, banner
        if explicit:
            log(f"{p} did not answer like a plotter ({banner or 'no reply'}); using it anyway as AxiDraw")
            return p, "ebb", banner
    return None, None, ""


# ---------------------------------------------------------------- bonjour

class Advertiser:
    def __init__(self, name, port, kind):
        self.proc = None
        txt = [f"kind={kind}", f"v={VERSION}"]
        if platform.system() == "Darwin" and shutil.which("dns-sd"):
            cmd = ["dns-sd", "-R", name, SERVICE_TYPE, "local", str(port)] + txt
        elif shutil.which("avahi-publish-service"):
            cmd = ["avahi-publish-service", name, SERVICE_TYPE, str(port)] + txt
        else:
            cmd = None
            log("Bonjour not available here; in the app, add the bridge by IP address.")
        if cmd:
            self.proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            log(f"Announced as \"{name}\" so the app can find it.")

    def stop(self):
        if self.proc:
            self.proc.terminate()


def local_ips():
    ips = set()
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("10.255.255.255", 1))
        ips.add(s.getsockname()[0])
        s.close()
    except OSError:
        pass
    return sorted(ips)


class StatusFile:
    """Writes what the bridge is doing to a small JSON file (the Inkpilot Wi-Fi box reads it)."""

    def __init__(self, path):
        self.path = path
        self.state = {"version": VERSION, "plotter": False, "kind": None, "banner": "", "port": None,
                      "app": False, "appAddress": None}
        self.write()

    def update(self, **kw):
        self.state.update(kw)
        self.write()

    def write(self):
        if not self.path:
            return
        try:
            os.makedirs(os.path.dirname(self.path), exist_ok=True)
            tmp = self.path + ".tmp"
            with open(tmp, "w") as f:
                json.dump(dict(self.state, updated=time.time()), f)
            os.replace(tmp, self.path)
        except OSError as e:
            log(f"Couldn't write status file: {e}")


# ---------------------------------------------------------------- main loop

def serve(args):
    status = StatusFile(args.status_file)
    path, kind, banner = find_plotter(args.port)
    if not path:
        log("Waiting for a plotter. Plug in the AxiDraw's USB cable and turn on its power.")
    while not path:
        time.sleep(3)
        path, kind, banner = find_plotter(args.port)
    log(f"Found {'AxiDraw' if kind == 'ebb' else 'GRBL plotter'} on {path}: {banner}")
    status.update(plotter=True, kind=kind, banner=banner, port=path)

    name = args.name or f"{'AxiDraw' if kind == 'ebb' else 'Plotter'} on {socket.gethostname().split('.')[0]}"
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(("0.0.0.0", args.tcp_port))
    server.listen(2)
    adv = Advertiser(name, args.tcp_port, kind) if not args.no_advertise else None
    for ip in local_ips():
        log(f"Ready. In the app: Plotter tab > Connect. Or Add by IP address: {ip}, port {args.tcp_port}")

    serial = Serial(path)
    client = None

    def shutdown(*_):
        if adv:
            adv.stop()
        status.update(plotter=False, app=False, appAddress=None)
        if serial:
            serial.close()
        sys.exit(0)

    signal.signal(signal.SIGINT, shutdown)
    signal.signal(signal.SIGTERM, shutdown)

    while True:
        rlist = [server, serial.fd] + ([client] if client else [])
        try:
            ready, _, _ = select.select(rlist, [], [], 1.0)
        except (OSError, ValueError):
            ready = []
        for r in ready:
            if r is server:
                conn, addr = server.accept()
                conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                if client:
                    log("A new app connected; closing the previous connection.")
                    client.close()
                client = conn
                log(f"App connected from {addr[0]}")
                status.update(app=True, appAddress=addr[0])
            elif r == serial.fd:
                data = serial.read()
                if data and client:
                    try:
                        client.sendall(data)
                    except OSError:
                        client = None
                        status.update(app=False, appAddress=None)
            elif client and r is client:
                try:
                    data = client.recv(4096)
                except OSError:
                    data = b""
                if not data:
                    log("App disconnected.")
                    client.close()
                    client = None
                    status.update(app=False, appAddress=None)
                    if kind == "ebb":
                        serial.write(b"SP,1\r")  # lift the pen so it can't bleed into the paper
                    continue
                serial.write(data)
        # plotter unplugged?
        if not os.path.exists(path):
            log("Plotter disconnected. Waiting for it to come back...")
            serial.close()
            if client:
                client.close()
                client = None
            status.update(plotter=False, app=False, appAddress=None)
            while not os.path.exists(path):
                time.sleep(1)
            time.sleep(1)
            serial = Serial(path)
            status.update(plotter=True)
            log("Plotter is back.")


def main():
    ap = argparse.ArgumentParser(description="Share a USB pen plotter with the Inkpilot app over Wi-Fi.")
    ap.add_argument("--port", help="serial port (default: find it)")
    ap.add_argument("--tcp-port", type=int, default=DEFAULT_TCP_PORT)
    ap.add_argument("--name", help="name shown in the app")
    ap.add_argument("--list", action="store_true", help="list serial ports and what's on them")
    ap.add_argument("--no-advertise", action="store_true",
                    help="don't announce on the network (the Inkpilot Wi-Fi box announces itself)")
    ap.add_argument("--status-file", help="write the bridge's state to this JSON file")
    args = ap.parse_args()
    if args.list:
        for p in candidate_ports():
            kind, banner = identify(p)
            print(f"{p}: {kind or 'unknown'} {banner}")
        return
    log(f"Inkpilot Bridge {VERSION}")
    serve(args)


if __name__ == "__main__":
    main()