#!/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);

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

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
    );
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;
}

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;
    }
}

$espopt{'port'} = $port;

my @espcmd = unquote_cmd($esptool);
foreach my $o (sort grep (!/^flash_/, keys %espopt)) {
    if (defined($espopt{$o})) {
	push(@espcmd, "--$o", $espopt{$o});
    }
}
push(@espcmd, 'write_flash', '-z');
foreach my $o (sort grep (/^flash_/, keys %espopt)) {
    if (defined($espopt{$o})) {
	push(@espcmd, "--$o", $espopt{$o});
    }
}

push(@espcmd, @espfiles);

print STDERR join(' ', @espcmd), "\n";
$err = system(@espcmd);
if ($err == -1) {
    print STDERR "$0: $espcmd[0]: $!\n";
} elsif ($err) {
    print STDERR "$0: $espcmd[0]: exit $err\n";
}
exit $err if ($err || $esponly);

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->parity('none');
$tty->stopbits(1);
$tty->datatype('raw');
$tty->stty_ignbrk(0);
#$tty->stty_brkint(0);
$tty->stty_parmrk(0);
$tty->stty_istrip(0);
$tty->stty_inlcr(0);
$tty->stty_igncr(0);
$tty->stty_icrnl(0);
$tty->stty_opost(0);
$tty->stty_echo(0);
$tty->stty_echonl(0);
$tty->stty_icanon(0);
$tty->stty_isig(0);
#$tty->stty_iexten(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 {
    my($tty,$bufref,$timeout) = @_;
    my $d = $$bufref;

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

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

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

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

    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 : 1000);
	last if ($d eq '');

	my $dc = unpack('C', $d);
	if ($dc == 022) {
	    $winspc = 0;
	    print STDERR "WRST, window: $winspc\n";
	} elsif ($dc == 024) {
	    if (defined($winspc)) {
		$winspc += 256;
		print STDERR "WGO, window: $winspc\n";
	    }
	} elsif ($dc == 006) {
	    $last_ack = shift(@pktends) || $last_ack;
	    print STDERR "ACK to $last_ack\n";
	} else {
	    print STDERR $a[$dc];
	    if ($dc == 025 || $dc == 031 || $dc == 037) {
		print STDERR "\n$0: $port: resetting upload to $last_ack\n";
		$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";
	    }
	}
    }

    print STDERR "Send offset: $offset, window: $winspc, chunk: $chunk\n";
    if (!$chunk) {
	if ($bytes > $offset) {
	    tty_write($tty, "\026"); # Request to resynchronize credits
	}
	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
$tty->close;

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

exit 0;