#!/usr/bin/perl
#  You may distribute under the terms of either the GNU General Public License
#  or the Artistic License (the same terms as Perl itself)
#
#  (C) Paul Evans, 2019-2022 -- leonerd@leonerd.org.uk

use v5.20;
use warnings;

use Commandable;
use Commandable::Finder::SubAttributes 0.08 ':attrs';

use Device::AVR::UPDI;
use Getopt::Long;
use List::Util qw( first );

=head1 NAME

F<avr-updi> - communicate with an F<AVR> microcontroller over F<UPDI>

=head1 SYNOPSIS

   $ avr-updi [--port PORT] [--part PART] COMMAND ARGS... [COMMAND ARGS ...]

=head1 DESCRIPTION

This program contains a number of sub-commands for interacting with an F<AVR>
microcontroller using a F<UPDI> interface.

=head1 COMMON OPTIONS

=head2 --port, -P DEVICE

Optional. Provides the USB device where the UPDI adapter is connected. If
absent a default of F</dev/ttyUSB0> will apply.

=head2 --part, -p PART

Required. Gives the name of the ATmega or ATtiny chip that is expected. Parts
may be specified in the following ways:

   ATmega4809
   atmega4809
   m4809

   ATtiny814
   attiny814
   t814

Specifically, these are the same forms as recognised by F<avr-gcc>'s C<-mmcu>
option and F<avrdude>'s C<-p> option, for convenience in Makefiles and build
scripts.

=head2 --baud, -B BAUD

Optional. Sets a different baud rate for communications. If not supplied, will
default to 115200.

If communications are unreliable, try setting a slower speed.

=head2 --erase, -e

Send the CHIPERASE key as well as the NVMPROG key when enabling programming
for the C<write-fuses> operation. Normally this flag is not required when
writing fuses to a newly-programmed chip, but may be necessary to recover from
a bad checksum or bad fuse value.

=head2 --binary

Files read from or written to will be in raw binary format, instead of Intel
hex.

=cut

STDOUT->autoflush;

GetOptions(
   'port|P=s' => \( my $PORT = "/dev/ttyUSB0" ),
   'part|p=s' => \( my $PART ),
   'baud|B=i' => \( my $BAUD ),

   'erase|e' => \( my $ERASE ),

   'binary'   => \( my $BINARY ),
) or exit 1;

my $updi;
sub make_updi
{
   $updi = Device::AVR::UPDI->new(
      dev => $PORT,
      part => $PART,
      baud => $BAUD,
   );

   $updi->init_link->get;
}

sub check_signature
{
   unless( $updi->read_asi_sys_status->get & (1<<3) ) {
      $updi->enable_nvmprog->get;
   }

   my $sig = $updi->read_signature->get;
   if( $updi->partinfo->signature ne $sig ) {
      printf STDERR "Signature %v02X does not match expected %v02X\n",
         $sig, $updi->partinfo->signature;
      exit 1;
   }
}

sub open_input_file
{
   my ( $file ) = @_;

   if( $BINARY ) {
      return main::Fileformat::Binary->open_in( $file );
   }
   else {
      return main::Fileformat::IntelHex->open_in( $file );
   }
}

sub open_output_file
{
   my ( $file ) = @_;

   if( $BINARY ) {
      return main::Fileformat::Binary->open_out( $file );
   }
   else {
      return main::Fileformat::IntelHex->open_out( $file );
   }
}

sub reset_chip
{
   $updi->request_reset( 1 )->get;
   $updi->request_reset( 0 )->get;
}

my $RESET_REQUIRED;

Commandable::Finder::SubAttributes->new_for_main
   ->configure( allow_multiple_commands => 1 )
   ->find_and_invoke_ARGV;

reset_chip if $RESET_REQUIRED;

=head1 SUBCOMMANDS

Multiple commands may be specified simultaneously; they will be executed in
the given sequence. This may be useful for writing flash and EEPROM at the
same time

   $ avr-updi -p t814 write-flash flash.hex write-eeprom eeprom.hex

=cut

sub _generic_read
{
   my ( $type, $pagesize, $size ) = @_;

   my $readmethod = "read_${type}_page";

   my $data = "";
   my $addr = 0;
   while( $addr < $size ) {
      $data .= $updi->$readmethod( $addr, $pagesize )->get;

      $addr += $pagesize;
      printf STDERR "\rRead %d (%.2f%%)", $addr, 100 * $addr / $size;
   }
   print STDERR "\n";
   print STDERR "Done\n";

   return $data;
}

sub _generic_write
{
   my ( $type, $data, $pagesize, $size ) = @_;

   my $writemethod = "write_${type}_page";
   my $readmethod  = "read_${type}_page";

   die "Too much data to write in $size bytes\n"
      if $size < length $data;

   $size = length $data;

   printf STDERR "Need to write %d bytes\n", $size;

   my $addr = 0;
   while( $addr < $size ) {
      my $chunk = substr $data, $addr, $pagesize;

      $updi->$writemethod( $addr, $chunk )->get;

      $addr += length $chunk;
      printf STDERR "\rWritten %d (%.2f%%)", $addr, 100 * $addr / $size;
   }
   print STDERR "\n";

   # Verify
   $addr = 0;
   while( $addr < $size ) {
      my $chunk = substr $data, $addr, $pagesize;

      my $readback = $updi->$readmethod( $addr, length $chunk )->get;
      if( $readback ne $chunk ) {
         my $offset = 0;
         $offset++ while substr( $readback, $offset, 1 ) eq substr( $chunk, $offset, 1 );

         printf STDERR "Read: %v02X\n", $readback;
         printf STDERR "Exp : %v02X\n", $chunk;
         die sprintf "Verify failed at addr=0x%04X\n", $addr + $offset;
      }

      $addr += length $chunk;
      printf STDERR "\rVerified %d (%.2f%%)", $addr, 100 * $addr / $size;
   }
   print STDERR "\n";

   $RESET_REQUIRED++;

   print STDERR "Done\n";
}

sub _print_decoded_fields
{
   my ( $reg, @fields ) = @_;

   while( @fields ) {
      my $name = shift @fields;
      my $mask = shift @fields;
      my $val = $reg & $mask;

      # shift down
      $mask >>= 1, $val >>= 1 while !( $mask & 1 );

      if( $mask > 1 ) {
         printf " %s=%X", $name, $val;
      }
      elsif( $val ) {
         printf " %s", $name;
      }
      else {
         printf " !%s", lc $name;
      }
   }
   print "\n";
}

# Debugging; undocumented
sub run_show_updi_regs
{
   make_updi;

   my @regs = map { $updi->ldcs( $_ )->get } 0 .. 0x0C;

   printf "UPDI.STATUSA        = %02X", $regs[0];
   _print_decoded_fields $regs[0], UPDIREV => 0xF0;
   printf "     STATUSB        = %02X", $regs[1];
   _print_decoded_fields $regs[1], PSIG => 0x0F;
   printf "     CTRLA          = %02X", $regs[2];
   _print_decoded_fields $regs[2], IBDLY => (1<<7), PARD => (1<<5), DTD => (1<<4), RSD => (1<<3), GTVAL => 0x07;
   printf "     CTRLB          = %02X", $regs[3];
   _print_decoded_fields $regs[3], NACKDIS => (1<<4), CCDETDIS => (1<<3), UPDIDIS => (1<<2);

   # [4] to [6] reserved
   #
   printf "     ASI_KEY_STATUS = %02X", $regs[7];
   _print_decoded_fields $regs[7], UROWWRITE => (1<<5), NVMPROG => (1<<4), CHIPERASE => (1<<3);
   printf "     ASI_RESET_REQ  = %02X\n", $regs[8];

   printf "     ASI_CTRLA      = %02X", $regs[9];
   _print_decoded_fields $regs[9], UPDICLKDIV => 0x03;
   printf "     ASI_SYS_CTRLA  = %02X", $regs[10];
   _print_decoded_fields $regs[10], UROWDONE => (1<<1), CLKREQ => (1<<0);
   printf "     ASI_SYS_STATUS = %02X", $regs[11];
   _print_decoded_fields $regs[11], ERASE_FAILED => (1<<6), SYSRST => (1<<5), INSLEEP => (1<<4),
      NVMPROG => (1<<3), UROWPROG => (1<<2), LOCKSTATUS => (1<<0);
   printf "     ASI_CRC_STATUS = %02X", $regs[12];
   _print_decoded_fields $regs[12], CRC_STATUS => 0x07;
}

=head2 reset

Sends a reset request.

   $ avr-updi reset

Note that this command does not need the C<--part> to be specified.

=cut

sub command_reset
   :Command_description("Sends a reset request")
{
   $PART //= "ATtiny814"; # part doesn't matter just for reset
   make_updi;

   print STDERR "Issuing Reset over UPDI\n";
   reset_chip;
}

=head2 read-sib

Reads the System Information Block

   $ avr-updi read-sib

Note that this command does not need the C<--part> to be specified.

=cut

sub command_read_sib
   :Command_description("Reads the System Information Block")
{
   $PART //= "ATtiny814"; # part doesn't matter
   make_updi;

   my $sib = $updi->read_sib->get;

   foreach my $key ( sort keys %$sib ) {
      printf "%12s: %s\n", uc $key, $sib->{$key};
   }
}

=head2 read-signature

Reads the 3-byte signature from the chip and prints it.

   $ avr-updi read-signature -p t814

=cut

sub command_read_signature
   :Command_description("Reads the 3-byte signature from the chip and prints it")
{
   make_updi;

   unless( $updi->read_asi_sys_status->get & (1<<3) ) {
      $updi->enable_nvmprog->get;
   }

   my $sig = $updi->read_signature->get;

   printf "Signature: %v02X\n", $sig;
}

=head2 erase

Erases the entire chip. Normally not required because a C<write-flash>
operation will do this anyway, but this command may be handy for erasing a
chip to reset it back to factory state.

   $ avr-updi erase -p t814

=cut

sub command_erase
   :Command_description("Erases the entire chip")
{
   make_updi;

   $updi->erase_chip( no_reset => 1 )->get;
}

=head2 read-flash

Reads from the flash portion of non-volatile memory ("NVM").

   $ avr-updi read-flash -p t814 flash-save.hex [LENGTH]

Optionally reads only the given length.

=cut

sub command_read_flash
   :Command_description("Reads from the flash portion of non-volatile memory (\"NVM\")")
   :Command_arg("file", "path to file data will be written to")
   :Command_arg("length?", "optional length to read")
{
   my ( $file, $readlen ) = @_;

   make_updi;
   check_signature;

   my $outfile = open_output_file( $file );

   my $length = $updi->partinfo->size_flash;

   if( defined $readlen and $readlen =~ m/^\d/ ) {
      $length = $readlen;
      $length = hex $length if $length =~ m/^0/;
   }

   my $data = _generic_read( flash =>
      $updi->partinfo->pagesize_flash, $length );

   $outfile->output( $data );
}

=head2 write-flash

Writes to the flash portion of non-volatile memory ("NVM").

   $ avr-updi write-flash -p t814 firmware.hex

=cut

sub command_write_flash
   :Command_description("Writes to the flash portion of non-volatile memory (\"NVM\")")
   :Command_arg("file", "path to file data will be read from")
{
   my ( $file ) = @_;

   make_updi;

   $updi->erase_chip( no_reset => 1 )->get;

   check_signature; # TODO: Check if enable_nvmprog logic is correct

   my $data = open_input_file( $file )->input;

   _generic_write( flash => $data,
      $updi->partinfo->pagesize_flash, $updi->partinfo->size_flash );
}

=head2 read-eeprom

Reads from the EEPROM portion of non-volatile memory ("NVM").

   $ avr-updi read-eeprom -p t814 eeprom-save.hex

=cut

sub command_read_eeprom
   :Command_description("Reads from the EEPROM portion of non-volatile memory (\"NVM\")")
   :Command_arg("file", "path to file data will be written to")
{
   my ( $file ) = @_;

   make_updi;
   check_signature;

   my $outfile = open_output_file( $file );

   my $data = _generic_read( eeprom =>
      $updi->partinfo->pagesize_eeprom, $updi->partinfo->size_eeprom );

   $outfile->output( $data );
}

=head2 write-eeprom

Writes to the EEPROM portion of non-volatile memory ("NVM").

   $ avr-updi write-eeprom -p t814 data.hex

=cut

sub command_write_eeprom
   :Command_description("Writes to the EEPROM portion of non-volatile memory (\"NVM\")")
   :Command_arg("file", "path to file data will be read from")
{
   my ( $file ) = @_;

   make_updi;

   check_signature;

   my $data = open_input_file( $file )->input;

   _generic_write( eeprom => $data,
      $updi->partinfo->pagesize_eeprom, $updi->partinfo->size_eeprom );
}

=head2 read-fuses

Reads fuse values.

   $ avr-updi read-fuses -p t814
   WDTCFG : 00
     WDTCFG.PERIOD    : OFF
     WDTCFG.WINDOW    : OFF
   BODCFG : 00
     BODCFG.ACTIVE    : DIS
     BODCFG.LVL       : BODLEVEL0
     BODCFG.SAMPFREQ  : 1KHz
     BODCFG.SLEEP     : DIS
   OSCCFG : 02
   ...

Decoded values of fields are also printed.

=cut

sub command_read_fuses
   :Command_description("Reads fuse values")
{
   make_updi;
   check_signature;

   my @FUSES = @{ $updi->partinfo->fusenames };
   my $fuseinfo = $updi->fuseinfo;

   foreach my $idx ( 0 .. $#FUSES ) {
      my $regname = $FUSES[$idx];
      next if !defined $regname;

      my $regvalue = $updi->read_fuse( $idx )->get;
      printf "%-7s: %02X\n", $regname, $regvalue;

      foreach my $fusename ( @{ $fuseinfo->{fuses} } ) {
         my $fuse = $fuseinfo->{fusemap}{$fusename};
         next unless $fuse->{reg} eq $regname;

         my $val = $regvalue & $fuse->{mask};

         if( $fuse->{values} ) {
            my $chosen = first { $_->{value} == $val } $fuse->{values}->@*;
            printf "  %s.%-10s: %s\n", $regname, $fusename, $chosen->{name};
         }
         else {
            printf "  %s.%-10s: %d\n", $regname, $fusename, $val ? 1 : 0;
         }
      }
   }
}

=head2 write-fuses

Writes fuse values.

   $ avr-updi write-fuses -p t814 BODCFG=02 SYSCFG0.RSTPINCFG=GPIO

Fuses may be specified as numerical values for entire registers, or symbolic
names for individual fields.

=cut

sub command_write_fuses
   :Command_description("Writes fuse values")
   :Command_arg("fuses...", "the fuses to set (as NAME=value ...)")
{
   my ( $fuseargs ) = @_;

   make_updi;

   if( $ERASE ) {
      print STDERR "Erasing chip...\n";
      $updi->erase_chip( no_reset => 1 )->get;
   }

   check_signature;

   my @FUSES = @{ $updi->partinfo->fusenames };
   my $fuseinfo = $updi->fuseinfo;

   my @values;

   while( @$fuseargs ) {
      my ( $name, $value ) = split m/=/, $fuseargs->[0] or last;
      shift @$fuseargs;

      if( $name =~ m/\./ ) {
         my ( $regname, $fusename ) = $name =~ m/^(.*?)\.(.*)$/;
         my $idx = first { defined $FUSES[$_] and $FUSES[$_] eq $regname } 0 .. $#FUSES;
         defined $idx or die "Unrecognised fuse register $regname";

         my $fuse = $fuseinfo->{fusemap}{$fusename} or
            die "Unrecognised fuse $fusename";
         $fuse->{reg} eq $regname or
            die "Fuse $fusename does not live in register $regname";

         my $mask = 0+$fuse->{mask};

         if( $fuse->{values} ) {
            my $chosen = first { $_->{name} eq $value } $fuse->{values}->@* or
               die "Fuse $fusename does not have a value $value";
            $value = $chosen->{value};
         }
         else {
            $value = $value ? $mask : 0;
         }

         $values[$idx] //= $updi->read_fuse( $idx )->get;

         $values[$idx] &= ~$mask;
         $values[$idx] |= $value;
      }
      else {
         my $idx = first { defined $FUSES[$_] and $FUSES[$_] eq $name } 0 .. $#FUSES;
         defined $idx or die "Unrecognised fuse";

         $value = hex $value if $value =~ m/^0x/;

         $values[$idx] = $value;
      }
   }

   foreach my $idx ( 0 .. $#values ) {
      next unless defined( my $value = $values[$idx] );
      printf "Setting %s to 0x%02X\n", $FUSES[$idx], $value;

      $updi->write_fuse( $idx, $value )->get;
   }

   $RESET_REQUIRED++;
}

# IO formats
package main::Fileformat {
   use base 'IO::Handle';
   sub open_out {
      my $class = shift;
      my $fh;
      if( $_[0] eq "-" ) {
         $fh = IO::Handle->new_from_fd( STDOUT->fileno, "w" );
      }
      else {
         open $fh, ">", $_[0] or die "Cannot write $_[0] - $!\n";
      }
      return bless $fh, $class;
   }
   sub open_in {
      my $class = shift;
      my $fh;
      if( $_[0] eq "-" ) {
         $fh = IO::Handle->new_from_fd( STDIN->fileno, "r" );
      }
      else {
         open $fh, "<", $_[0] or die "Cannot read $_[0] - $!\n";
      }
      return bless $fh, $class;
   }
}

package main::Fileformat::Binary {
   use base 'main::Fileformat';
   sub output {
      my $self = shift;
      my ( $bytes ) = @_;
      $self->print( $bytes );
   }
   sub input {
      my $self = shift;
      local $/;
      return scalar <$self>;
   }
}

package main::Fileformat::IntelHex {
   use base 'main::Fileformat';
   sub output {
      my $self = shift;
      my ( $bytes ) = @_;
      my $addr = 0;
      foreach my $chunk ( $bytes =~ m/(.{1,16})/gs ) {
         my $clen = length $chunk;
         my $cksum = $clen + ( $addr & 0xff ) + ( $addr >> 8 );
         $self->printf( ":%02X%04X%02X", $clen, $addr, 0 );
         foreach my $byte ( split //, $chunk ) {
            $byte = ord $byte;
            $cksum += $byte;
            $self->printf( "%02X", $byte );
         }
         $self->printf( "%02X\n", ( -$cksum ) & 0xff );
         $addr += $clen;
      }
      $self->print( ":00000001FF\n" );
   }
   sub input {
      my $self = shift;
      my $bytes = "";
      while( my $line = <$self> ) {
         $line =~ s/\r?\n$//; # chomp doesn't do CRLF on Linux
         next unless my ( $clen, $addr, $type, $data, $cksum ) =
            $line =~ m/^:([0-9a-f]{2})([0-9a-f]{4})([0-9a-f]{2})([0-9a-f]*)([0-9a-f]{2})$/i;
         # TODO: check checksum
         $_ = hex $_ for $clen, $addr, $type;
         last if $type == 1; # EOF
         next if $type != 0; # unrecognised record
         warn "Bad record length on line $.\n" and next if
            length $data != 2 * $clen;
         $data = pack "H*", $data;
         if( $addr > length $bytes ) {
            # TODO: at write time this could be optimised away after chip erase
            $bytes .= "\xFF" x ( $addr - length $bytes );
         }
         substr( $bytes, $addr, length $data ) = $data;
      }
      return $bytes;
   }
}

__END__

=head1 AUTHOR

Paul Evans <leonerd@leonerd.org.uk>