#!/usr/bin/python
# Software License Agreement (BSD License)
#
# Copyright (c) 2012, I Heart Engineering
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above
#    copyright notice, this list of conditions and the following
#    disclaimer in the documentation and/or other materials provided
#    with the distribution.
#  * Neither the name of the I Heart Engineering nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
# Requires pyusb-1.0
# sudo easy_install pyusb
#################

import subprocess
import usb.core
import usb.util
import getpass

#################
# Globals

MAX_ISERIAL_LEN = 8;
create_iserial = "";
arm_iserial = "";
armed = False;

#################
# Functions

def is_root():
    return (getpass.getuser() == "root")

def get_ftdi_count():
    dev = usb.core.find(find_all=True,idVendor=0x0403, idProduct=0x6001)
    if dev is None:
        return 0
    return len(dev)

def generate_udev_rule():
    print """\
# Turtlebot FTDI rules Autogenerated by turtlebot_config
SUBSYSTEMS=="usb", ENV{ID_MODEL}=="", IMPORT{program}="usb_id --export %p"
SUBSYSTEMS=="usb", ENV{ID_MODEL_FROM_DATABASE}=="", IMPORT{program}="usb-db %p"
SUBSYSTEMS=="usb", ATTRS{idVendor}!="", ATTRS{idProduct}!="", ENV{ID_VENDOR_ID}="$attr{idVendor}", ENV{ID_MODEL_ID}="$attr
{idProduct}"

ATTRS{idVendor}=="0403", ATTRS{idProduct}=="6001", MODE="664", OWNER="turtlebot", GROUP="turtlebot"
"""
    print 'ATTRS{idVendor}=="0403", ATTRS{idProduct}=="6001", ENV{ID_SERIAL_SHORT}=="%s", SYMLINK+="irobot_create"' % create_iserial
    if (armed):
        print 'ATTRS{idVendor}=="0403", ATTRS{idProduct}=="6001", ENV{ID_SERIAL_SHORT}=="%s", SYMLINK+="turtlebot_arm"' % arm_iserial

def dialog_run_as_root():
    p = subprocess.Popen(['zenity','--warning','--text=This configuration tool must be run as root','--title=TurtleBot Config Tool'], stdout=subprocess.PIPE)
    output = p.communicate()[0]
    p.wait()

def dialog_unplug_all():
    rc = subprocess.call(['zenity','--question','--ok-label=OK','--cancel-label=Cancel','--text=Please unplug all USB cables.','--title=TurtleBot Config Tool'], stdout=subprocess.PIPE)
    if rc:
        exit()

def dialog_plug_create():
    rc = subprocess.call(['zenity','--question','--ok-label=OK','--cancel-label=Cancel','--text=Connect iRobot Create USB serial cable.','--title=TurtleBot Config Tool'], stdout=subprocess.PIPE)
    if rc:
        exit()

def dialog_armed():
    rc = subprocess.call(['zenity','--question','--text=Is this robot equipped with a TurtleBot Arm?','--title=TurtleBot Config Tool'])
    return not rc

def dialog_plug_arm():
    rc = subprocess.call(['zenity','--question','--ok-label=OK','--cancel-label=Cancel','--text=Connect TurtleBot Arm USB cable.','--title=TurtleBot Config Tool'], stdout=subprocess.PIPE)
    if rc:
        exit()


#################

if (not is_root()):
    dialog_run_as_root()
    exit()

while (get_ftdi_count() > 0):
    dialog_unplug_all()

while (get_ftdi_count() != 1):
    dialog_plug_create()

dev = usb.core.find(find_all=True,idVendor=0x0403, idProduct=0x6001)
for i in range(len(dev)):
    create_iserial = usb.util.get_string(dev[i],MAX_ISERIAL_LEN,dev[i].iSerialNumber)

armed = dialog_armed()

if armed:
    while (get_ftdi_count() != 2):
        dialog_plug_arm()

    if (armed == True):
        dev = usb.core.find(find_all=True,idVendor=0x0403, idProduct=0x6001)
        for i in range(len(dev)):
            temp_iserial = usb.util.get_string(dev[i],MAX_ISERIAL_LEN,dev[i].iSerialNumber)
            if (temp_iserial != create_iserial):
                arm_iserial = temp_iserial

generate_udev_rule()
