#!/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 = 'esptool.py';
$esptool = $ENV{'ESPTOOL'} || $esptool;

my $esp_retries = 10;

my @FW_MAGIC = (undef, 0x7a07fbd6, 0xa924ed0b);

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 $FDF_PRETARGET = 0x0002;

my $STRING_MAX_LEN = 4095;

my $boardinfo_addr = 0;
my $boardinfo_len  = 4096;

# 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 match_version($$) {
    my($ver,$pattern) = @_;

    return 1;			# FIX THIS
}

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) {
	    @output = ();
	    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;
}

sub target_string_valid($)
{
    my($v) = @_;

    return $v =~ /^(\S+) v((?:0|[1-9][0-9]*)(?:\.0|\.[1-9][0-9]*)*)(?: ([a-zA-Z0-9]*))?$/;
}

sub match_version($$)
{
    my($version,$pattern) = @_;

    my @vv = target_string_valid($version);

    return 0 unless (defined($vv[0]));
    
    my $v_board = $vv[0];
    my @v_ver   = split(/\./, $vv[1]);
    my $v_flags = $vv[2];

    return 1 if ($pattern eq $v_board); # Board only matchall pattern
    
    if ($pattern !~ /^(\S+) v((?:0|[1-9][0-9]*)(?:\.\.0|\.[1-9][0-9]*)*)(?:(\-)((?:0|[1-9][0-9]*)(?:\.\.0|\.[1-9][0-9]*)*))?(?: ([\+\-a-zA-Z0-9]*))?$/) {
	return 0;
    }

    return 0 if ($1 ne $v_board);

    my @p_min = split(/\./, $2);
    my @p_max = split(/\./, $3 eq '' ? $2 : $4);
    my $p_flags = $5;

    while (scalar(@p_min) || scalar(@p_max)) {
	my $mi = shift(@p_min);
	my $ma = shift(@p_max);
	my $ve = shift(@v_ver) || 0;

	return 0 if (defined($mi) && $ve < $mi);
	return 0 if (defined($ma) && $ve > $ma);
    }

    my $flag_pol = 1;

    foreach my $c (split(//, $p_flags)) {
	if ($c eq '-') {
	    $flag_pol = 0;
	} elsif ($c eq '+') {
	    $flag_pol = 1;
	} else {
	    return 0 if ((index($v_flags, $c) != -1) != $flag_pol);
	}
    }

    return 1;
}
    
sub get_target_board($$)
{
    my($port,$boardinfo) = @_;

    run_esptool($port, { 'before' => 'default_reset', 'after' => 'hard_reset' },
		'read_flash', { },
		''.$boardinfo_addr, ''.$boardinfo_len, $boardinfo);

    open(my $bi, '<:raw', $boardinfo) or return undef;
    my $bid;
    read($bi, $bid, $boardinfo_len);
    close($bi);
    return undef if (length($bid) != $boardinfo_len);

    my @bh = unpack('VVVVZ[256]', $bid);
    if ($bh[0] == 0x6682df97 && $bh[1] == 0xe2a0d506) {
	# Standard format board_info structure
	substr($bid, 12, 4) = "\0\0\0\0"; # Clear the CRC field
	if ($bh[2] >= 16 && $bh[2] <= $boardinfo_len &&
	    crc32(substr($bid, 0, $bh[2])) == $bh[3]) {
	    return $bh[4];
	}
    } elsif ($bid =~ /^([[:print:]]+)\0/) {
	# Preliminary board_info (only version string)
	return $1;
    }

    return undef;
}
    
my @args = @ARGV;
my $esponly = 0;
my $file;
my $target_board = undef;
my $setver = 0;

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

    if ($file eq '--esponly') {
	$esponly = 1;
    } elsif ($file eq '--setver') {
	$target_board = shift(@args);
	$setver = defined($target_board);
    } 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 [--esponly][--setver versionoptions] 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";

my $td = File::Temp->newdir(CLEANUP => 1);
my $boardinfo = File::Spec->catfile($td, 'boardinfo.bin');

my @espfiles = ();

if (defined($target_board)) {
    # Forcibly write target board version
    if (!target_string_valid($target_board)) {
	die "$0: $port: invalid firmware target string: $target_board\n";
    }

    open(my $bi, '>:raw', $boardinfo)
	or die "$0: $port: $boardinfo: $!\n";
    my $bid = $target_board . "\0";
    $bid .= "\xff" x ($boardinfo_len - length($bid));
    print $bi $bid;
    close($bi);
    push(@espfiles, ''.$boardinfo_addr, $boardinfo);
} else {
    # Get the board version from target flash
    
    $target_board = get_target_board($port, $boardinfo);
    if (!defined($target_board) || !target_string_valid($target_board)) {
	die "$0: $port: board version not programmed, specify with --setver\n";
    }
}

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

my @chunks = ();
my $version_match = undef;

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],
	      'vmatch' => 0, 'vmask' => 0, 'vmin' => 0, 'vmax' => 0xffff };
    if ($c->{'magic'} == $FW_MAGIC[2]) {
	if (read($fw, $hdr, 16) != 16) {
	    print STDERR "$0: $file: short header read\n";
	    $err = 1;
	    last;
	}
	my @h = unpack('VVvvV', $hdr);
	$c->{'hdr'}   .= $hdr;
	$c->{'vmatch'} = $h[0];
	$c->{'vmask'}  = $h[1];
	$c->{'vmin'}   = $h[2];
	$c->{'vmax'}   = $h[3];
    } elsif ($c->{'magic'} != $FW_MAGIC[1]) {
	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;

    if ($t eq 'target') {
	my $is_match = match_version($target_board, $d);
	if ($c->{'magic'} == $FW_MAGIC[1] || $is_match) {
	    $version_match = $c;
	}
	print STDERR "$0: $file: supports hardware: $d",
	    ($is_match ? ' (match)' : ''), "\n";
    } elsif ($t eq 'end' || $t eq 'note' || ($c->{'flags'} & $FDF_PRETARGET)) {
	# Not filtered
    } else {
	if (!defined($version_match)) {
	    print STDERR "$0: $file: hardware version $target_board not supported\n";
	    $err = 1;
	    last;
	}
	
	if ($c->{'vmin'} > $version_match->{'vmax'} ||
	    $c->{'vmax'} < $version_match->{'vmin'} ||
	    (($c->{'vmatch'} ^ $version_match->{'vmatch'}) & $c->{'vmask'})) {
	    # Not applicable to this board
	    next;
	}
    }
    
    push(@chunks, $c);

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

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

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);
    } elsif ($t eq 'note' || ($t eq 'target' && $c != $version_match)) {
	# Skip
    } 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 {!/^flash_/} %espopt },
	    'write_flash', { hgrep {/^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;