#!/usr/bin/env python3
#
#   Razer device commandline configuration tool
#
#   Copyright (C) 2007-2011 Michael Buesch <m@bues.ch>
#
#   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 sys
import os
import getopt
import time
import re
from pyrazer import *
from configparser import *

try:
	PYRAZER_SETUP_PY == True
except:
	print("ERROR: Found an old 'pyrazer' module.")
	print("You should uninstall razercfg from the system (see README)")
	print("and re-install it properly.")
	sys.exit(1)


razer = None

def getRazer():
	global razer
	if razer is None:
		razer = Razer()
	return razer

class Operation:
	def parseProfileValueStr(self, parameter, idstr, nrValues=1):
		# Parse profile:value[:value]... string. Default to active
		# profile, if not given. May raise ValueError.
		split = parameter.split(":")
		if len(split) == nrValues:
			profile = None
			values = split
		elif len(split) == nrValues + 1:
			profile = int(split[0].strip())
			values = split[1:]
			if profile == 0:
				profile = getRazer().getActiveProfile(idstr) + 1
			if profile < 0:
				raise ValueError
		else:
			raise ValueError
		return (profile, values)

class OpSleep(Operation):
	def __init__(self, seconds):
		self.seconds = seconds

	def run(self, idstr):
		time.sleep(self.seconds)

class OpScan(Operation):
	def run(self, idstr):
		scanDevices()

class OpReconfigure(Operation):
	def run(self, idstr):
		reconfigureDevices()

class OpGetFwVer(Operation):
	def run(self, idstr):
		verTuple = getRazer().getFwVer(idstr)
		print("%s: Firmware version %d.%02d" %\
			(idstr, verTuple[0], verTuple[1]))

class OpGetProfile(Operation):
	def run(self, idstr):
		profileId = getRazer().getActiveProfile(idstr)
		print("Active profile: %u" % (profileId + 1))

class OpGetFreq(Operation):
	def run(self, idstr):
		freqs = getRazer().getSupportedFreqs(idstr)
		minfo = getRazer().getMouseInfo(idstr)
		curProf = getRazer().getActiveProfile(idstr)
		if minfo & Razer.MOUSEINFOFLG_GLOBAL_FREQ:
			curFreq = getRazer().getCurrentFreq(idstr)
			self.printFreq(freqs, curFreq)
		if minfo & Razer.MOUSEINFOFLG_PROFILE_FREQ:
			profiles = getRazer().getProfiles(idstr)
			for profile in profiles:
				sys.stdout.write("Profile %2u%s:   " %\
					(profile + 1,
					 "*" if profile == curProf else " "))
				curFreq = getRazer().getCurrentFreq(idstr, profile)
				self.printFreq(freqs, curFreq)

	def printFreq(self, freqs, curFreq):
		output = []
		for freq in freqs:
			pfx = "*" if freq == curFreq else " "
			output.append("%s%u Hz" % (pfx, freq))
		print(", ".join(output))

class OpGetRes(Operation):
	def run(self, idstr):
		mappings = getRazer().getSupportedDpiMappings(idstr)
		profiles = getRazer().getProfiles(idstr)
		curProf = getRazer().getActiveProfile(idstr)
		for profile in profiles:
			sys.stdout.write("Profile %2u%s:   " %\
				(profile + 1,
				 "*" if profile == curProf else " "))
			curMapping = getRazer().getDpiMapping(idstr, profile)
			output = []
			pm = [m for m in mappings if m.profileMask == 0 or\
					      m.profileMask & (1 << profile)]
			for mapping in pm:
				pfx = "*" if mapping.id == curMapping else " "
				r = [ "%u" % r for r in mapping.res if r is not None ]
				rStr = "/".join(r)
				output.append("%s%u (%s DPI)" % (pfx, mapping.id + 1, rStr))
			print(", ".join(output))

class OpPrintLeds(Operation):
	def run(self, idstr):
		minfo = getRazer().getMouseInfo(idstr)
		if minfo & Razer.MOUSEINFOFLG_GLOBAL_LEDS:
			sys.stdout.write("Global LEDs:   ")
			leds = getRazer().getLeds(idstr)
			self.printLeds(leds)
		if minfo & Razer.MOUSEINFOFLG_PROFILE_LEDS:
			profiles = getRazer().getProfiles(idstr)
			curProf = getRazer().getActiveProfile(idstr)
			for profile in profiles:
				sys.stdout.write("Profile %2u%s LEDs:   " %\
					(profile + 1,
					 "*" if profile == curProf else " "))
				leds = getRazer().getLeds(idstr, profile)
				self.printLeds(leds)

	def printLeds(self, leds):
		output = []
		for led in leds:
			state = "on" if led.state else "off"
			mode  = led.mode.toString()
			color = ""
			if led.color is not None:
				color = "color#%02X%02X%02X" %\
					(led.color.r, led.color.g, led.color.b)
			output.append("%s => %s/%s/%s" % (led.name, state, mode, color))
		print(",  ".join(output))

class OpSetProfile(Operation):
	def __init__(self, param):
		self.param = param

	def run(self, idstr):
		try:
			profileId = int(self.param) - 1
			error = getRazer().setActiveProfile(idstr, profileId)
			if error:
				raise RazerEx("Failed to set active profile (%s)" %\
					      Razer.strerror(error))
		except (ValueError) as e:
			raise RazerEx("Invalid parameter to --profile option")

class OpSetLedState(Operation):
	def __init__(self, param):
		self.param = param

	def setLed(self, idstr, led, state):
		led.state = state
		error = getRazer().setLed(idstr, led)
		if error:
			raise RazerEx("Failed to set LED state (%s)" %\
					  Razer.strerror(error))

	def run(self, idstr):
		try:
			(profile, config) = self.parseProfileValueStr(self.param, idstr, 2)
			ledName = config[0].strip().lower()
			newState = razer_str2bool(config[1])
			if profile is None:
				profile = Razer.PROFILE_INVALID
			else:
				profile -= 1
			leds = getRazer().getLeds(idstr, profile)
			if 'all' == ledName:
				for led in leds:
					self.setLed(idstr, led, newState)
			else:
				led = [led for led in leds if led.name.lower() == ledName.lower()][0]
				self.setLed(idstr, led, newState)
		except (IndexError, ValueError):
			raise RazerEx("Invalid parameter to --setled option")

class OpSetLedColor(Operation):
	def __init__(self, param):
		self.param = param

	def run(self, idstr):
		try:
			(profile, config) = self.parseProfileValueStr(self.param, idstr, 2)
			ledName = config[0].strip()
			newColor = RazerRGB.fromString(config[1])
			if profile is None:
				profile = Razer.PROFILE_INVALID
			else:
				profile -= 1
			leds = getRazer().getLeds(idstr, profile)
			led = [led for led in leds if led.name.lower() == ledName.lower()][0]
			led.color = newColor
			error = getRazer().setLed(idstr, led)
			if error:
				raise RazerEx("Failed to set LED color (%s)" %\
					      Razer.strerror(error))
		except (IndexError, ValueError):
			raise RazerEx("Invalid parameter to --setledcolor option")

class OpSetLedMode(Operation):
	def __init__(self, param):
		self.param = param

	def run(self, idstr):
		try:
			(profile, config) = self.parseProfileValueStr(self.param, idstr, 2)
			ledName = config[0].strip().lower()
			newMode = RazerLEDMode.fromString(config[1].lower())
			if profile is None:
				profile = Razer.PROFILE_INVALID
			else:
				profile -= 1
			leds = getRazer().getLeds(idstr, profile)
			led = [led for led in leds if led.name.lower() == ledName.lower()][0]
			led.mode = newMode
			error = getRazer().setLed(idstr, led)
			if error:
				raise RazerEx("Failed to set LED mode (%s)" %\
					      Razer.strerror(error))
		except (KeyError, IndexError, ValueError):
			raise RazerEx("Invalid parameter to --setledmode option")


class OpSetRes(Operation):
	def __init__(self, param):
		self.param = param

	def setDpiMapping(self, idstr, profile, axisId, value):
		# Get profile mappings.
		mappings = [m for m in getRazer().getSupportedDpiMappings(idstr) \
					if m.profileMask == 0 or m.profileMask & (1 << profile)]
		if value >= 100:
			# Value is in DPI.
			mappings = [m for m in mappings if value in m.res]
		else:
			# Value is a mapping ID.
			mappings = [m for m in mappings if (value - 1) == m.id]
		try:
			mappingId = mappings[0].id
		except IndexError:
			raise RazerEx("Invalid resolution %d" % value)
		error = getRazer().setDpiMapping(idstr, profile - 1, mappingId, axisId=axisId)
		if error:
			raise RazerEx("Failed to set resolution to %u (%s)" %\
					(mappingId, Razer.strerror(error)))

	def run(self, idstr):
		try:
			(profile, values) = self.parseProfileValueStr(self.param, idstr)
			resolutions = []
			for arg in values[0].split(','):
				m = re.match('^\d+$', arg)
				if m is not None:
					resolutions.append((None, int(arg)))
					continue
				m = re.match('^(\d+)x(\d+)$', arg)
				if m is not None:
					resolutions.append((0, int(m.group(1))))
					resolutions.append((1, int(m.group(2))))
					continue
				raise ValueError
		except ValueError:
			raise RazerEx("Invalid parameter to --res option")
		if profile is None:
			# No profile number was specified. Get the current one.
			profile = getRazer().getActiveProfile(idstr) + 1
		for axisId, value in resolutions:
			self.setDpiMapping(idstr, profile, axisId, value)

class OpSetFreq(Operation):
	def __init__(self, param):
		self.param = param

	def run(self, idstr):
		try:
			(profile, freq) = self.parseProfileValueStr(self.param, idstr)
			freq = int(freq[0])
		except ValueError:
			raise RazerEx("Invalid parameter to --freq option")
		if profile is None:
			profile = Razer.PROFILE_INVALID
		else:
			profile -= 1
		error = getRazer().setFrequency(idstr, profile, freq)
		if error:
			raise RazerEx("Failed to set frequency to %d Hz (%s)" %\
					(freq, Razer.strerror(error)))

class OpFlashFw(Operation):
	def __init__(self, filename):
		self.filename = filename

	def run(self, idstr):
		p = RazerFirmwareParser(self.filename)
		data = p.getImage()
		print("Flashing firmware on %s ..." % idstr)
		print("!!! DO NOT DISCONNECT ANY DEVICE !!!")
		print("Sending %d bytes..." % len(data))
		error = getRazer().flashFirmware(idstr, data)
		if error:
			raise RazerEx("Failed to flash firmware (%s)" % Razer.strerror(error))
		print("Firmware successfully flashed.")

# List of operations
class DevOps:
	def __init__(self, idstr):
		self.idstr = idstr
		self.ops = []

	def add(self, op):
		self.ops.append(op)

	def runAll(self):
		for op in self.ops:
			op.run(self.idstr)

def scanDevices():
	getRazer().rescanMice()
	mice = getRazer().getMice()
	for mouse in mice:
		print(mouse)

def reconfigureDevices():
	getRazer().rescanDevices()
	getRazer().reconfigureDevices()

def exit(exitcode):
	sys.exit(exitcode)

def prVersion():
	print("Razer device configuration tool")
	print("Version", RAZER_VERSION)

def usage():
	prVersion()
	print("")
	print("Usage: razercfg [OPTIONS] [-d DEV DEVOPS] [-d DEV DEVOPS]...")
	print("")
	print("-h|--help            Print this help text")
	print("-v|--version         Print the program version number")
	print("-B|--background      Fork into the background")
	print("-s|--scan            Scan for devices and print the bus IDs")
	print("-K|--reconfigure     Force-reconfigure all detected devices")
	print("")
	print("-d|--device DEV      Selects the device with the bus ID \"DEV\"")
	print("    Use the special value \"mouse\" for DEV to select")
	print("    the first razer mouse device found in the system.")
	print("    If this option is omitted, the first Razer device found is selected.")
	print("")
	print("-S|--sleep SECS      Sleep SECS seconds.")
	print("")
	print("Device operations (DEVOPS):")
	print("These options apply to the device that is specified with -d")
	print("")
	print("Options for mice:")
	print("-V|--fwver                          Print the firmware version number")
	print("-p|--profile PROF                   Changes the active profile")
	print("-P|--getprofile                     Prints the active profile")
	print("-r|--res [PROF:]RES[xRES]           Changes the scan resolution")
	print("-R|--getres                         Prints the resolutions")
	print("-f|--freq [PROF:]FREQ               Changes the scan frequency")
	print("-F|--getfreq                        Prints the frequencies")
	print("-L|--leds                           List the device supported LED identifiers")
	print("-l|--setled [PROF:]LED:(on|off)     Toggle the LED with the identifier \"LED\"")
	print("                                    Use the special identifier \"all\"")
	print("                                    to toggle all supported LEDs.")
	print("-c|--setledcolor [PROF:]LED:rrggbb  Set LED color to RGB 'rrggbb'")
	print("-m|--setledmode  [PROF:]LED:MODE    Set LED mode to MODE ('static', 'spectrum'")
	print("                                                               or 'breathing')")
	print("")
	print("-X|--flashfw FILE                   Flash a firmware image to the device")
	print("")
	print("The profile number \"PROF\" may be 0 for the current profile. If omitted,")
	print("the global settings are changed (not possible for every device).")

def findDevice(deviceType=None):
	if deviceType is None or deviceType == "mouse":
		getRazer().rescanMice()
		mice = getRazer().getMice()
		if mice:
			return mice[0] # Return the first idstr
		if deviceType:
			raise RazerEx("No Razer mouse found in the system")
	raise RazerEx("No Razer device found in the system")

def parse_args():
	devOpsList = []
	currentDevOps = None

	try:
		(opts, args) = getopt.getopt(sys.argv[1:],
			"hvBsKd:r:Rf:FLl:VS:X:c:p:Pm:",
			[ "help", "version", "background",
			  "scan", "reconfigure", "device=", "res=",
			  "getres", "freq=", "getfreq", "leds", "setled=",
			  "fwver", "config=", "sleep=", "flashfw=",
			  "setledcolor=", "setledmode=",
			  "profile=", "getprofile", ])
	except getopt.GetoptError:
		usage()
		exit(1)

	for (o, v) in opts:
		if o in ("-h", "--help"):
			usage()
			exit(0)
		if o in ("-v", "--version"):
			prVersion()
			exit(0)
		if o in ("-B", "--background"):
			if os.fork() != 0:
				exit(0) # Exit parent
		if o in ("-s", "--scan"):
			ops = currentDevOps
			if not currentDevOps:
				ops = DevOps(None)
			ops.add(OpScan())
			if not currentDevOps:
				devOpsList.append(ops)
			continue
		if o in ("-K", "--reconfigure"):
			ops = currentDevOps
			if not currentDevOps:
				ops = DevOps(None)
			ops.add(OpReconfigure())
			if not currentDevOps:
				devOpsList.append(ops)
			continue
		if o in ("-d", "--device"):
			if v == "mouse": # magic; select the first mouse
				v = findDevice("mouse")
			if currentDevOps and currentDevOps.ops:
				devOpsList.append(currentDevOps)
			currentDevOps = DevOps(v)
			continue
		if o in ("-p", "--profile"):
			if not currentDevOps:
				currentDevOps = DevOps(findDevice())
			currentDevOps.add(OpSetProfile(v))
			continue
		if o in ("-P", "--getprofile"):
			if not currentDevOps:
				currentDevOps = DevOps(findDevice())
			currentDevOps.add(OpGetProfile())
			continue
		if o in ("-r", "--res"):
			if not currentDevOps:
				currentDevOps = DevOps(findDevice())
			currentDevOps.add(OpSetRes(v))
			continue
		if o in ("-R", "--getres"):
			if not currentDevOps:
				currentDevOps = DevOps(findDevice())
			currentDevOps.add(OpGetRes())
			continue
		if o in ("-f", "--freq"):
			if not currentDevOps:
				currentDevOps = DevOps(findDevice())
			currentDevOps.add(OpSetFreq(v))
			continue
		if o in ("-F", "--getfreq"):
			if not currentDevOps:
				currentDevOps = DevOps(findDevice())
			currentDevOps.add(OpGetFreq())
			continue
		if o in ("-L", "--leds"):
			if not currentDevOps:
				currentDevOps = DevOps(findDevice())
			currentDevOps.add(OpPrintLeds())
			continue
		if o in ("-l", "--setled"):
			if not currentDevOps:
				currentDevOps = DevOps(findDevice())
			currentDevOps.add(OpSetLedState(v))
			continue
		if o in ("-c", "--setledcolor"):
			if not currentDevOps:
				currentDevOps = DevOps(findDevice())
			currentDevOps.add(OpSetLedColor(v))
			continue
		if o in ("-m", "--setledmode"):
			if not currentDevOps:
				currentDevOps = DevOps(findDevice())
			currentDevOps.add(OpSetLedMode(v))
			continue
		if o in ("-V", "--fwver"):
			if not currentDevOps:
				currentDevOps = DevOps(findDevice())
			currentDevOps.add(OpGetFwVer())
			continue
		if o in ("-S", "--sleep"):
			ops = currentDevOps
			if not currentDevOps:
				ops = DevOps(None)
			try:
				v = float(v)
			except ValueError:
				raise RazerEx("Value for -S|--sleep must be a floating point value")
			ops.add(OpSleep(v))
			if not currentDevOps:
				devOpsList.append(ops)
			continue
		if o in ("-X", "--flashfw"):
			if not currentDevOps:
				raise RazerEx("Must specify a device (-d) before -X|--flashfw")
			currentDevOps.add(OpFlashFw(v))
			continue
	if currentDevOps and currentDevOps.ops:
		devOpsList.append(currentDevOps)
	if not devOpsList:
		usage()
		exit(1)
	return devOpsList

def main():
	try:
		devOpsList = parse_args()
		for devOps in devOpsList:
			devOps.runAll()
	except (RazerEx) as e:
		print(e)
		return 1
	return 0

if __name__ == "__main__":
	exit(main())