#!/usr/bin/env python
# This program is a shell-like program that can control the nxt brick
#
# Copyright (C) 2006 Brian Dam Pedersen
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.


import os
import os.path
import glob
import stat
import sys
from usb import USBError
import nxt.locator
import nxt.error
from nxt.brick import FileReader,FileWriter

########### Functions and variables that are necessary for the shell itself

def cleanstr(x):
    return x.replace('\x00','')

global brick
brick=None

def connect(dummy=None):
    global brick
    try:
        br=nxt.locator.find_one_brick()
        brick=br.connect()
    except nxt.locator.BrickNotFoundError:
        print "Could not connect to brick. Make sure it is on and wired up"
        return
    except USBError:
        "Whoa - that was quick. Found brick and lost it again. Check cable"
        return
    name=cleanstr(brick.get_device_info()[0])
    print 'Connected to "%s"' % name

def run_shell(commands):
    global brick
    connect()
    while 1:
        cmd=raw_input('# ').strip().lower().split(' ',1)
        if len(cmd)==1:
            arg=None
        else:
            arg=cmd[1]
        if not cmd[0]:
            continue
        else:
            if cmd[0] == 'help':
                if arg:
                    if not arg in commands:
                        print "Unknown command: %s" % arg
                    else:
                        print "\n%s" % commands[arg][2]
                else:
                    cmds=commands.keys()
                    colwidth=max([len(x) for x in cmds])+3
                    print "The following commands are available:\n"
                    colcnt=min(80/colwidth,4)
                    cc=0
                    for command in cmds:
                        cc+=1
                        sys.stdout.write("%s%s"%(command," "*(colwidth-len(command))))
                        if cc==colcnt:
                            cc=0
                            print ''
                    if cc > 0:
                        print ''
                    print "\nWrite 'help <cmd>' on the prompt for more information on a specific command"
            elif not cmd[0] in commands:
                print "Unknown command: %s" % cmd[0]
            else:
                issue=commands[cmd[0]]
                if issue[1] and brick==None:
                    print "Command requires connection. Please reconnect by issuing the 'connect' command"
                else:
                    try:
                        issue[0](arg)
                    except USBError:
                        brick=None
                        print "Lost connection..."
        print ''
            

########### Local shell command functionals
def sh_lls(param,silent=0):
    if not param:
        param = "*"
    if param[-1]=="/":
        param+="*"
    plist=glob.glob(param)
    if silent: return plist
    tsz=0
    filler=max([len(x) for x in plist])+2
    if plist:
        print "..%s<DIR>"%(" "*(filler-2))
    screensiz=1
    for f in plist:
        st=os.stat(f)
        if not stat.S_ISDIR(st[stat.ST_MODE]):
            sz=os.stat(f)[stat.ST_SIZE]
            tsz+=sz
            lf=filler-len(f)
            print "%s%s%10d bytes"%(f," "*lf,sz)
        else:
            lf=filler-len(f)
            print "%s%s           <DIR>"%(f," "*lf)
        screensiz+=1
        if screensiz==23:
            raw_input('Press enter to continue ...')
            screensiz=0
    print "Total %d files, %d bytes" % (len(plist),tsz)

def sh_lcd(param):
    if not param:
        print os.getcwd()
    else:
        try:
            os.chdir(param)
        except Exception, x:
            print x

def sh_quit(param):
    sys.exit(0)

########### NXT functions

def nxt_ls(param,silent=0):
    global brick
    if not param:
        param='*.*'
    try:
        fdesc=brick.find_first(param)
    except nxt.error.FileNotFound:
        print "No such file found on the brick"
        return []
    flist=[]
    flist.append((cleanstr(fdesc[1]),fdesc[2]))
    total=fdesc[2]
    while 1:
        try:
            fdesc=brick.find_next(fdesc[0])
        except nxt.error.FileNotFound:
            break
        flist.append((cleanstr(fdesc[1]),fdesc[2]))
        total+=fdesc[2]
    left=brick.get_device_info()[3]

    if silent:
        return flist

    filler=max([len(x) for x in [y[0] for y in flist]])+2
    screensiz=0
    for f in flist:
        lf=filler-len(f[0])
        print "%s%s%6d bytes"%(f[0]," "*lf,f[1])
        if screensiz==23:
            raw_input('Press enter to continue ...')
            screensiz=0
    print "Total %d files, %d bytes. %d bytes left on brick" % (len(flist),total,left)

def nxt_rm(param):
    global brick
    if not param:
        print "You must specify a file"
        return
    try:
        fdesc=brick.find_first(param)
    except nxt.error.FileNotFound:
        print "No such file found on the brick"
        return
    ans=raw_input("Do you really want to delete %s ? [y/n] " % param).lower()
    if ans=='y':
        brick.delete(param)
        print "%s deleted" % param

def nxt_put(param):
    global brick
    if not param:
        print "You must specify a file or wildcard"
        return
    flist=sh_lls(param,1)
    allperm=False
    oom =False
    while flist:
        f=flist[0]
        flist=flist[1:]
        brickfilename=os.path.split(f)[1]
        if oom:
            print "Skipped writing %s" % f
            continue
        try:
            brick.find_first(param)
            # If we got here we have to ask for permission
            if not allperm:
                if flist:
                    ans=raw_input('Overwrite file %s on brick (Yes/No/All) ? '%f).lower()
                    if ans=='a':
                        ok_to_write=True
                        allperm=True
                    elif ans=='y':
                        ok_to_write=True
                    else:
                        ok_to_write=False
                else:
                    ans=raw_input('Overwrite file %s on brick (Yes/No) ? '%f).lower()
                    if ans=='y':
                        ok_to_write=True
                    else:
                        ok_to_write=False
            else:
                ok_to_write=True
            # Delete file if we got an OK
            if ok_to_write:
                brick.delete(brickfilename)
        except nxt.error.FileNotFound:
            ok_to_write=True
        if not ok_to_write:
            print "Skipped writing %s" % f
            continue

        # First we check for space
        sz=os.stat(f)[stat.ST_SIZE]
        left=brick.get_device_info()[3]
        if left < sz:
            print "Not enough space left on brick, skipping the rest of the files !!!"
            print "Skipped writing %s" % f
            oom=True
            continue
        # Write the file to the brick
        localfile=open(f,'rb')
	writer = FileWriter(brick, brickfilename, localfile)
	writer.__enter__()
	try:
            sys.stdout.write("Writing %s .. " % f)
            size = 0
            for n_bytes in writer:
                size += n_bytes
            print '%d bytes' % size
	finally:
            writer.__exit__(None, None, None)
                
        
        
        

def nxt_get(param):
    if not param:
        print "You must specify a file or wildcard"
        return
    flist=[x[0] for x in nxt_ls(param,1)]
    for f in flist:
        localfile=open(f,'wb')
	reader = FileReader(brick, f)
	reader.__enter__()
	try:
            sys.stdout.write("Reading %s .. " % f)
            size=0
            for data in reader:
                size+=len(data)
                localfile.write(data)
            print '%d bytes' % size
	finally:
            reader.__exit__(None, None, None)

        
    
        
            
########### Shell command registration table
# Meaning of tuple:
# 0: Handler function
# 1: Requires connection (True/False)
shell_commands = {\
    'lls' : (sh_lls,False,"Lists the files in the current directory on the PC.\nNormal UNIX wildcards are allowed"),\
    'lcd' : (sh_lcd,False,"Changes directory on the PC. Write 'lcd' with no arguments to get current\ndirectory" ),\
    'quit': (sh_quit,False,"Quits this program"),\
    'connect' : (connect,False,"Connects to the brick in case the connection is lost"),\
    'ls' : (nxt_ls,True,"Lists the files on the brick. * and ? are allowed as wildcards"),\
    'rm' : (nxt_rm,True,"Remove one or more files from the brick. * and ? are allowed as wildcards"),\
    'get': (nxt_get,True,"Fetches one or more files from the brick to the current directory on the PC.\n* and ? are allowed as wildcards"),\
    'put': (nxt_put,True,"Uploads one or more files from the current directory on the PC to the brick.\n* and ? are allowed as wildcards")}


########### Do me
if __name__=="__main__":
    print "nxtsh version 0.1 (c) Brian Dam Pedersen, 2007. \nThis software is FREE software licenced under the GNU GPL"
    print "\nType 'help' at the prompt for more information\n"
    run_shell(shell_commands)