#!/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-2025 -- leonerd@leonerd.org.uk

use v5.20;
use warnings;

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

use Device::AVR::UPDI;
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;

my %configvars;
if( -e ( my $config = "$ENV{HOME}/.avrupdirc" ) ) {
   open my $fh, "<", $config or die "Cannot read $config - $!\n";
   while( <$fh> ) {
      next if m/^\s*#/;
      chomp;
      m/^\s*(\w+)\s*=\s*(.*)\s*$/ and
         $configvars{$1} = $2;
   }
}

our $PORT
   :GlobalOption("port|P=",
                 "Path to the device where the UPDI adapter is connected")
   = join ",", grep { defined } $configvars{port}, "/dev/ttyUSB0";
if( $PORT =~ m/,/ ) {
   my $chosen_port = first { -e $_ } map { glob $_ } split m/,/, $PORT;
   defined $chosen_port or
      die "Cannot find a UPDI port out of $PORT\n";
   $PORT = $chosen_port;
}

our $PART
   :GlobalOption("part|p=",
                 "Name of the ATmega or ATtiny chip that is expected");

our $BAUD
   :GlobalOption("baud|B=u",
                 "Sets a different baud rate for communications")
   = $configvars{baud} // 115200;

our $ERASE
   :GlobalOption("erase|e",
                 "Send the CHIPERASE key as well as the NVMPROG key");

our $BINARY
   :GlobalOption("binary",
                 "Files read from or written to will be in raw binary format");

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, $baseaddr ) = @_;
   $baseaddr //= 0;

   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 to %06X\n", $size, $baseaddr;

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

      $updi->$writemethod( $baseaddr + $addr, $page )->get;

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

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

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

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

      $addr += length $page;
      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;

   check_signature;
}

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

   # Hard to detect this for real but in practice just trimming the 0xFF bytes
   # off the end is good enough
   $data =~ s/\xFF+$//;

   printf "Trimmed to %d bytes\n", length $data;

   $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 ) = @_;

   my @chunks = open_input_file( $file )->input_chunks;

   make_updi;

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

   check_signature; # TODO: Check if enable_nvmprog logic is correct

   foreach my $chunk ( @chunks ) {
      my ( $baseaddr, $data ) = @$chunk;

      # Some UPDI chips seem to have trouble writing an odd number of bytes.
      # We'll pad the data with an extra 0xFF if so
      $data .= "\xFF" if length($data) % 2;

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

=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 ) = @_;

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

   make_updi;

   check_signature;

   _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>;
   }
   sub input_chunks {
      my $self = shift;
      return [ 0, $self->input ];
   }
}

package main::Fileformat::IntelHex {
   use base 'main::Fileformat';

   # Because this file format was invented for Intel x86 in the 16-bit era,
   # we have to account for its crazy 16bit offset + 16bit segment == 20bit
   # linear addressing madness.

   sub print_record {
      my $self = shift;
      my ( $addr, $type, $data ) = @_;
      my $clen = length $data;
      my $cksum = $clen + ( $addr & 0xFF ) + ( $addr >> 8 ) + $type;
      my $encdata = "";
      ( $encdata .= sprintf "%02X", ord $_), ( $cksum += ord $_ )
         for split //, $data;
      $self->printf( ":%02X%04X%02X%s%02X\n",
         $clen, $addr, $type, $encdata, ( -$cksum ) & 0xFF );
   }
   sub output_chunk {
      my $self = shift;
      my ( $addr, $bytes ) = @_;

      my $segaddr = 0;

      foreach my $data ( $bytes =~ m/(.{1,16})/gs ) {
         my $hiaddr = $addr & 0xFFFF0000;
         if( $segaddr << 4 != $hiaddr ) {
            $segaddr = $hiaddr >> 4;
            $self->print_record( 0, 0x02, pack "S>", $segaddr );
         }
         $self->print_record( $addr & 0xFFFF, 0x00, $data );
         $addr += length $data;
      }
   }
   sub output {
      my $self = shift;
      my ( $bytes ) = @_;

      my $addr = 0;
      my $chunk = "";

      foreach my $data ( $bytes =~ m/(.{1,16})/gs ) {
         if( $data !~ m/[^\xff]/ ) {
            # data contains only 0xFF bytes.

            if( length $chunk ) {
               # Write the current chunk then reset it

               my $chunkaddr = $addr;
               $addr += length $chunk;

               $chunk =~ s/\xFF+$//;
               $self->output_chunk( $chunkaddr, $chunk );

               $chunk = "";
            }

            $addr += length $data;
         }
         else {
            $chunk .= $data;
         }
      }

      if( length $chunk ) {
         $self->output_chunk( $addr, $chunk );
      }

      $self->print_record( 0, 0x01, "" );
   }

   sub input {
      my $self = shift;

      my $bytes = "";
      foreach my $chunk ( $self->input_chunks ) {
         my ( $addr, $chunkbytes ) = @$chunk;

         # TODO: this presumes chunks arrive in address order with no overlaps
         # but maybe gaps
         $bytes .= "\xFF" x ( $addr - length $bytes );
         $bytes .= $chunkbytes;
      }
      return $bytes;
   }
   sub input_chunks {
      my $self = shift;
      my @chunks; my $lastchunk;
      my $segaddr = 0;
      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;
         warn "Bad record length on line $.\n" and next if
            length $data != 2 * $clen;
         if( $type == 0x01 ) {
            # EOF
            last;
         }
         elsif( $type == 0x02 ) {
            # Extended Segment Address; sets upper bits
            warn "Expected 16 bits of Extended Segment Address on line $.\n" and next
               if length $data != 4;
            $segaddr = hex $data;
            next;
         }
         elsif( $type != 0 ) {
            warn sprintf "TODO: unrecognised input record type 0x%02X on line %d\n",
               $type, $.;
            next;
         }
         $addr |= ( $segaddr << 4 );
         $data = pack "H*", $data;
         if( $lastchunk and ( $lastchunk->[0] + length $lastchunk->[1] ) == $addr ) {
            # merge
            $lastchunk->[1] .= $data;
         }
         else {
            push @chunks, $lastchunk = [ $addr, $data ];
         }
      }
      return @chunks;
   }
}

__END__

=head1 CONFIG FILE

If the file F<$HOME/.avrupdirc> exists, it is opened and read as lines of
C<name = value> declarations, to define some default configuration. The
following variables are recognised. Lines beginning with C<#> are ignored as
comments.

=over 4

=item C<baud>

Overrides the baud rate.

=item C<port>

Overrides the path to the UPDI device. This value may consist of multiple
paths separated by commas. The first path in the list that names a device
which actually exists is used.

=back

=head1 AUTHOR

Paul Evans <leonerd@leonerd.org.uk>
