#!/usr/bin/perl

use strict;
use integer;
use PerlIO::gzip;
use File::Temp;
use File::Spec;
use Time::HiRes qw(usleep tv_interval);
use Digest::CRC qw(crc32);
use v5.10;			# For "state"

my $esptool = ($^O eq 'MSWin32') ? 'esptool.exe' : 'esptool.py';
$esptool = $ENV{'ESPTOOL'} || $esptool;

my $esp_retries = 5;

my $FW_MAGIC = 0x7a07fbd6;

my %datatypes = (
    'end'       => 0,		# End of data
    'data'      => 1,		# FPGA flash data
    'target'    => 2,		# Firmware target string
    'note'      => 3,		# Informative string
    'espota'    => 4,		# ESP32 OTA image
    'fpgainit'  => 5,		# FPGA bypass (transient) image during update
    'esppart'   => 6,		# ESP32 partition table
    'espsys'    => 7,		# ESP32 boot loader, OTA control partition...
    'esptool'   => 8,		# esptool.py options for flashing
    'boardinfo' => 9		# board_info base address
    );
my @type;
foreach my $t (keys(%datatypes)) {
    $type[$datatypes{$t}] = $t;
}

my $FDF_OPTIONAL = 0x0001;

my $STRING_MAX_LEN = 4095;

# For debug; DC1-4 replaced with functional names
my @ascii = qw(NUL SOH STX ETX EOT ENQ ACK BEL BS HT LF VT FF CR SO SI
	       DLE XON WRST XOFF WGO NAK SYN ETB CAN EM SUB ESC FS GS RS US);
foreach my $i (9, 10, 13, 32..126) {
    $ascii[$i] = chr($i);
}
$ascii[127] = 'DEL';
for (my $i = 128; $i < 256; $i++) {
    $ascii[$i] = sprintf("%02X", $i);
}
my @a = map { length($_) == 1 ? $_ : '<'.$_.'>' } @ascii;

# Simple base64 encode using 3F-7E, bytewise bigendian
sub mybaseencode {
    my $nbits  = 0;
    my $bitbuf = 0;
    my $out    = '';

    foreach my $s (@_) {
	foreach my $c (unpack('C*', $s)) {
	    $nbits += 8;
	    $bitbuf = ($bitbuf << 8) + $c;
	    while ($nbits >= 6) {
		$nbits -= 6;
		$out .= pack('C', 63 + (($bitbuf >> $nbits) & 63));
	    }
	}
    }

    if ($nbits) {
	$out .= pack('C', 63 + ($bitbuf & ((1 << $nbits) - 1)));
    }

    return $out;
}

sub getint($) {
    my($s) = @_;

    return undef
	unless ($s =~ /^(([1-9][0-9]+)|(0(x[0-9a-f]+|[0-7]*)))([kmgtpe]?)$/i);

    my $o = oct($3) + $2;
    my $p = lc($5);

    if ($p eq 'k') {
	$o <<= 10;
    } elsif ($p eq 'm') {
	$o <<= 20;
    } elsif ($p eq 'g') {
	$o <<= 30;
    } elsif ($p eq 't') {
	$o <<= 40;
    } elsif ($p eq 'p') {
	$o <<= 50;
    } elsif ($p eq 'e') {
	$o <<= 60;
    }
    return $o;
}

sub filelen($) {
    my($f) = @_;
    my @s = stat($f);

    return $s[7];
}

sub unquote_cmd($) {
    my($s) = @_;
    my @a;

    $s =~ s/[\r\n]+/ /g;
    while ($s =~ /^\s*(?:\"((?:[^\"]+|\"\")*)\"|(\S+))(\s.*)?$/) {
	push(@a, $1.$2);
	$s = $3;
    }

    return @a;
}

# Similar to grep, but for a hash; also filters out 
sub hgrep(&%) {
    my($mfunc, %hash) = @_;

    return map { $_ => $hash{$_} } grep(&$mfunc, keys %hash);
}

# Wrapper for running esptool, returns a hash with info output
# or dies on failure
sub hash2opt($)
{
    my($h) = @_;

    return () unless (defined($h));
    return map { $h->{$_} ne '' ? ('--'.$_, $h->{$_}) : () } sort keys %{$h};
}

sub run_esptool($$$$@)
{
    my($port,$common_options,$cmd,$cmd_options,@args) = @_;

    my @espcmd = unquote_cmd($esptool);
    push(@espcmd, '--port', $port);
    push(@espcmd, hash2opt($common_options));
    push(@espcmd, unquote_cmd($cmd));
    push(@espcmd, hash2opt($cmd_options));
    push(@espcmd, @args);

    my $retries = $esp_retries;
    my $ok = 0;
    my %outinfo;
    my @output;

    while (!$ok && $retries--) {
	%outinfo = ();
	@output  = ();

	print STDERR 'Command: ', join(' ', @espcmd), "\n";
	print STDERR "Running $espcmd[0] $cmd... ";

	my $esp;
	if (!open($esp, '-|', @espcmd)) {
	    print STDERR $!, "\n";
	    exit 1;
	}

	while (defined(my $l = <$esp>)) {
	    if ($l =~ /^Chip is (\S+)/) {
		$outinfo{'chip'} = $1;
	    } elsif ($l =~ /^MAC: ([0-9a-f:]+)/) {
		$outinfo{'mac'} = $1;
	    }
	    push(@output, $l);
	}

	$ok = close($esp);
	if ($ok) {
	    print STDERR "ok\n";
	    last;
	} elsif ($retries) {
	    print STDERR "failed, retrying\n";
	    usleep(1000000);
	} else {
	    print STDERR "failed, giving up\n";
	}
    }

    print STDERR @output;

    die "$0: $espcmd[0] $cmd failed\n" unless ($ok);
    return %outinfo;
}

my @args = @ARGV;
my $esponly = 0;
my $file;
while (1) {
    $file = shift(@args);
    last if ($file !~ /^\-/);

    if ($file eq '--esponly') {
	$esponly = 1;
    } elsif ($file eq '--') {
	$file = shift(@args);
	last;
    } else {
	undef $file;		# Invalid argument, print usage
	last;
    }
}

my $port = shift(@args);
if (!defined($port)) {
    die "Usage: $0 file.fw port [esptool options...]\n";
}

if (!File::Spec->file_name_is_absolute($port)) {
    if (-c "/dev/$port") {
	$port = "/dev/$port";
    } elsif (-c "/dev/tty$port") {
	$port = "/dev/tty$port";
    } elsif ($^O eq 'MSWin32') {
	$port = "\\\\.\\$port";
    } else {
	die "$0: no such serial port: $port\n";
    }
}
print STDERR "Using serial port device $port\n";

open(my $fw, '<:gzip', $file)
    or die "$0: $file: $!\n";

my @chunks = ();

my $err = 0;

my $hdr;
while (read($fw, $hdr, 16) == 16) {
    # magic type flags data_len addr
    my @h = unpack('VvvVV', $hdr);
    my $c = { 'hdr' => $hdr, 'magic' => $h[0], 'type' => $h[1],
		  'flags' => $h[2], 'len' => $h[3], 'addr' => $h[4] };
    if ($c->{'magic'} != $FW_MAGIC) {
	print STDERR "$0: $file: bad chunk magic\n";
	$err = 1;
	last;
    }

    my $t = $type[$c->{'type'}];

    my $d;
    if (read($fw, $d, $c->{'len'}) != $c->{'len'}) {
	print STDERR "$0: $file: short chunk read\n";
	$err = 1;
	last;
    }
    $c->{'data'} = $d;
    push(@chunks, $c);

    last if ($t eq 'end'); # End of stream
}

close($fw);
exit $err if ($err);

my $td = File::Temp->newdir(CLEANUP => 1);

my @espfiles = ();
my %espopt = ('before' => 'default_reset', 'after' => 'hard_reset',
	      'baud' => 115200, 'port' => undef, 'chip' => undef,
	      'flash_mode' => undef, 'flash_freq' => undef,
	      'flash_size' => undef);

# Create a compressed data buffer without the ESP32 chunks
my $fpgadata;
open(my $fpgafh, '>:gzip', \$fpgadata) or die;

my $nc = 0;
foreach my $c ( @chunks ) {
    my $t = $type[$c->{'type'}];
    if ($t eq 'esptool') {
	my $s = $c->{'data'};
	$s =~ s/[\r\n]+/ /g;
	while ($s =~ /^\s*(\S+)\s+(\S+)(.*)/) {
	    $espopt{$1} = $2;
	    $s = $3;
	}
    } elsif ($t =~ /^esp/ && $c->{'addr'}) {
	my $addr = sprintf('%x', $c->{'addr'});
	my $tff = File::Spec->catfile($td, $t.$nc.'_'.$addr.'.bin');
	open(my $tf, '>', $tff) or die "$0: $tff: $!\n";
	print $tf $c->{'data'};
	close($tf);
	push(@espfiles, '0x'.$addr, $tff);
    } else {
	print $fpgafh $c->{'hdr'};
	print $fpgafh $c->{'data'};
    }
    $nc++;
}

close($fpgafh);

foreach my $e (keys %espopt) {
    my $ev = $ENV{"ESP\U$e"};
    if (defined($ev)) {
	$espopt{$e} = $ev;
    }
}

run_esptool($port, { hgrep sub {!/^flash_/}, %espopt },
	    'write_flash', { hgrep sub {/^flash_/}, %espopt },
	    '-z', @espfiles);

my $SerialPort = eval {
    require Device::SerialPort;
    Device::SerialPort->import(qw(:PARAM));
    'Device::SerialPort';
} || eval {
    require Win32::SerialPort;
    Win32::SerialPort->import(qw(:STAT));
    'Win32::SerialPort';
} || die "$0: need Device::SerialPort or Win32::SerialPort\n";

print STDERR "Waiting for reinit...\n";
usleep(4000000);

my $tty = $SerialPort->new($port);
die "$0: $port: $!\n" if (!$tty);

$tty->buffers($tty->buffer_max);
$tty->reset_error;

$tty->user_msg(1);
$tty->error_msg(1);

$tty->handshake('none');
$tty->databits(8);
$tty->baudrate(115200);
$tty->stopbits(1);
$tty->parity('none');
$tty->stty_istrip(0);
$tty->stty_inpck(0);
$tty->datatype('raw');
$tty->stty_icanon(0);
$tty->stty_opost(0);
$tty->stty_ignbrk(0);
$tty->stty_inlcr(0);
$tty->stty_igncr(0);
$tty->stty_icrnl(0);
$tty->stty_echo(0);
$tty->stty_echonl(0);
$tty->stty_isig(0);
$tty->read_char_time(0);
$tty->write_settings;

# In case DTR and/or RTS was asserted on the physical serial port.
# Note that DTR needs to be deasserted before RTS!
# However, deasserting DTR on the ACM port prevents it from working,
# so only do this if we don't see CTS (which is always enabled on ACM)...
if ($tty->can_modemlines && !($tty->modemlines & $tty->MS_CTS_ON)) {
    usleep(100000);
    $tty->dtr_active(0);
    usleep(100000);
    $tty->rts_active(0);
    usleep(100000);
}

sub tty_read {
    state %old_timeout;
    my($tty,$bufref,$timeout) = @_;
    my $d = $$bufref;

    if ($d eq '') {
	my $c;
	if (!defined($old_timeout{$tty}) || $timeout != $old_timeout{$tty}) {
	    $tty->read_const_time($timeout);
	    $old_timeout{$tty} = $timeout;
	}
	($c,$d) = $tty->read(256);
	return '' if (!$c);
    }

    my $r = substr($d,0,1);
    $$bufref = substr($d,1);

    return $r;
}
my $ttybuf = '';

sub tty_write($$) {
    my($tty,$data) = @_;
    my $bytes = length($data);
    my $offs  = 0;

    while ($bytes) {
	my $cnt = $tty->write(substr($data,$offs,$bytes));
	usleep(10000) unless ($cnt);
	$offs  += $cnt;
	$bytes -= $cnt;
    }

    return $offs;
}

my $found = 0;
my $tt = time();
my $start_enq = $tt;
tty_write($tty, "\005");	# ENQ
while (($tt = time()) - $start_enq < 30) {
    my $d = tty_read($tty, \$ttybuf, 1000);
    if ($d eq '') {
	tty_write($tty, "\005");	# ENQ
    } else {
	my $ix = index("\026\004\027", $d);
	if ($ix < 0) {
	    print STDERR $d;
	    next;
	}
	$found = $found == $ix ? $found+1 : 0;
	last if ($found == 3);
    }
}

if ($found < 3) {
    die "$0: $port: no MAX80 ESP32 detected\n";
}

$start_enq = $tt;
my $last_req;
my $winspc;
while (!defined($winspc)) {
    $tt = time();
    if ($tt - $start_enq >= 10) {
	die "$0: $port: failed to start FPGA firmware upload\n";
    }
    if ($tt != $last_req) {
	tty_write($tty, "\034\001: /// MAX80 FW UPLOAD \~\@\~ \$\r\n\035");
	$last_req = $tt;
    }
    my $d;
    while (1) {
	$d = tty_read($tty, \$ttybuf, 100);
	last if ($d eq '');
	my $dc = unpack('C', $d);
	print STDERR $a[$dc];
	if ($dc == 036) {
	    $winspc = 0;
	    last;
	}
    }
}

my $bytes  = length($fpgadata);
my $offset = 0;
my $maxchunk = 64;
my $maxahead = 256;
my $last_ack = 0;
my @pktends  = ();
my $last_enq = 0;
my $last_ack_time = 0;

print STDERR "\nStarting packet transmit...\n";

while ($last_ack < $bytes) {
    my $chunk;
    my $now;

    while (1) {
	$chunk = $bytes - $offset;
	$chunk = $winspc if ($chunk > $winspc);
	if ($offset + $chunk > $last_ack + $maxahead) {
	    $chunk = $last_ack + $maxahead - $offset;
	}
	$chunk = 0 if ($chunk <= 0);
	$chunk = $maxchunk if ($chunk > $maxchunk);

	my $d = tty_read($tty, \$ttybuf, $chunk ? 0 : $maxchunk/10);
	$now = time();
	last if ($d eq '');

	my $dc = unpack('C', $d);
	if ($dc == 022) {
	    $winspc = 0;
	} elsif ($dc == 024) {
	    if (defined($winspc)) {
		$winspc += 256;
	    }
	} elsif ($dc == 006) {
	    $last_ack = shift(@pktends) || $last_ack;
	    if ($now != $last_ack_time) {
		printf STDERR "%s: %s: %d/%d (%d%%) uploaded\n",
		    $0, $port, $last_ack, $bytes, $last_ack*100/$bytes;
		$last_ack_time = $now;
	    }
	} else {
	    print STDERR $a[$dc];
	    if ($dc == 025 || $dc == 031 || $dc == 037) {
		$offset = $last_ack;
		@pktends = ();
		undef $winspc;	# Wait for WRST before resuming
	    } elsif ($dc == 030) {
		print STDERR "\n";
		die "$0: $port: upload aborted by target system\n";
	    }
	}
    }

    if (!$chunk) {
	if ($bytes > $offset) {
	    if ($now != $last_enq) {
		tty_write($tty, "\026"); # SYN: request window resync
		$last_enq = $now;
	    }
	}
	next;
    }

    my $data = substr($fpgadata, $offset, $chunk);
    my $hdr = pack("CvVV", $chunk-1, 0, $offset, crc32($data));

    tty_write($tty, "\002".mybaseencode($hdr, $data));

    push(@pktends, $offset + $chunk);
    $offset += $chunk;
    $winspc -= $chunk;
}

tty_write($tty, "\004");		# EOT

# Final messages out
while (1) {
    my $d = tty_read($tty, \$ttybuf, 1000);
    last if ($d eq '');
    print STDERR $d;
}
$tty->close;

print STDERR "\n$0: $port: firmware upload complete\n";

exit 0;