#define MODULE "tty"
#define DEBUG 0

#define BAUD_RATE 115200

#include "tty.h"
#include "config.h"
#include "fw.h"

#include <USB.h>
#include <HardwareSerial.h>

#include <rom/crc.h>

#define SOH  '\001'
#define STX  '\002'
#define ETX  '\003'
#define EOT  '\004'
#define ENQ  '\005'
#define ACK  '\006'
#define XON  '\021'		// DC1
#define WRST '\022'		// DC2 - window = 0
#define XOFF '\023'		// DC3
#define WGO  '\024'		// DC4 - window + 256 bytes
#define NAK  '\025'
#define SYN  '\026'
#define ETB  '\027'		// Not in upload mode
#define CAN  '\030'
#define EM   '\031'		// Packet offset too high
#define FS   '\034'
#define GS   '\035'
#define RS   '\036'
#define US   '\037'

#define WGO_CHUNK	 256
#define STREAMBUF_SIZE	2048
#define BUF_SLACK       (STREAMBUF_SIZE >> 1)

static char enq_str[] = "\026\035MAX80 ready\004\r\n";
static const char fwupload_start[] =
    "\034\001: /// MAX80 FW UPLOAD ~@~ $\r\n\035";

void TTY::reset()
{
    rx.rlen     = 0;
    rx.state    = rx_state::normal;
    rx.b64_bits = 0;
}

static void null_flush(void)
{
}

TTY::TTY(Stream &port)
{
    _port = &port;
    rx_sbuf = true ? xStreamBufferCreate(STREAMBUF_SIZE, 1) : NULL;
    _flush = null_flush;
    reset();
}

TTY::~TTY()
{
    if (rx_sbuf)
	vStreamBufferDelete(rx_sbuf);
}

int TTY::rxdata(void *buf, size_t len)
{
    if (!rx_sbuf)
	return 0;

    int rcv = 0;
    while (!rcv) {
	rcv = xStreamBufferReceive(rx_sbuf, buf, len, 0);
	if (rcv)
	    break;

	uint32_t now = millis();
	if (now - last_rx > 500) {
	    last_rx = now;
	    tx_credits_reset = true;
	}

	if (tx_credits_reset) {
	    // Drain input before WRST
	    //flush();
	    if (port().write(WRST)) {
		MSG("Resetting window, last_ack = %u\n", rx.last_ack);
		tx_credits_reset = false;
		tx_credits = STREAMBUF_SIZE - BUF_SLACK;
	    } else {
		// Uhm... wait a tiny bit and then try again to sent WRST?
		static bool failed_once = false;
		if (!failed_once) {
		    MSG("Failed to reset window?!\n");
		    failed_once = true;
		}
		rcv = xStreamBufferReceive(rx_sbuf, buf, len, 1);
	    }
	} else {
	    while (tx_credits >= WGO_CHUNK && port().write(WGO))
		tx_credits -= WGO_CHUNK;

	    rcv = xStreamBufferReceive(rx_sbuf, buf, len, 1);
	    tx_credits += rcv;
	}
    }

    CMSG("got %d\n", rcv);
    return rcv;
}

int TTY::rxdata(token_t me, void *buf, size_t len)
{
    TTY *tty = (TTY *)me;
    return tty->rxdata(buf, len);
}

void TTY::_upload_begin()
{
    MSG("_upload_begin\n");

    if (rx_sbuf)
	xStreamBufferReset(rx_sbuf);
    else
	rx_sbuf = xStreamBufferCreate(STREAMBUF_SIZE, 1);

    MSG("rx_sbuf = %p\n", rx_sbuf);

    if (!rx_sbuf)
	goto can;

    tx_credits_reset = false;
    tx_credits = STREAMBUF_SIZE - BUF_SLACK;
    port().write(RS);
    rx.state = rx_state::stxwait;
    rx.last_ack = 0;
    rx.rlen  = 0;
    rx.b64_bits = 0;

    MSG("firmware_update_start()\n");
    if (firmware_update_start(TTY::rxdata, (token_t)this, true))
	goto can;

    return;

can:
    port().write(CAN);
    rx.state = rx_state::normal;
    return;
}

void TTY::_onerr()
{
    if (rx.state != rx_state::normal) {
	port().write(NAK);
    }
}

static int filter_echo(int byte)
{
    if (byte >= ' ' && byte < 127)
	return byte;
    if (byte == '\t' || byte == '\r' || byte == '\n')
	return byte;

    return -1;
}

// Decode a base64 data byte (using simple offset-63)
// Call with -1 or any invalid value to invalidate the input buffer
int TTY::_decode_data(int input)
{
    unsigned int buf = rx.b64_buf;
    unsigned int inval = input - '?';

    if (inval > 63) {
	rx.b64_bits = 0;
	return -2;		// Invalid input
    }

    rx.b64_buf = buf = (buf << 6) + inval;
    rx.b64_bits += 6;

    if (rx.b64_bits >= 8) {
	rx.b64_bits -= 8;
	return (uint8_t)(buf >> rx.b64_bits);
    } else {
	return -1;
    }
}

// Change this to be a buffer...
void TTY::_onrx()
{
    int byte, data;

    while ((byte = port().read()) >= 0) {
	last_rx = millis();

	switch (rx.state) {
	case rx_state::normal:
	    if (rx.rlen < sizeof fwupload_start &&
		byte == fwupload_start[rx.rlen]) {
		if (!fwupload_start[++rx.rlen]) {
		    _upload_begin();
		    byte = -1;
		}
	    } else {
		rx.rlen = 0;
		switch (byte) {
		case ENQ:
		    port().write(enq_str);
		    byte = ETB;
		    break;
		case ETX:
		case EOT:
		case CAN:
		    byte = ETB;	// Not in file upload state
		    break;
		default:
		    // Echo if printable
		    byte = filter_echo(byte);
		    break;
		}
	    }
	    break;

	case rx_state::stxwait:
	    switch (byte) {
	    case STX:
		rx.rlen = 0;
		rx.hdr_raw[rx.rlen++] = byte;
		rx.b64_bits = 0;
		rx.state = rx_state::hdr;
		byte = -1;
		break;
	    case ETX:
	    case CAN:
		MSG("Received <%02X> waiting for STX\n", byte);
		reset();
		byte = CAN;
		break;
	    case ENQ:
		rx.rlen = 0;
		byte = GS;		 // In upload wait for STX state
		break;
	    case SYN:
		rx.rlen = 0;
		tx_credits_reset = true; // Request to resync credits
		byte = -1;
		break;
	    case EOT:
		// Upload complete, no more data
		byte = ETB;
		break;
	    default:
		byte = -1; // No echo
		break;
	    }
	    break;

	case rx_state::hdr:
	    data = _decode_data(byte);
	    byte = -1;
	    if (data < -1) {
		rx.state = rx_state::stxwait;
		rx.rlen = 0;
		tx_credits_reset = true;
		byte = US;	// Framing error
	    } else if (data == -1) {
		// Nothing to do
	    } else if (rx.rlen >= sizeof rx.hdr_raw) {
		// ERROR THIS SHOULD NEVER HAPPEN
		MSG("Header buffer overrun!!!\n");
		reset();
		byte = CAN;
	    } else {
		rx.hdr_raw[rx.rlen++] = data;
		if (rx.rlen == sizeof rx.hdr) {
		    // Start of data packet
		    MSG("Start packet hdr %d length %d offset %d last_ack %d\n",
			   rx.rlen, rx.hdr.len+1, rx.hdr.offs, rx.last_ack);
		    rx.state = rx_state::data;
		    rx.rlen = 0;
		}
	    }
	    break;

	case rx_state::data:
	    data = _decode_data(byte);
	    byte = -1;
	    if (data < -1) {
		rx.state = rx_state::stxwait;
		rx.rlen = 0;
		tx_credits_reset = true;
		byte = US;	// Framing error
	    } else if (data == -1) {
		// Nothing to do
	    } else if (rx.rlen >= sizeof rx_data) {
		// ERROR THIS SHOULD NEVER HAPPEN
		MSG("Packet data buffer overrun!!!\n");
		reset();
		byte = CAN;
	    } else {
		rx_data[rx.rlen++] = data;
		// rx.hdr.len = packet data len - 1
		if (rx.rlen > rx.hdr.len) {
		    int have = rx.rlen;
		    uint32_t crc;

		    if (have != rx.hdr.len + 1) {
			MSG("Invalid packet length (should not happen...)\n");
			byte = NAK;
		    } else if ((crc = crc32_le(0, rx_data, have))
			       != rx.hdr.crc) {
			MSG("Packet CRC error hdr %08x data %08x\n", rx.hdr.crc, crc);
			byte = NAK;
		    } else if (rx.hdr.offs > rx.last_ack) {
			MSG("Invalid packet offsets [%d..%d) at %d\n",
			       rx.hdr.offs, rx.hdr.offs + have, rx.last_ack);
			byte = EM;
		    } else if (rx.hdr.offs + have <= rx.last_ack) {
			// Ack and discard packet below current window (transmitter is catching up)
			byte = ACK;
		    } else {
			int sent = 0;

			uint32_t skip = rx.last_ack - rx.hdr.offs;
			have -= skip;
			sent = xStreamBufferSend(rx_sbuf, rx_data+skip, have, 0);
			rx.last_ack += sent;

			if (sent != have) {
			    MSG("Packet underrun, got %d, expected %d\n", sent, have);
			    byte = NAK;
			} else {
			    MSG("%d bytes received OK\n", sent);
			    byte = ACK;
			}
		    }
		    if (byte != ACK)
			tx_credits_reset = true;
		    rx.state = rx_state::stxwait;
		    rx.rlen  = 0;
		}
	    }

	    break;
	}

	if (byte >= 0)
	    port().write(byte);
    }
}

void TTY::_onconnect()
{
    do_log_config_status = true;
    port().write(XON);
}

void TTY::_onbreak()
{
    reset();
}

void TTY::_ondisconnect()
{
    reset();
}

static TTY *ttys[2];
#define uart_tty (ttys[0])
#define usb_tty  (ttys[1])

static TaskHandle_t usb_tty_task;

#define USB_NOTIFY_INDEX     0   /* Seems to be the only available... */
#define USB_TASK_STACK	  4096
#define USB_TASK_PRIORITY    5

void TTY::usb_task_handler(void *pvt)
{
    (void)pvt;

    while (1) {
	uint32_t notify_value;

	xTaskNotifyWaitIndexed(USB_NOTIFY_INDEX, 0, 1,
			       &notify_value, portMAX_DELAY);
	usb_tty->_onrx();
    }
}

void TTY::usb_onevent(void *arg, esp_event_base_t event_base,
		      int32_t event_id, void *event_data)
{
    if (event_base == ARDUINO_USB_CDC_EVENTS) {
	switch (event_id) {
	case ARDUINO_USB_CDC_CONNECTED_EVENT:
	case ARDUINO_USB_CDC_LINE_STATE_EVENT:
	    usb_tty->_onconnect();
	    break;
	case ARDUINO_USB_CDC_DISCONNECTED_EVENT:
	    usb_tty->_ondisconnect();
	    break;
	case ARDUINO_USB_CDC_RX_EVENT:
	    xTaskNotifyIndexed(usb_tty_task, USB_NOTIFY_INDEX, 1, eSetBits);
	    break;
	case ARDUINO_USB_CDC_RX_OVERFLOW_EVENT:
	    usb_tty->_onerr();
	    break;
	default:
	    // Do nothing
	    break;
	}
    }
}

static void usb_flush()
{
    Serial.flush();
}

void TTY::uart_onrx(void)
{
    uart_tty->_onrx();
}

void TTY::uart_onerr(hardwareSerial_error_t err)
{
    switch (err) {
    case UART_BREAK_ERROR:
	uart_tty->_onbreak();
	break;
    default:
	uart_tty->_onerr();
	break;
    }
}

static void uart_flush()
{
    Serial0.flush(true);
}

void TTY::init()
{
    uart_tty = new TTY(Serial0);
    uart_tty->_flush = uart_flush;
    Serial0.begin(BAUD_RATE);
    Serial0.onReceive(uart_onrx, false);
    Serial0.onReceiveError(uart_onerr);
    static const char uart_tty_msg[] = "[TTY] UART initialized\r\n";
    uart_tty->port().write(uart_tty_msg, sizeof uart_tty_msg - 1);

    usb_tty  = new TTY(Serial);
    usb_tty->_flush = usb_flush;
    if (xTaskCreate(usb_task_handler, "usbttyd",
		    USB_TASK_STACK, usb_tty,
		    USB_TASK_PRIORITY, &usb_tty_task) == pdPASS) {
	Serial.onEvent(usb_onevent);
	xTaskNotifyIndexed(usb_tty_task, USB_NOTIFY_INDEX, 1, eSetBits);
    }
    Serial.enableReboot(true);
    static const char usb_tty_msg[] = "[TTY] USB initialized\r\n";
    usb_tty->port().write(usb_tty_msg, sizeof usb_tty_msg - 1);
}

void TTY::ping()
{
}

void logmsg(const char *module, const char *fmt, ...)
{
    char buf[128];
    int len = 0;
    va_list ap;

    if (module)
	len = snprintf(buf, sizeof buf-1, "%-7s : ", module);

    va_start(ap, fmt);
    len += vsnprintf(buf+len, sizeof buf-len, fmt, ap);
    va_end(ap);

    if (len > sizeof buf - 2) {
	len = sizeof buf - 2;
	buf[len-1] = '\n';
    }
    if (buf[len-1] == '\n') {
	buf[len-1] = '\r';
	buf[len++] = '\n';
    }
    buf[len] = '\0';

    for (int i = 0; i < ARRAY_SIZE(ttys); i++) {
	if (ttys[i])
	    ttys[i]->port().write(buf, len);
    }
}
