#!/usr/bin/env python3

import os
from time import time
from binascii import hexlify
import enum
import struct
import zlib
import sys
import sqlite3

import serial
from cobs import cobs


class CtrlPacketTypes(enum.Enum):
    RESET = 1
    ACK = 2
    RETRANSMIT = 3

def unpack_head(fmt, data):
    split = struct.calcsize(fmt)
    return [ *struct.unpack(fmt, data[:split]), data[split:] ]

def ctrl_packet(ptype, pid=0):
    return cobs.encode(struct.pack('BB', ptype.value, pid)) + b'\0'

ctrl_reset = lambda: ctrl_packet(CtrlPacketTypes.RESET)
ctrl_ack = lambda pid: ctrl_packet(CtrlPacketTypes.ACK, pid)
ctrl_retransmit = lambda pid: ctrl_packet(CtrlPacketTypes.RETRANSMIT, pid)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()

    parser.add_argument('-b', '--baudrate', type=int, default=250000)
    parser.add_argument('port', nargs='?', default=None)
    parser.add_argument('dbfile')
    args = parser.parse_args()

    if args.port is None:
        try:
            candidate, = os.listdir('/dev/serial/by-id')
            args.port = os.path.join('/dev/serial/by-id', candidate)
            print(f'No port given, guessing {args.port}')

        except:
            print('No port given and could not guess port. Exiting.')
            sys.exit(1)
    
    ser = serial.Serial(args.port, args.baudrate, timeout=1.0)
    db = sqlite3.connect(args.dbfile)
    db.execute('CREATE TABLE IF NOT EXISTS measurements (run_id INTEGER, rx_ts INTEGER, seq INTEGER, data BLOB)')
    db.execute('''CREATE TABLE IF NOT EXISTS errors (
                run_id INTEGER,
                rx_ts INTEGER,
                type TEXT,
                seq INTEGER,
                pid INTEGER,
                pid_expected INTEGER,
                crc32 INTEGER,
                crc32_expected INTEGER,
                data BLOB)''')
    run_id, = db.execute('SELECT IFNULL(MAX(run_id), -1) + 1 FROM measurements').fetchone()

    ser.flushInput()
    ser.write(ctrl_reset())
    ser.flushOutput()

    last_pid = None
    lines_written = 0
    cur = db.cursor()
    while True:
        #ser.write(cobs.encode(b'\x01\xff') + b'\0')
        data = ser.read_until(b'\0')
        if not data or data[-1] != 0x00:
            #print(f'{time():>7.3f} Timeout: resetting')
            #ser.write(cobs.encode(b'\x01\xff') + b'\0') # reset
            continue

        try:
            if len(data) <= 1: # delimiting zero for retransmission
                cur.execute('INSERT INTO errors(run_id, rx_ts, type) VALUES (?, ?, "retransmission")',
                        (run_id, int(time()*1000)))
                continue
            crc32, payload = unpack_head('I', cobs.decode(data[:-1]))
            pid, seq, data = unpack_head('xBH', payload)
            ts = time()

            # Calculate byte-wise CRC32
            our_crc = zlib.crc32(bytes(b for x in payload for b in (0, 0, 0, x)))
            #log.append((time(), seq, crc32, our_crc, pid, data))
            print(f'{ts:>7.3f} {seq:05d} {crc32:08x} {our_crc:08x} {pid} {hexlify(data).decode()}', end='')

            error = False
            suppress_ack = False
            if crc32 != our_crc:
                print(' CRC ERROR', end='')
                suppress_ack = True
                error = True

            if last_pid is not None and pid != (last_pid+1)%8:
                print(' PID ERROR', end='')
                error = True
            else:
                last_pid = pid

            if not suppress_ack:
                ser.write(ctrl_ack(pid))
                ser.flushOutput()

            if not error:
                cur.execute('INSERT INTO measurements VALUES (?, ?, ?, ?)', (run_id, int(ts*1000), seq, data))
            else:
                cur.execute('INSERT INTO errors VALUES (?, ?, "pid", ?, ?, ?, ?, ?, ?)',
                        (run_id, int(ts*1000), seq, pid, (last_pid+1)%8, crc32, our_crc, data))

            print()
            lines_written += 1
            if lines_written == 80:
                lines_written = 0
                print('\033[2J\033[H', end='')
                db.commit()

        except Exception as e:
            print(e, len(data))
            ser.write(ctrl_ack(0)) # FIXME delet this