#!/usr/bin/python
# usb_host_ctl Copyright (c) 2012 Stuart Pook (http://www.pook.it/)
# 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 3 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.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# copied from the script Pali Rohar wrote for usbmode_1.0-2_armel.deb
# http://talk.maemo.org/showthread.php?t=85658
# http://atrey.karlin.mff.cuni.cz/~pali/projects/maemo/bq2415x.html

# test connected to a powered hub 2.0 with
# Bus 001 Device 006: ID 0461:4d51 Primax Electronics, Ltd 
# Bus 001 Device 005: ID 0424:2228 Standard Microsystems Corp. 
# Bus 001 Device 004: ID 0424:2602 Standard Microsystems Corp. 
# Bus 001 Device 003: ID 0424:2512 Standard Microsystems Corp. 
# Bus 001 Device 001: ID 1d6b:0002  
# -s high -m 1 -e 2 -n 0 -B 0 -c -u 2 ->no device arrived none low
# -s high -m 1 -e 2 -n 1 -B 0 -c -u 2 -> connection rapid but no USB devices
# -s high -m 1 -e 2 -n 0 -B 1 -c -u 2 -> no device arrived none low
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 2 -> worked on first try
# -s high -m 1 -e 1 -n 1 -B 1 -c -u 2 -> worked on 3rd try (lsusb missed the mouse)
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 2 -> worked on 3rd try (lsusb ok)
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 2 -> worked on 3rd try (lsusb ok)
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 2 -> worked on 3rd try (lsusb ok)
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 1 -> worked on 1st try but lsusb missing a device
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 1 -> worked on 3rd try (lsusb ok)
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 1 -> worked on 3rd try (lsusb ok)
# -s high -m 1 -e 1 -n 1 -B 1 -c -u 1 -> worked on 5th try (lsusb missing 2 devices)
# -s high -m 1 -e 1 -n 1 -B 1 -c -u 1 -> worked on 5th try (lsusb missing 2 devices)
# -s high -m 1 -e 2 -n 1 -B 1 -fc -u 1 -> worked on 3th try (lsusb ok)
# -s high -m 1 -e 2 -n 1 -B 1 -fc -u 1 -> worked on 3th try (lsusb ok)
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 1 -> worked on 3rd try (lsusb ok)
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 1-> worked on 3rd try (lsusb ok)
# -s high -m 1 -e 2 -n 1 -B 0 -c -u 1 -> no device arrived none low
# -s high -m 1 -e 2 -n 0 -B 1 -c -u 1 -> worked 1st try (lsusb ok)
# reboot
# -s high -m 1 -e 2 -n 0 -B 1 -c -u 1 -> worked 1st try (lsusb ok)
# -s high -m 1 -e 2 -n 0 -B 1 -c -u 1 -> failed no device arrived none low
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 1 -> worked 1st try (lsusb ok)
# -s high -m 1 -e 1 -n 1 -B 1 -c -u 1 -> worked on 5th try (lsusb missing 2 devices)
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 1 -> worked on 3rd try (lsusb ok)
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 4  -> worked on 3rd try (lsusb ok)
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 4  -> worked on 3rd try (lsusb ok)
# -s high -m 1 -e 2 -n 1 -B 1 -c -u 1 -> worked on 3rd try (lsusb ok)

from __future__ import with_statement 
import os, sys, optparse, subprocess, pipes, errno, time, select, signal

def sleep(d, options):
	d = d * options.multiply_delays
	if options.verbose: print "sleep", d
	time.sleep(d)
	
def lsmod(options):
	modules = []
	with open("/proc/modules") as f:
			for l in f:
				modules.append(l.split()[0])
	return modules

def open_bq24150(file, mode, options):
	try:
		return open("/sys/class/power_supply/bq24150-0/" + file, mode)
	except IOError, (e, strerror):
		if e != errno.ENOENT:
			raise
		return open("/sys/class/power_supply/bq24150a-0/" + file, mode)

def get_bq24150(file, options):
	with open_bq24150(file, "r", options) as f:
		r = f.readline().strip()
#		if options.verbose: print "bq24150", file, r
		return r

def set_bq24150(file, value, options):
	with open_bq24150(file, "w", options) as f:
		f.write(value + "\n")

def set_charger_mode(mode, options):
	set_bq24150("mode", mode, options)
	if options.verbose: print "set bq24150", "mode", mode

def get_charger_mode(options):
	r = get_bq24150("mode", options)
#	if options.verbose: print "get_charger_mode", r
	return r

def call_print(command, options):
	if options.verbose:
		for c in command:
			print pipes.quote(c),
		print
			
def call(command, options):
	call_print(command, options)
	subprocess.check_call(command)
	if options.verbose: print command[0], "done"
	
def call_ignore(command, options):
	call_print(command, options)
	r = subprocess.call(command)
	if options.verbose: print command[0], "done", r
	return r
	
usb_mode_file = "/sys/devices/platform/musb_hdrc/mode"

def set_usb_mode(mode, options):
	with open(usb_mode_file, "w") as f:
		f.write(mode + "\n")
	if options.verbose: print "set_usb_mode", mode

def get_usb_mode(options):
	with open(usb_mode_file) as f:
		return f.readline().strip()

def rmmod(module, options):
	call(["modprobe", "-r"] + [module], options)
		
def insmod(module, options):
	call(["modprobe"] + [module], options)

def gadget_unload(options):
	if call_ignore(["osso-usb-mass-storage-is-used.sh"], options) == 0:
		call(["osso-usb-mass-storage-disable.sh"], options)
	modules = []
	for module in lsmod(options):
		if module == "g_nokia":
			if options.verbose: print "found", module
			call(["initctl", "emit", "G_NOKIA_REMOVE"], options)
			call_ignore(["killall", "pnatd", "obexd", "syncd"], options)
			call(["pcsuite-disable.sh"], options)
	for module in lsmod(options):
		if module[0:2] == "g_":
			rmmod(module, options)

def is_running(cmd, options):
	proc = "/proc"
	for d in os.listdir(proc):
		if d.isdigit():
			try:
				exe = os.readlink(os.path.join(proc, d, "exe"))
			except OSError:
				pass
			else:
				if exe.endswith("/" + cmd):
					return True
	return False
		
def bme(options):
	for m in lsmod(options):
		if m == "bq2415x_charger" or  m == "bq27x00_battery":
			rmmod(m, options)
	if not is_running("bme_RX-51", options):
		sleep(2, options) # avoid start: bme respawning too fast, stopped
		cmd = ["start"]
		if not options.verbose:
			cmd.append("-q")
		call(cmd + ["bme"], options)
	elif options.verbose:
		print "bme already running"
	
def peripheral(options):
	gadget_unload(options)
	try:
		if get_charger_mode(options).find("boost") != -1:
			set_charger_mode("none", options)
	except IOError:
		pass
	set_usb_mode("peripheral", options)
	try:
		set_charger_mode("auto", options)
	except IOError:
		pass
	call(["modprobe", "g_file_storage", "stall=0", "luns=2", "removable"], options)
	bme(options)
	
def enumerate(options):
	with open("/proc/driver/musb_hdrc", "w") as f:
		f.write("F\n")
	if options.verbose: print "enumerate"

def usb_speed(options):
	with open("/sys/devices/platform/musb_hdrc/hostdevice") as f:
		r = f.readline().strip();
#	if options.verbose: print "usb_speed (usb_device)", r
	return r
	
def kernel(options):
	if is_running("bme_RX-51", options):
		call_ignore(["stop", "-q", "bme"], options)

	list = set(["bq2415x_charger", "bq27x00_battery"])
	for l in lsmod(options):
		list.discard(l)
	for m in list:
		insmod(m, options)
	
def check_usb(options):
	usb = "/dev/bus/usb/001/001"
	if options.verbose: print "reading", usb
	sys.stdout.flush()
	fds = os.pipe()
	pid = os.fork()
	if pid == 0:
		os.close(fds[0])
		with open(usb) as f:
			f.read(100)
		sys.exit(0)
	os.close(fds[1])
	poller = select.poll()
	poller.register(fds[0])
	before = time.time()
	delay = options.usb_delay
	if poller.poll(delay * 1000) == []:
		if options.verbose: print "check_usb timeout after", delay, "seconds"
		os.kill(pid, signal.SIGKILL)
		r = False
	else:
		if options.verbose: print "check_usb ok", time.time() - before, "seconds"
		r = True
	os.close(fds[0])
	while os.waitpid(0, os.WNOHANG) != (0, 0):
		pass
	return r
	
def error(message, options):
	print  >> sys.stderr, message
	if options.interactive:
		time.sleep(60 * 60)
	sys.exit(1)
	
def hostc(options):
	kernel(options)
	
	# this is wrong, I think that I should check /sys/devices/platform/musb_hdrc/hostdevice as well
	if usb_speed(options) != "none" and check_usb(options):
		if options.verbose: print "already attached"
		error("already attached", options)

	mode = get_charger_mode(options)
	# this appears to be wrong, I have nothing plugged in and mode is "auto (host)"
	if mode.find("none") == -1  and mode.find("boost") == -1:
		error("mode " + mode + "charger attached?", options)
		
	gadget_unload(options)
	set_charger_mode("none", options)
	call(["modprobe", "g_file_storage", "stall=0", "luns=2", "removable"], options)
	rtry = options.start
	good = ""
	
	for i in range(10):
		if options.verbose: print "trying", rtry
		set_charger_mode("none", options)
		sleep(options.charge_none_delay, options)
		set_usb_mode("host" + rtry[0], options)
	
		set_charger_mode("boost", options)
		sleep(options.charge_boost_delay, options)
		enumerate(options)
		sleep(options.enumerate_delay, options)
	
		if check_usb(options):
			if options.verbose:
				speed = usb_speed(options)
				print "hostc ok at", speed, "after", i + 1
			return

		speed = usb_speed(options)
		if options.verbose: print "speed", speed
		if not options.force_speed:
			if speed == "none":
				next = good
			elif speed == "low":
				next = "low"
			elif speed == "full/low":
				if rtry == "low":
					next = "full"
				else:
					next = "low"
			elif speed == "full":
				next = "full"
			elif speed == "full/high":
				if rtry == "full":
					next = "high"
				else:
					next = "full"
			elif speed == "high":
				next = "high"
			else:
				error("unknown speed in /sys" + speed, options)
				
			if speed.find(rtry) != -1:
				good = rtry
				
			if rtry == next and speed != "none":
				error("correct speed but no device %s %d" % ( speed, rtry), options)
				
			if next == "":
				error("no device arrived %s %d" % ( speed, rtry), options)
			
			rtry = next
		if options.verbose: print
		
	error("tried too many times %d" % (rtry), options)
	
def usb_host_ctl(options):
	if options.verbose: print options.mode
	if options.mode == "peripheral":
		peripheral(options)
	elif options.mode == "charge":
		hostc(options)
		set_charger_mode("auto", options)
	elif options.mode == "boost":
		hostc(options)
	elif options.mode == "enumerate":
		enumerate(options)
		if check_usb(options):
			if options.verbose:
				speed = usb_speed(options)
				print "enumerate ok at", speed
	if options.lsusb:
		call(["lsusb"], options)
		
def ls(dir, options):
	print os.listdir(dir)
		
def mount(options):
	devices = []
	with open("/proc/partitions") as mounts:
		for mount in mounts:
			fields = mount.split()
			if len(fields) == 4:
				device = fields[3]
				if len(device) == 4 and device.startswith("sd"):
					if call_ignore(["pmount", device], options) == 0:
						ls(os.path.join("/media", device), options)

def mounted(options):
	with open("/proc/mounts") as mounts:
		for mount in mounts:
			fields = mount.split()
			device = fields[0]
			if device.startswith("/dev/sd"):
				print device, fields[1], fields[2]

def umount(options):
	devices = []
	with open("/proc/mounts") as mounts:
		for mount in mounts:
			device = mount.split()[0]
			if device.startswith("/dev/sd"):
				devices.append(device)
	for device in devices:
		call_ignore(["umount", device], options)
		
def interactive_commands(options):
	help = "commands: [l]susb [e]numerate [m]ount [u]mount mounte[d] [c]harge [5]00mA host charge, [8]00mA charge, 1800m[A] charge, [i]nfo [q]uit"
	prompt = help
	while True:
		print get_bq24150("status", options), "@", \
			get_bq24150("current_limit", options), "mA", \
			prompt,
		prompt = "?"
		f = sys.stdin.readline()
		if not f:
			break
		f = f.strip()
		if f == "l":
			call(["lsusb"], options)
		elif f == "e":
			enumerate(options)
		elif f == "m":
			mount(options)
		elif f == "d":
			mounted(options)
		elif f == "u":
			umount(options)
		elif f == "c":
			set_charger_mode("auto", options)
		elif f == "5":
			set_bq24150("mode", "host", options)
		elif f == "8":
			set_bq24150("mode", "host", options)
			set_bq24150("current_limit", "800", options)
		elif f == "A":
			set_bq24150("mode", "host", options)
			set_bq24150("current_limit", "1800", options)
		elif f == "i":
			print "speed " + usb_speed(options) +\
				", mode " + get_usb_mode(options) +\
				", bq24150: mode " + get_bq24150("mode",options) +\
				", current_limit " + get_bq24150("current_limit", options) + \
				", status" + get_bq24150("status", options)
		elif f == 'q':
			break
		elif f == "":
			pass
		else:
			print "unknown command", f
			prompt = help
			
def ask_for_quit(options):
	print "hit return to quit"
	sys.stdin.readline()

def main():
	parser = optparse.OptionParser()
	boost = "boost"
	parser.set_defaults(mode=boost)
	
	parser.add_option("-i", "--interactive", action="store_true", help="interactive mode", default=False)
	parser.add_option("-f", "--force_speed", action="store_true", help="don't try any other speeds", default=False)
	parser.add_option("-v", "--verbose", action="store_true", dest="verbose", default=False)
	parser.add_option("-s", "--start", type = "choice", choices=["low","full","high"], help="start speed in host mode [%default]", default = "high")
	parser.add_option("-u", "--usb_delay", type = "float", help="timeout in seconds waiting for USB device to respond [%default]", default = 1)
	parser.add_option("-m", "--multiply_delays", type = "float", help="multiply all delays [%default]", default = 1)
	parser.add_option("-e", "--enumerate_delay", type="float", help = "delay after enumerate [%default]", default = 2)
	parser.add_option("-n", "--charge_none_delay", type="float", help = "delay after charge none [%default]", default = 1)
	parser.add_option("-B", "--charge_boost_delay", type="float", help = "delay after charge boost [%default]", default = 1)
	parser.add_option("-l", "--lsusb", action="store_true", help = "run lsusb")

	parser.add_option("-E", "--enumerate", action="store_const", dest="mode", const="enumerate", help="enumerate only")
	parser.add_option("-p", "--peripheral", action="store_const", dest="mode", const="peripheral", help = "the N900 is a USB client charging if it can")
	parser.add_option("-b", "--boost", "--hostboost", action="store_const", dest="mode", const=boost, help = "host mode with the N900 providing power")
	parser.add_option("-c", "--charge", "--hostcharge", action="store_const", dest="mode", const="charge", help = "host mode with the N900 charging itself if it can")
	(options, args) = parser.parse_args()
	
	usb_host_ctl(options)
	if options.interactive and not options.mode == "peripheral":
		interactive_commands(options)
		umount(options)
		peripheral(options)
		ask_for_quit(options)

if __name__ == "__main__":
	main()
