#!/usr/bin/perl
# LiNXT
# Allows basic functionality with the NXT brick
# By jacques and Ben Slote <bslote@gmail.com>
# irc.freenode.net #nxthacks
# This project is released under the Perl Artistic License

use Device::USB;
use Getopt::Long;
use strict;
use warnings;

# Global USB stuff
my $USB_ID_VENDOR_LEGO = 0x0694;
my $USB_ID_PRODUCT_NXT = 0x0002;

my $USB_INTERFACE = 0;

my $USB_OUT_ENDPOINT = 0x01;
my $USB_IN_ENDPOINT  = 0x82;

my $USB_TIMEOUT = 1000;

my $DIRECT_COMMAND = 0x00;
my $SYSTEM_COMMAND = 0x01;

my $RESPONSE       = 0x00;
my $NO_RESP0NSE    = 0x80;

my $STATUS_SUCCESS	= 0x00;

# System commands
my $GET_FIRMWARE_VERSION_COMMAND = 0x88;
my $GET_DEVICE_INFO_COMMAND      = 0x9b;
my $GET_BATTERY_LEVEL_COMMAND    = 0x0b;
my $CLOSE_COMMAND                = 0x84;
my $FIND_FIRST_FILE_COMMAND      = 0x86;
my $FIND_NEXT_FILE_COMMAND       = 0x87;
my $OPEN_READ_COMMAND            = 0x80;
my $READ_COMMAND                 = 0x82;
my $OPEN_WRITE_COMMAND		 = 0x8B;
my $WRITE_COMMAND                = 0x83;
my $SET_NAME_COMMAND		 = 0x98;

# Handle command line options
my %opts = ();
if ($#ARGV > -1)
{
	GetOptions(
					'help|h'	=> \$opts{help},
					'battery|b'	=> \$opts{battery},
					'download|d'	=> \$opts{download},
					'info|i'	=> \$opts{info},
					'list|l'	=> \$opts{list},
					'upload|u=s'	=> \$opts{upload},
					'setname|s=s'	=> \$opts{setname}
	 	  );
}
else { disp_help(); }

# If they just want help, give it to them and don't bother with USB
if ($opts{help})  { disp_help(); }
# Initiate the USB stuff
# Will want to varify that an option is specified first
my $devref = init_usb();

if ($opts{battery}) { battery($devref); }
if ($opts{download}) { list_files($devref, 1); } # FIXME: Eventually allow single file
if ($opts{info}) { info($devref); }
if ($opts{list}) { list_files($devref, 0); }
if ($opts{upload}) { upload_file($opts{upload}, $devref); }
if ($opts{setname}) { setname($opts{setname}, $devref); }

cleanup($devref);
exit(0);

sub disp_help
{
	print "Usage:\n";
	print "  libnxt [option]\n\n";
	print "  -h,--help	Displays help menu.\n";
	print "  -b,--battery	Displays battery level.\n";
	print "  -d,--download	Downloads all files residing on the brick.\n";
	print "  -i,--info	Displays device information.\n";
	print "  -l,--list	Lists all files on the brick.\n";
	print "  -u,--upload <filename> Uploads file to the brick.\n";
	print "  -s,--setname <name> Sets the name of the brick.\n";
	
	exit(0);
}

sub battery
{
	my ($dev) = @_;
	my ($outbuf, $inbuf);
	
	# Get battery level
	$outbuf = pack("CC", $DIRECT_COMMAND, $GET_BATTERY_LEVEL_COMMAND);
	$inbuf = doUSB($outbuf, $dev);
	my ($reply, $command, $status, $mV) = unpack("C3v", $inbuf);
	
	printf("Battery level: %dmV\n", $mV);
}	

sub info
{
	my ($dev) = @_;
	my ($outbuf, $inbuf);
	
	# Get firmware version
	$outbuf = pack("CC", $SYSTEM_COMMAND, $GET_FIRMWARE_VERSION_COMMAND);
	$inbuf = doUSB($outbuf, $dev);
	
	my ($reply, $command, $status, $protocol_minor_version, $protocol_major_version,
	    $firmware_minor_version, $firmware_major_version) = unpack("C7", $inbuf);
	
	# Get device info 
	$outbuf = $outbuf = pack("CC", $SYSTEM_COMMAND, $GET_DEVICE_INFO_COMMAND);
	$inbuf = doUSB($outbuf, $dev);
	
	my ($nxt_name, $bt0, $bt1, $bt2, $bt3, $bt4, $bt5, $bt_signal, $free_user_flash);
	($reply, $command, $status, $nxt_name, $bt0, $bt1, $bt2, $bt3, $bt4, $bt5,
 	 $bt_signal, $free_user_flash) = unpack("C3Z15C6VV", $inbuf);
	
	printf("Firmware information:\n");
	printf("\tProtocol version: %d.%d\n", $protocol_major_version, $protocol_minor_version);
	printf("\tFirmware version: %d.02%d\n", $firmware_major_version, $firmware_minor_version);
	printf("Device information:\n");
	printf("\tNXT name: %s\n", $nxt_name);
	printf("\tBluetooth address: %02x:%02x:%02x:%02x:%02x:%02x\n", $bt0, $bt1, $bt2, $bt3, $bt4, $bt5);
	printf("\tBluetooth signal strength: %d\n", $bt_signal);
	printf("\tFree user flash: %d\n", $free_user_flash);
}

sub list_files
{
	my ($dev, $download) = @_;
	my ($outbuf, $inbuf);
	
	# Find first file
	$outbuf = pack("CCa19", $SYSTEM_COMMAND, $FIND_FIRST_FILE_COMMAND, "*.*");
	$inbuf = doUSB($outbuf, $dev);
	
	my ($reply, $command, $status, $handle, $filename, $filesize,
	    $fs0, $fs1, $fs2, $fs3) = unpack("C4Z20VX4C4", $inbuf);
	
	printf("%s %d\n", $filename, $filesize);
	
	# Download if needed
	if ($download) { readFile($filename, $filesize, $dev); }
	
	# Find the rest of the files
	while ($status == $STATUS_SUCCESS)
	{
		$outbuf = pack("C3", $SYSTEM_COMMAND, $FIND_NEXT_FILE_COMMAND, $handle);
		$inbuf = doUSB($outbuf, $dev);
		($reply, $command, $status, $handle, $filename, $filesize,
	         $fs0, $fs1, $fs2, $fs3) = unpack("C4Z20VX4C4", $inbuf);
		
		if ($status != $STATUS_SUCCESS) { last; } # Not sure I like this fix
		
		printf("%s %d\n", $filename, $filesize);
		
		# Download if needed
		if ($download) { readFile($filename, $filesize, $dev); }
	}
	
	# Close the handle
	doUSB(pack("C3", $SYSTEM_COMMAND, $CLOSE_COMMAND, $handle), $dev);
}

sub upload_file
{
	my ($file_name, $dev) = @_;
	my ($outbuf, $inbuf, $filesize);
	
	# Grab the file size
	if (-e $file_name) { $filesize = -s $file_name; }
	else { die "$file_name does not exist.\n" }

	# Open write data
	$outbuf = pack("CCZ20V", $SYSTEM_COMMAND, $OPEN_WRITE_COMMAND, $file_name, $filesize);
	$inbuf = doUSB($outbuf, $dev);
	my ($reply, $command, $status, $file_handle) = unpack("C4", $inbuf);
	
	# Make sure it worked
	if ($status)
	{
		printf("Open write data failed! Error %#04x\n", $status);
		exit(0);
	}
	else { writeFile($file_name, $filesize, $file_handle, $dev) }
}

sub setname
{
	my ($name, $dev) = @_;
	my ($outbuf, $inbuf);
	
	# Make sure the name isn't too long
	if (length($name) > 15)
	{
		print "New brick name can be no longer than 15 characters.\n";
		exit(0);
	}
	else
	{
		$outbuf = pack("CCa16", $SYSTEM_COMMAND, $SET_NAME_COMMAND, $name);
		doUSB($outbuf, $dev);
	}
}
	

sub init_usb
{

	my $usb = Device::USB->new();
	my $dev = $usb->find_device($USB_ID_VENDOR_LEGO, $USB_ID_PRODUCT_NXT);
	die "Device not found.\n" unless defined $dev;
	
	$dev->open();
	$dev->reset();
	
	print "Device found: ", $dev->filename(), ": ";
	printf("ID %04x:%04x\n", $dev->idVendor(), $dev->idProduct());
	print "SerialNumber: ", $dev->serial_number(), "\n";
	
	$dev->claim_interface($USB_INTERFACE);
	
	return $dev;
	
}

sub readFile
{
	my ($filename, $filesize, $dev) = @_;
	my ($outbuf, $inbuf, $reply, $command, $status, $read_handle);
	
	# Open file for read
	$outbuf = pack("CCZ20", $SYSTEM_COMMAND, $OPEN_READ_COMMAND, $filename);
	$inbuf = doUSB($outbuf, $dev);
	
	($reply, $command, $status, $read_handle, $filesize) = unpack("C4V", $inbuf);
	
	# Read the file
	my ($payload_len , $total_bytes_read, $to_read, $bytes_read) = (57, 0);
	my $data = '';
	my $file_data = '';
	
	while ($total_bytes_read < $filesize)
	{
		$to_read = $filesize - $total_bytes_read > $payload_len ?
			$payload_len : $filesize - $total_bytes_read;
		
		$outbuf = pack("C3v", $SYSTEM_COMMAND, $READ_COMMAND, $read_handle, $to_read);
		$inbuf = doUSB($outbuf, $dev);
		($reply, $command, $status, $read_handle, $bytes_read, $data) = unpack("C4va$to_read", $inbuf);
		$total_bytes_read += $bytes_read;
		$file_data .= $data;
	}
	
	# Close the file handle
	$outbuf = pack("C3", $SYSTEM_COMMAND, $CLOSE_COMMAND, $read_handle);
	doUSB($outbuf, $dev);
	
	# Write the file
	writeOut($filename, $file_data);
}

sub writeOut
{
	my ($filename, $data) = @_;
	
	open(FILE, ">$filename") or die "Can't open file \"$filename\": $!\n";
	if (defined($data)) { print FILE $data; }
	close(FILE) or die "Can't close file \"$filename\": $!\n";
}

sub writeFile
{
	my ($filename, $filesize, $file_handle, $dev) = @_;
	my ($outbuf, $inbuf, $err);
	
	# Open the local file for reading
	open(FILE, $filename);
	
	# Write the data (61 bytes at a time)
	while (read(FILE, my $buf, 61))
	{
		my $blen = length($buf);
		$outbuf = pack("C3a$blen", $SYSTEM_COMMAND, $WRITE_COMMAND, $file_handle, $buf);
		$inbuf = doUSB($outbuf, $dev);
		my ($reply, $command, $status, $file_handle, $lsb, $msb) = unpack("C4vX2C2", $inbuf);
		if ($status) { $err = $status; }
	}
	
	# Close the handles
	close(FILE);
	
	$outbuf = pack("C3", $SYSTEM_COMMAND, $CLOSE_COMMAND, $file_handle);
	doUSB($outbuf, $dev);
	
	if ($err) { print "Encounted error %#04x while writing %s", $err, $filename }
}

sub cleanup
{
	my ($dev) = @_;
	$dev->release_interface($USB_INTERFACE);
}

sub doUSB 
{
	my ($outbuf, $dev) = @_;

        $dev->bulk_write($USB_OUT_ENDPOINT, $outbuf, length($outbuf), $USB_TIMEOUT);
	my $inbuf = "\0" x 64;
	$dev->bulk_read($USB_IN_ENDPOINT, $inbuf, length($inbuf), $USB_TIMEOUT);

	return($inbuf);
}