#!/usr/bin/perl

use strict;
use integer;
use PerlIO::gzip;
use File::Temp;
use File::Spec;

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 $file = shift(@args);

if (!defined($file)) {
    die "Usage: $0 fwfile [esptool options...]\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'}];
    last if ($t eq 'end'); # End of stream

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

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

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);
    }
    $nc++;
}

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

my @espcmd = unquote_cmd($esptool);
foreach my $o (sort grep (!/^flash_/, keys %espopt)) {
    if (defined($espopt{$o})) {
	push(@espcmd, "--$o", $espopt{$o});
    }
}
push(@espcmd, @args);
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);

if (defined($ENV{'VERBOSE'})) {
    print STDERR join(' ', @espcmd), "\n";
}

$err = system(@espcmd);
if ($err == -1) {
    print STDERR "$0: $espcmd[0]: $!\n";
}

exit !!$err;