"""sandtable: drive the sand table from this computer over USB.

Run `python3 -m sandtable --help` from the software folder.
"""
import argparse
import glob
import os
import random
import sys
import time

from . import config, patterns, thr
from .motion import HEADER, Planner

USB_IDS = {0x2341: 'Arduino', 0x2A03: 'Arduino', 0x1A86: 'CH340', 0x0403: 'FTDI', 0x10C4: 'CP210x', 0x303A: 'ESP32'}


# ---------------------------------------------------------------- helpers
def find_port(cfg, explicit=None):
    want = explicit or cfg['serial']['port']
    if want and want != 'auto':
        return want
    from serial.tools import list_ports
    cands = [p for p in list_ports.comports() if p.vid in USB_IDS]
    if len(cands) == 1:
        return cands[0].device
    if not cands:
        sys.exit('No Arduino found on USB. Plug it in, or give --port (see `ports`).')
    sys.exit('Several possible ports: ' + ', '.join(p.device for p in cands) + '. Pick one with --port.')


def connect(cfg, args):
    from .grbl import Controller
    dev = find_port(cfg, args.port)
    print(f'connecting to {dev} ...', flush=True)
    try:
        c = Controller.open(dev, cfg['serial']['baud'], cfg['controller']['type'])
    except Exception as e:      # pyserial raises SerialException with errno 13, not PermissionError
        if getattr(e, 'errno', None) != 13 and 'Permission denied' not in str(e):
            raise
        sys.exit(f'No permission for {dev}.\n'
                 f'  Quick fix (until it is unplugged):  sudo chmod a+rw {dev}\n'
                 f'  Lasting fix:  sudo usermod -aG dialout $USER  (then log out and back in)')
    print(f"  {c.banner or 'controller'}  state: {c.status.state}")
    return c


def alert(title, msg):
    """Make sure a person notices: a beep and a big message here, plus a desktop notification."""
    print(f'\a\n*** {title} ***\n{msg}\n', flush=True)
    if os.environ.get('SANDTABLE_QUIET'):          # set by the tests: no desktop pop-ups
        return
    try:
        import subprocess
        subprocess.run(['notify-send', '-u', 'critical', '-a', 'Gyre', title, msg],
                       timeout=5, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except Exception:
        pass


def check_floor(cfg):
    m, h = cfg['machine'], cfg['homing']
    if m['rho_min_mm'] < h['switch_rho_mm'] + 0.5:
        sys.exit(f"machine.rho_min_mm ({m['rho_min_mm']}) must be at least 0.5 mm more than "
                 f"homing.switch_rho_mm ({h['switch_rho_mm']}), or patterns press the home switch")


def guarded(fn):
    """Run fn(); if the switch guard or the controller stops the machine, tell the user loudly."""
    from .grbl import GrblError, SwitchTripped
    try:
        return fn()
    except SwitchTripped as e:
        alert('Gyre stopped: home switch pressed',
              f'{e}.\nThe motors were stopped at once. Check the arm, belts and carriage, then run '
              f'`play` again; it homes first.')
        sys.exit(1)
    except GrblError as e:
        alert('Gyre stopped', str(e))
        sys.exit(1)
    except OSError as e:        # pyserial's SerialException is an OSError: the USB link dropped
        alert('Gyre: USB connection lost',
              f'{e}\nThe Uno restarts when it is plugged back in, so the table stops where it is. '
              f'Run `play` again; it homes first.')
        sys.exit(1)


def ask(prompt, choices):
    while True:
        a = input(f'{prompt} [{"/".join(choices)}] ').strip().lower()
        if a and a[0] in choices:
            return a[0]


def fmt_time(s):
    s = int(s)
    return f'{s // 3600}h{s % 3600 // 60:02d}m' if s >= 3600 else f'{s // 60}m{s % 60:02d}s'


def expand(items):
    files = []
    for it in items:
        if os.path.isdir(it):
            files += sorted(glob.glob(os.path.join(it, '**', '*.thr'), recursive=True))
        else:
            files += sorted(glob.glob(it)) or [it]
    return files


def load_any(item):
    """A .thr file, or a generator written as gen:NAME[:key=value,key=value]."""
    if item.startswith('erase:'):           # erase only: end at the centre or the edge
        r = 0.0 if item.endswith('centre') else 1.0
        return 'erase', [(0.0, r), (0.0, r)]
    if item.startswith('gen:'):
        _, name, *rest = item.split(':', 2)
        params = dict(kv.split('=', 1) for kv in rest[0].split(',')) if rest and rest[0] else {}
        return name, patterns.generate(name, **params)
    return os.path.basename(item), thr.load(item)


def run_job(c, lines, label, est):
    """Stream with progress; Ctrl-C pauses and asks what to do."""
    t0 = time.monotonic()

    def progress(i, n, st):
        el = time.monotonic() - t0
        left = max(0.0, est * (1 - i / n)) if n else 0
        sys.stdout.write(f'\r  {label}: {100 * i / n:5.1f}%  {fmt_time(el)} elapsed, ~{fmt_time(left)} left   ')
        sys.stdout.flush()

    from .grbl import RT_STATUS
    i, n_err, last = 0, len(c.errors), 0.0
    while i < len(lines):
        try:
            for j in range(i, len(lines)):
                c.send(lines[j])
                i = j + 1
                if len(c.errors) > n_err:
                    c.hold()
                    raise SystemExit(f'\nthe controller rejected a line ({c.errors[-1]}): {lines[j]!r}; paused')
                now = time.monotonic()
                if now - last > 0.25:
                    c.realtime(RT_STATUS)
                    last = now
                    progress(i, len(lines), c.status)
            c.wait_idle(on_status=lambda st: progress(len(lines), len(lines), st))
        except KeyboardInterrupt:
            c.hold()
            print('\n  paused (feed hold).')
            if ask('  resume or stop?', 'rs') == 'r':
                c.resume()
                continue
            print('  stopping ...')
            c.stop()
            raise SystemExit('stopped. The position is kept; run `play` again to continue with another pattern.')
    print()


# ---------------------------------------------------------------- commands
def cmd_ports(cfg, args):
    from serial.tools import list_ports
    usb = [p for p in list_ports.comports() if p.vid is not None]     # skip the built-in /dev/ttyS* ports
    if not usb:
        print('No USB serial ports. Is the Arduino plugged in?')
    for p in usb:
        tag = USB_IDS.get(p.vid, '')
        print(f'{p.device:16s} {tag:8s} {p.description}')


FIRMWARE = os.path.join(os.path.dirname(config.HERE), 'firmware', 'grbl_v1.1h.20190825.hex')


def cmd_flash(cfg, args):
    """Put GRBL 1.1h on the Arduino Uno (replaces whatever program is on it)."""
    import shutil
    import subprocess
    acli = shutil.which('arduino-cli') or os.path.expanduser('~/.local/bin/arduino-cli')
    if not os.path.exists(acli):
        sys.exit('arduino-cli not found. See "Flashing GRBL" in software/README.md.')
    if cfg['controller']['type'] != 'grbl':
        sys.exit('flash is for the Arduino Uno; the ESP32 board gets FluidNC instead (see pcb/README.md).')
    dev = find_port(cfg, args.port)
    print(f'This replaces the program on the Arduino at {dev} with GRBL 1.1h ({os.path.basename(FIRMWARE)}).')
    print('Turn the 12 V supply off first; only the USB cable is needed.')
    if not args.yes and ask('Flash it?', 'yn') != 'y':
        return
    for board in ('arduino:avr:uno', 'arduino:avr:nano:cpu=atmega328old'):   # Uno, then old-bootloader clones
        print(f'  uploading as {board} ...', flush=True)
        r = subprocess.run([acli, 'upload', '-b', board, '-p', dev, '--input-file', FIRMWARE],
                           capture_output=True, text=True)
        if r.returncode == 0:
            break
        msg = (r.stderr or r.stdout).strip().splitlines()
        print('    failed: ' + (msg[-1] if msg else 'unknown error'))
        if 'ermission' in r.stderr:
            sys.exit(f'No permission for {dev}. Run: sudo usermod -aG dialout $USER  (then log out and back in)')
    else:
        sys.exit('Upload failed both ways. Check the USB cable (some are charge-only) and the port.')
    c = None
    for attempt in range(4):          # the board restarts after the upload; give it a moment
        time.sleep(1.5)
        try:
            c = connect(cfg, args)
            break
        except Exception as e:
            print(f'  not answering yet ({e}); retrying ...')
    if c is None or 'grbl' not in c.banner.lower():
        sys.exit('Uploaded, but GRBL did not answer. Try unplugging and replugging the USB cable, then `info`.')
    c.command('$RST=*')          # clear anything an older program left in the EEPROM
    print('  GRBL is on the Arduino, with default settings. Next: `setup`, then `test`.')
    c.close()


def cmd_web(cfg, args):
    from .grbl import Controller
    from .web import serve
    check_floor(cfg)
    if args.port:
        cfg['serial']['port'] = args.port
    try:
        c, ok = connect(cfg, args), True
    except SystemExit as e:             # no Uno yet: start anyway, the page connects when it appears
        print(f'  {e}\n  Starting the page anyway; it connects when the Uno is plugged in.')
        c, ok = Controller(None, cfg['controller']['type']), False
    serve(c, cfg, host=args.host, port=args.http_port, open_browser=not args.no_browser,
          open_path='/jog' if args.cmd == 'ui' else '/', alert=alert, connected=ok)


def cmd_info(cfg, args):
    c = connect(cfg, args)
    st = c.query()
    print(f'  position: theta {st.wpos[0]:.2f} deg, rho {st.wpos[1]:.2f} mm (not homed yet after connecting)')
    if cfg['controller']['type'] == 'grbl':
        have = c.settings()
        want = config.grbl_settings(cfg)
        diff = {k: (have.get(k), v) for k, v in want.items() if _num(have.get(k)) != _num(v)}
        print(f"  GRBL settings: {len(have)} read, {len(diff)} differ from this machine's values")
        for k, (h, w) in diff.items():
            print(f'    {k}: {h} -> {w}')
        if diff:
            print('  Run `setup` to write them.')
    c.close()


def _num(v):
    try:
        return round(float(v), 3)
    except (TypeError, ValueError):
        return v


def cmd_setup(cfg, args):
    if cfg['controller']['type'] != 'grbl':
        sys.exit('setup writes GRBL settings; FluidNC is set up with pcb/fluidnc/config.yaml instead.')
    c = connect(cfg, args)
    have = c.settings()
    want = config.grbl_settings(cfg)
    diff = {k: v for k, v in want.items() if _num(have.get(k)) != _num(v)}
    if not diff:
        print('  GRBL settings already match.')
        return
    for k, v in diff.items():
        print(f'  {k}: {have.get(k)} -> {v}')
    if not args.yes and ask('  write these to the Uno?', 'yn') != 'y':
        return
    for k, v in diff.items():
        c.command(f'{k}={v}')
    print('  written.')
    c.close()


def cmd_test(cfg, args):
    """Check motor directions: +theta should turn the arm counter-clockwise seen from above,
    +rho should move the magnet outward."""
    c = connect(cfg, args)
    c.unlock()
    print('\nThe motors are free until the first move. Slide the magnet carriage to roughly the')
    print('middle of the arm by hand, so a 15 mm move either way is safe.')
    input('Press Enter when ready ... ')
    c.command('G21')
    c.command('G91')
    c.command('G1 X90 F1200')
    c.wait_idle()
    ccw = ask('Did the arm turn a quarter turn COUNTER-clockwise (looking down through the glass)?', 'yn')
    c.command('G1 X-90 F1200')
    c.command('G1 Y15 F600')
    c.wait_idle()
    out = ask('Did the magnet move OUTWARD, away from the centre?', 'yn')
    c.command('G1 Y-15 F600')
    c.command('G90')
    c.wait_idle()
    flip = (1 if ccw == 'n' else 0) | (2 if out == 'n' else 0)
    if not flip:
        print('Both directions are right.')
    elif cfg['controller']['type'] == 'grbl':
        mask = int(float(c.settings().get('$3', '0'))) ^ flip
        c.command(f'$3={mask}')
        print(f'Flipped {"theta " if flip & 1 else ""}{"rho" if flip & 2 else ""} in GRBL ($3={mask}). '
              'Run `test` again to confirm.')
    else:
        print('Add :low to the direction_pin of ' + ' and '.join(
            a for a, b in (('x (theta)', 1), ('y (rho)', 2)) if flip & b) + ' in config.yaml on the board.')
    c.close()


def cmd_home(cfg, args):
    c = connect(cfg, args)
    st = c.home_rho(cfg)
    print(f'  homed: rho {st.wpos[1]:.2f} mm')
    c.close()


def plan_items(cfg, args, items):
    pl = Planner(cfg, 0.0, cfg['homing']['switch_rho_mm'])
    jobs = []
    for item in items:
        name, pts = load_any(item)
        lines, segs, est = pl.program(pts, erase=not args.no_erase, speed=args.speed,
                                      rotate_deg=args.rotate, mirror=args.mirror)
        jobs.append((name, lines, segs, est))
    return jobs


def cmd_play(cfg, args):
    items = expand(args.items)
    if not items:
        sys.exit('nothing to play')
    if args.dry_run or args.preview or args.gcode:
        jobs = plan_items(cfg, args, items)
        total = sum(j[3] for j in jobs)
        for name, lines, segs, est in jobs:
            print(f'  {name}: {len(lines)} lines, about {fmt_time(est)}')
        print(f'  total about {fmt_time(total)}')
        if args.preview:
            from .preview import render, sheet
            inv, R = cfg['machine']['invert_theta'], cfg['machine']['rho_max_mm']
            if len(jobs) == 1:
                render(jobs[0][2], R, args.preview, invert_theta=inv)
            else:
                sheet([(j[0], j[2]) for j in jobs], R, args.preview, invert_theta=inv)
            print(f'  preview: {args.preview}')
        if args.gcode:
            with open(args.gcode, 'w') as f:
                f.write('\n'.join(HEADER + [l for j in jobs for l in j[1]]) + '\n')
            print(f'  G-code: {args.gcode}')
        if args.dry_run:
            return
    check_floor(cfg)
    c = connect(cfg, args)
    guarded(lambda: _play(c, cfg, args, items))
    c.close()


def _play(c, cfg, args, items):
    try:
        _play_list(c, cfg, args, items)
    finally:
        if cfg['motion'].get('motors_off_when_done', True):
            try:
                c.motors_off()                  # done, stopped or failed: let the motors rest
                print('  motors off')
            except Exception:
                pass                            # e.g. the USB link has gone


def _play_list(c, cfg, args, items):
    if not args.no_home:
        print('  homing rho ...', flush=True)
        c.home_rho(cfg)
    st = c.query()
    pl = Planner(cfg, st.wpos[0], st.wpos[1])
    for h in HEADER:
        c.command(h)
    c.arm_guard(cfg)        # from here on, a press of the home switch stops everything
    rnd = random.Random()
    while True:
        order = list(items)
        if args.shuffle:
            rnd.shuffle(order)
        for n, item in enumerate(order, 1):
            name, pts = load_any(item)
            lines, segs, est = pl.program(pts, erase=not args.no_erase, speed=args.speed,
                                          rotate_deg=args.rotate, mirror=args.mirror)
            print(f'[{n}/{len(order)}] {name}  (~{fmt_time(est)})')
            run_job(c, lines, name, est)
            if args.pause:
                time.sleep(args.pause)
        if not args.loop:
            break


def cmd_gen(cfg, args):
    if args.name == 'list' or not args.name:
        for name, (fn, defaults) in patterns.GENERATORS.items():
            print(f'  {name:11s} {", ".join(f"{k}={v}" for k, v in defaults.items()) or "(no settings)"}')
            print(f'              {fn.__doc__.strip().splitlines()[0]}')
        return
    params = dict(kv.split('=', 1) for kv in args.params)
    pts = patterns.generate(args.name, **params)
    print(f'  {args.name}: {len(pts)} points')
    if args.save:
        thr.save(args.save, pts, header=f'sandtable generator {args.name} {" ".join(args.params)}')
        print(f'  saved {args.save}')
    item = 'gen:' + args.name + (':' + ','.join(args.params) if args.params else '')
    if args.play or args.preview or args.dry_run or args.gcode:
        args.items = [item]
        cmd_play(cfg, args)


def cmd_erase(cfg, args):
    args.items = ['erase:' + args.to]
    args.no_erase = False
    cmd_play(cfg, args)


def cmd_move(cfg, args):
    check_floor(cfg)
    c = connect(cfg, args)

    def go():
        if not args.no_home:
            c.home_rho(cfg)
        rho = min(max(args.rho, cfg['machine']['rho_min_mm']), cfg['machine']['rho_max_mm'])
        c.command('G90')
        c.arm_guard(cfg)
        c.command(f'G1 X{args.theta:.3f} Y{rho:.3f} F600')
        c.wait_idle()
    guarded(go)
    c.close()


# ---------------------------------------------------------------- main
def play_options(p):
    p.add_argument('--no-erase', action='store_true', help='skip the erase spiral before each pattern')
    p.add_argument('--speed', type=float, help='ball speed in mm/s (default from sandtable.toml)')
    p.add_argument('--rotate', type=float, default=0.0, help='turn patterns by this many degrees')
    p.add_argument('--mirror', action='store_true', help='mirror patterns')
    p.add_argument('--loop', action='store_true', help='repeat the list forever')
    p.add_argument('--shuffle', action='store_true', help='random order')
    p.add_argument('--pause', type=float, default=0, help='seconds to wait between patterns')
    p.add_argument('--no-home', action='store_true', help="don't home rho first (only if nothing was reset)")
    p.add_argument('--dry-run', action='store_true', help="don't connect; just plan (with --preview / --gcode)")
    p.add_argument('--preview', metavar='PNG', help='draw what the ball will trace')
    p.add_argument('--gcode', metavar='FILE', help='also write the G-code to a file')


def main(argv=None):
    ap = argparse.ArgumentParser(prog='sandtable', description='Drive the sand table over USB.')
    ap.add_argument('--config', help='settings file (default: software/sandtable.toml)')
    ap.add_argument('--port', help='serial port, e.g. /dev/ttyACM0 (default: auto)')
    sub = ap.add_subparsers(dest='cmd', required=True)
    sub.add_parser('ports', help='list serial ports')
    sub.add_parser('info', help='connect and show the controller and its settings')
    for name, text in (('web', 'the web page: a library of designs to play, plus the jog panel'),
                       ('ui', 'the same web page, opened at the jog panel: home, jog, calibrate')):
        u = sub.add_parser(name, help=text)
        u.add_argument('--http-port', type=int, default=8765)
        u.add_argument('--host', default='127.0.0.1',
                       help='127.0.0.1 = this computer only (default); 0.0.0.0 = also phones on your network (no password)')
        u.add_argument('--no-browser', action='store_true')
    fl = sub.add_parser('flash', help='put GRBL 1.1h on the Arduino Uno (one time)')
    fl.add_argument('--yes', action='store_true')
    s = sub.add_parser('setup', help="write this machine's GRBL settings to the Uno")
    s.add_argument('--yes', action='store_true')
    sub.add_parser('test', help='check the motor directions (moves the arm a little)')
    sub.add_parser('home', help='home rho against the switch')
    p = sub.add_parser('play', help='play .thr files, folders of them, or gen:NAME[:k=v,...] patterns')
    p.add_argument('items', nargs='+')
    play_options(p)
    g = sub.add_parser('gen', help='make a pattern with a built-in generator (`gen list` shows them)')
    g.add_argument('name', nargs='?', default='list')
    g.add_argument('params', nargs='*', help='settings as key=value')
    g.add_argument('--save', metavar='THR', help='save it as a .thr file')
    g.add_argument('--play', action='store_true', help='draw it now')
    play_options(g)
    e = sub.add_parser('erase', help='spiral-erase the whole tray')
    e.add_argument('--to', choices=('centre', 'edge'), default='centre', help='where the ball ends up')
    play_options(e)
    m = sub.add_parser('move', help='move the ball to theta (deg) and rho (mm)')
    m.add_argument('theta', type=float)
    m.add_argument('rho', type=float)
    m.add_argument('--no-home', action='store_true')
    args = ap.parse_args(argv)
    cfg = config.load(args.config)
    {'ports': cmd_ports, 'info': cmd_info, 'flash': cmd_flash, 'ui': cmd_web, 'web': cmd_web, 'setup': cmd_setup, 'test': cmd_test, 'home': cmd_home,
     'play': cmd_play, 'gen': cmd_gen, 'erase': cmd_erase, 'move': cmd_move}[args.cmd](cfg, args)