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

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;
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;
    do {
	$d = tty_read($tty, \$ttybuf, 1000);
	if ($d eq "\036") {
	    $winspc = 0;
	    last;
	} else {
	    print $d;
	}
    } while ($d ne '');
}

my $bytes  = length($fpgadata);
my $offset = 0;
while ($offset < $bytes) {
    my $chunk = $bytes - $offset;
    $chunk = 64 if ($chunk > 64);
    $chunk = $winspc if ($chunk > $winspc);

    my $d = tty_read($tty, \$ttybuf, $chunk ? undef : 1000);
    if ($d ne '') {
	if ($d eq "\022") {
	    $winspc = 0;
	} elsif ($d eq "\024") {
	    $winspc += 256;
	} elsif ($d eq "\027" || $d eq "\030" || $d eq "\025") {
	    die "$0: $port: upload terminated by MAX80\n";
	} elsif ($d ne "\006") {
	    print STDERR $d;
	}
    }

    next unless($chunk);

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

    tty_write($tty, $hdr.$data);

    $offset += $chunk;
    $winspc -= $chunk;
}

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

exit 0;