"""Talk to GRBL 1.1 (Arduino Uno) or FluidNC over USB serial.

Streaming uses GRBL's character-counting method: keep sending lines while the bytes not yet
acknowledged fit in the controller's 128-byte receive buffer. That keeps the planner full, so
the ball doesn't stutter on short moves.
"""
import re
import time

RX_BUFFER = 127
STATUS_RE = re.compile(r'<([^|>]+)(.*)>')

# real-time commands (single bytes, acted on at once, never queued)
RT_STATUS = b'?'
RT_HOLD = b'!'
RT_RESUME = b'~'
RT_RESET = b'\x18'
RT_JOG_CANCEL = b'\x85'


class GrblError(RuntimeError):
    pass


class Aborted(GrblError):
    """A wait was cut short because `abort` was set (a STOP from the web page or the jog panel)."""


class SwitchTripped(GrblError):
    """The home switch was pressed while a pattern was running: the carriage has lost its place."""


class Status:
    def __init__(self):
        self.state = 'Unknown'      # Idle, Run, Hold, Jog, Alarm, ...
        self.sub = ''               # e.g. Hold:0 = stopped, Hold:1 = still slowing down
        self.mpos = (0.0, 0.0, 0.0)
        self.wco = (0.0, 0.0, 0.0)
        self.pins = ''
        self.when = 0.0

    @property
    def wpos(self):
        return tuple(m - w for m, w in zip(self.mpos, self.wco))

    def parse(self, line):
        m = STATUS_RE.match(line)
        if not m:
            return False
        state = m.group(1)
        self.state, _, self.sub = state.partition(':')
        self.pins = ''
        wpos = None
        for field in m.group(2).split('|'):
            k, _, v = field.partition(':')
            if k in ('MPos', 'WPos', 'WCO'):
                vals = (tuple(float(x) for x in v.split(',')) + (0.0, 0.0, 0.0))[:3]
                if k == 'MPos':
                    self.mpos = vals
                elif k == 'WCO':
                    self.wco = vals
                else:
                    wpos = vals
            elif k == 'Pn':
                self.pins = v
        if wpos is not None:
            self.mpos = tuple(a + b for a, b in zip(wpos, self.wco))
        self.when = time.monotonic()
        return True


class Controller:
    """One open connection. `port` is a pyserial-like object (the tests pass a simulator)."""

    def __init__(self, port, kind='grbl', log=print):
        self.port = port
        self.kind = kind
        self.log = log
        self.status = Status()
        self.pending = []           # byte lengths of lines sent but not yet acknowledged
        self.errors = []
        self.banner = ''
        self._buf = b''
        self.abort = False          # set from another thread (the jog panel's STOP) to end a wait
        self.guard = None           # None, or 'release' / 'armed' while patterns run (see arm_guard)
        self.guard_stuck_mm = 0.0
        self.status_every = 0.08    # while guarded, ask for a status report this often (s)
        self._presses = 0
        self._last_q = 0.0

    # ------------------------------------------------------------ open / close
    @classmethod
    def open(cls, device, baud=115200, kind='grbl', log=print):
        import serial
        p = serial.Serial(device, baud, timeout=0.02, write_timeout=2)
        c = cls(p, kind, log)
        c.wake()
        return c

    def reopen(self, port):
        """Carry on with a new port (after the USB link dropped and came back): fresh state, then
        wait for the Uno's start-up banner. The position is lost, so home again before drawing."""
        try:
            if self.port:
                self.port.close()
        except Exception:
            pass
        self.port = port
        self._buf, self.banner = b'', ''
        self.pending.clear()
        self.guard, self.abort = None, False
        self.status = Status()
        self.wake()

    def wake(self, wait=3.0):
        """Opening the port resets an Uno (about 2 s); wait for its start-up banner.
        (No blank lines are sent to wake it: their late 'ok's would upset the byte counting.)"""
        t_end = time.monotonic() + wait
        seen = None
        while time.monotonic() < t_end:
            for line in self._lines():
                if 'grbl' in line.lower() or 'fluidnc' in line.lower():
                    self.banner, seen = line, time.monotonic()
            if seen and time.monotonic() - seen > 0.3:
                break
        self.port.reset_input_buffer()
        self._buf = b''
        self.pending.clear()
        self.query()

    def close(self):
        self.port.close()

    # ------------------------------------------------------------ low level
    def _lines(self):
        data = self.port.read(self.port.in_waiting or 1)
        if data:
            self._buf += data
        out = []
        while b'\n' in self._buf:
            raw, self._buf = self._buf.split(b'\n', 1)
            line = raw.decode('ascii', 'replace').strip()
            if line:
                out.append(line)
        return out

    def poll(self):
        """Read whatever has arrived and act on it."""
        if self.guard and time.monotonic() - self._last_q > self.status_every:
            self.realtime(RT_STATUS)
        for line in self._lines():
            if line == 'ok':
                if self.pending:
                    self.pending.pop(0)
            elif line.startswith('error:'):
                if self.pending:
                    self.pending.pop(0)
                self.errors.append(line)
            elif line.startswith('<'):
                self.status.parse(line)
                if self.guard:
                    self._check_guard()
            elif line.startswith('ALARM'):
                self.status.state = 'Alarm'
                raise GrblError(f'{line}: the controller stopped. Check the table, then run `home` again.')
            elif line.startswith('[MSG:'):
                self.log(f'controller says {line}')
            elif 'grbl' in line.lower():
                self.banner = line

    def realtime(self, b):
        if b == RT_STATUS:
            self._last_q = time.monotonic()
        self.port.write(b)

    # ------------------------------------------------------------ switch guard
    def arm_guard(self, cfg):
        """Stop at once if the home switch is pressed while a pattern runs.

        Patterns stay at least `rho_min_mm` out, 1 mm past where the switch trips, so they never
        press it. A press can then only mean the carriage isn't where the controller thinks:
        missed steps, a slipping belt, or something in the way. Left alone, it would grind into
        the end stop, and later patterns would run off the other end of the arm.
        Right after homing the switch may still be pressed; that's fine until it has been seen
        released once (or the carriage is clearly past it, which means the switch is stuck)."""
        self.guard = 'release'
        self.guard_stuck_mm = cfg['homing']['switch_rho_mm'] + 5.0
        self._presses = 0

    def disarm_guard(self):
        self.guard = None

    def _check_guard(self):
        st = self.status
        if st.state in ('Alarm', 'Home'):
            return
        pressed = 'Y' in st.pins
        if self.guard == 'release':
            if not pressed:
                self.guard = 'armed'
            elif st.wpos[1] > self.guard_stuck_mm:
                self._trip(f'the home switch is still pressed at rho {st.wpos[1]:.1f} mm, so it is stuck '
                           f'or its wiring is shorted')
            return
        # two reports in a row (about 0.1 s apart), so one electrical blip doesn't stop a pattern
        self._presses = self._presses + 1 if pressed else 0
        if self._presses >= 2:
            self._trip(f'the home switch was pressed at rho {st.wpos[1]:.1f} mm, where the carriage should never '
                       f'reach it. The carriage has lost its place (missed steps, a slipping belt, or '
                       f'something blocking it)')

    def _trip(self, why):
        self.realtime(RT_RESET)     # stops the motors at once; the position is lost anyway
        self.guard = None
        self.status.state = 'Alarm'
        self.pending.clear()
        raise SwitchTripped(why)

    def query(self, timeout=1.0):
        """Ask for a status report and wait for it."""
        before = self.status.when
        self.realtime(RT_STATUS)
        t_end = time.monotonic() + timeout
        while time.monotonic() < t_end:
            self.poll()
            if self.status.when != before:
                return self.status
        raise GrblError('no status report from the controller (wrong port, or not GRBL?)')

    def send(self, line):
        """Queue one line, waiting while the controller's receive buffer is full."""
        data = (line.strip() + '\n').encode('ascii')
        while sum(self.pending) + len(data) > RX_BUFFER:
            if self.abort:
                raise Aborted('stopped')
            self.poll()
        self.port.write(data)
        self.pending.append(len(data))

    def command(self, line, timeout=10.0):
        """Send one line and wait for its ok (raises on an error)."""
        n_err = len(self.errors)
        self.send(line)
        t_end = time.monotonic() + timeout
        while self.pending:
            if self.abort:
                raise Aborted('stopped')
            self.poll()
            if time.monotonic() > t_end:
                raise GrblError(f'no answer to {line!r}')
        if len(self.errors) > n_err:
            raise GrblError(f'{line!r} -> {self.errors[-1]}')

    def text(self, line, timeout=3.0):
        """Send a $ query (only while idle) and collect the lines before its ok."""
        self.port.write((line + '\n').encode())
        out, t_end = [], time.monotonic() + timeout
        while time.monotonic() < t_end:
            for ln in self._lines():
                if ln == 'ok':
                    return out
                if ln.startswith('error'):
                    raise GrblError(f'{line} -> {ln}')
                if not ln.startswith('<'):
                    out.append(ln)
        raise GrblError(f'no answer to {line!r}')

    def settings(self):
        """$$ -> {'$100': '71.111', ...}"""
        out = {}
        for ln in self.text('$$'):
            k, _, v = ln.partition('=')
            if k.startswith('$'):
                out[k] = v.split()[0] if v else ''
        return out

    def wait_idle(self, every=0.15, on_status=None):
        """Wait until everything sent is acknowledged and the machine has stopped."""
        while True:
            if self.abort:
                raise Aborted('stopped')
            st = self.query()
            if on_status:
                on_status(st)
            if not self.pending and st.state == 'Idle':
                return st
            self.pause(every)

    def pause(self, secs):
        """Wait, but keep reading the controller (and, while guarded, keep checking the switch)."""
        t_end = time.monotonic() + secs
        while time.monotonic() < t_end:
            self.poll()
            time.sleep(0.005)

    # ------------------------------------------------------------ jobs
    def stream(self, lines, on_progress=None, status_every=0.25):
        """Send G-code lines. on_progress(sent, total, status) is called a few times a second."""
        total, last, n_err = len(lines), 0.0, len(self.errors)
        for i, line in enumerate(lines, 1):
            self.send(line)
            if len(self.errors) > n_err:
                self.hold()
                raise GrblError(f'the controller rejected a line ({self.errors[-1]}); stopped')
            now = time.monotonic()
            if now - last > status_every:
                self.realtime(RT_STATUS)
                last = now
                if on_progress:
                    on_progress(i, total, self.status)
        self.wait_idle(on_status=(lambda st: on_progress(total, total, st)) if on_progress else None)

    def hold(self):
        self.realtime(RT_HOLD)

    def resume(self):
        self.realtime(RT_RESUME)

    def stop(self):
        """Stop cleanly: feed hold, wait until the motion has ended (Hold:0), then reset to throw
        away the moves still queued. Resetting while stopped keeps the machine position, but GRBL
        clears the G92 offset that homing set, so it is put back: no re-homing needed."""
        self.abort = False
        guard, self.guard = self.guard, None
        self.realtime(RT_HOLD)
        t_end = time.monotonic() + 15
        while time.monotonic() < t_end:
            st = self.query()
            if st.state == 'Idle' or (st.state == 'Hold' and st.sub == '0'):
                break
            time.sleep(0.05)
        if st.state == 'Idle' and not self.pending:
            return st                       # nothing moving and nothing queued: no reset needed
        x, y = st.wpos[0], st.wpos[1]
        self.realtime(RT_RESET)
        time.sleep(0.5)
        t_end = time.monotonic() + 3
        while time.monotonic() < t_end:
            self._lines()
            time.sleep(0.05)
        self.pending.clear()
        if self.query().state != 'Alarm':
            self.command(f'G92 X{x:.3f} Y{y:.3f}')
        if guard:
            self.guard = 'release'
            self._presses = 0
        return self.query()

    def motors_off(self):
        """Switch the motors off with the sleep command ($SLP; FluidNC has it too). They stop
        holding, so the arm turns by hand and the position is lost: before moving again the
        controller needs waking (a reset, see wake_up) and the table homing."""
        self.guard = None
        if self.query().state != 'Sleep':
            self.command('$SLP')
        self.query()

    def wake_up(self):
        """Out of sleep: only a reset wakes it, and it then waits in ALARM until unlocked."""
        if self.query().state != 'Sleep':
            return
        self.realtime(RT_RESET)
        t_end, seen = time.monotonic() + 3, None
        while time.monotonic() < t_end:
            for line in self._lines():
                if 'grbl' in line.lower() or 'fluidnc' in line.lower():
                    seen = time.monotonic()
            if seen and time.monotonic() - seen > 0.2:
                break
            time.sleep(0.02)
        self.pending.clear()
        self.query()

    def unlock(self):
        self.wake_up()
        if self.query().state == 'Alarm':
            self.command('$X')

    # ------------------------------------------------------------ homing (rho only)
    def home_rho(self, cfg, log=print):
        """Home the carriage against the switch at the centre end of the arm, then set
        theta = 0 and rho = the switch position.

        GRBL on the Uno: done from here with jog moves and the switch state that GRBL puts in
        every status report ('Pn:Y'), so the stock firmware needs no changes.
        FluidNC: its own $H (configured in its config.yaml)."""
        h = cfg['homing']
        self.disarm_guard()
        if self.kind == 'fluidnc':
            self.command('$H', timeout=90)
            self.command('G92 X0')
            return self.wait_idle()
        self.unlock()
        self.command('G21')
        self.command('G90')

        def pressed():
            return 'Y' in self.query().pins

        if pressed():
            log('the switch is already pressed; moving off it first')
            self.command(f"$J=G91 G21 Y{h['pulloff_mm'] + 3:.2f} F{h['seek_mm_s'] * 60:.0f}")
            self.wait_idle()
            if pressed():
                raise GrblError('the home switch stays pressed: check its wiring (COM and NO to the Y- pins)')
        for speed, final in ((h['seek_mm_s'], False), (h['feed_mm_s'], True)):
            self.command(f"$J=G91 G21 Y-{h['max_travel_mm']:.1f} F{speed * 60:.0f}")
            t_end = time.monotonic() + h['max_travel_mm'] / speed + 5
            found = False
            while time.monotonic() < t_end:
                if self.abort:
                    raise Aborted('stopped')
                if pressed():
                    self.realtime(RT_JOG_CANCEL)
                    found = True
                    break
                time.sleep(0.005)
            if not found:
                self.realtime(RT_JOG_CANCEL)
                self.wait_idle()
                raise GrblError('home switch not found: is the 12 V supply on? Without it the motors stay still. '
                               'Otherwise check the switch wiring, and that +Y moves the magnet outward (run `test`)')
            self.wait_idle()
            if not final:
                self.command(f"$J=G91 G21 Y{h['pulloff_mm']:.2f} F{h['seek_mm_s'] * 60:.0f}")
                self.wait_idle()
        # G92 sets the position without writing the Uno's EEPROM (G10 would, on every homing)
        self.command(f"G92 X0 Y{h['switch_rho_mm']:.3f}")
        self.command('G90')
        return self.wait_idle()