#!/bin/env python
"""Utility to read/write logitech g600 mouse key maps.
Behaves like cp, ie, give it a source and destination.
In most cases, this requires root (ie, run sudo <this script>).
Mouse configurations are stored in a human readable json format.
Note: MOUSE is a special keyword that specifies the mouse rather than a file.

For example, to copy the current mouse config to a file called mouse_config.json:
$ sudo ./g600prog.py MOUSE mouse_config.json

Copying a custom_config.json file to the mouse:
$ sudo ./g600prog.py custom_config.json MOUSE

Leaving off the second argument causes it to print to stdout.
So, to print your current mouse config:
$ sudo ./g600prog.py MOUSE"""
from __future__ import print_function
import sys
import os
import argparse
import itertools
import json
import collections
import time
import usb.core
import usb.util


def main(argv):
    cfg = parseArgs(argv)
    if cfg.SOURCE == "MOUSE":
        mouseMapping = readMouseMappingFromMouse(cfg.debug)
    else:
        mouseMapping = readMouseMappingFromFile(cfg.SOURCE, cfg.debug)
    if cfg.bytes:
        mouseMappingBytes = G600MouseMappingBytes()
        mouseMappingBytes.fromModeRawBytesList(mouseMapping.toModeRawBytesList())
        mouseMapping = mouseMappingBytes
    if cfg.DESTINATION is None:
        print(mouseMapping)
    elif cfg.DESTINATION == "MOUSE":
        writeMouseMappingToMouse(mouseMapping, cfg.debug, cfg.dry_run)
    else:
        saveMouseMappingToFile(mouseMapping, cfg.DESTINATION, cfg.overwrite_file)


def readMouseMappingFromMouse(debug):
    print("Reading mouse config from mouse...")
    mouseMapping = G600MouseMapping()
    rawModeBytesList = readUsbMouseMappingRawBytes(debug)
    mouseMapping.fromModeRawBytesList(rawModeBytesList)
    print("... done reading mouse config from mouse")
    return mouseMapping


def readMouseMappingFromFile(fileName, debug):
    print("Reading mouse config from file >{}< ...".format(fileName))
    with open(fileName, 'r') as fileHandle:
        jsonObj = json.loads(fileHandle.read())
        mouseMapping = G600MouseMapping()
        if "configFormat" not in jsonObj:
            raise FromJsonError("missing configFormat!")
        if jsonObj["configFormat"] == "BytesFormat":
            mouseMappingBytes = G600MouseMappingBytes()
            mouseMappingBytes.simpleRepr = jsonObj
            mouseMapping.fromModeRawBytesList(mouseMappingBytes.toModeRawBytesList())
        elif jsonObj["configFormat"] == "HumanReadableFormat":
            mouseMapping.simpleRepr = jsonObj
        else:
            raise FromJsonError("Undefined configFormat >>{}<<".format(jsonObj["configFormat"]))
    print("... done reading mouse config from file")
    return mouseMapping


def saveMouseMappingToFile(mouseMapping, fileName, forceWrite):
    print("Saving the mouse config to file >{}< ...".format(fileName))
    if os.path.isfile(fileName) and not forceWrite:
        raise Exception("File already exists and overwrite-file flag not set")
    with open(fileName, "w") as fileHandle:
        fileHandle.write(mouseMapping.json)
    print("...done saving the mouse config to file")


def writeMouseMappingToMouse(mouseMapping, debug, dryRun):
    print("Writing the mouse config to the mouse...")
    rawModeBytesList = mouseMapping.toModeRawBytesList()
    writeUsbMouseMappingRawBytes(rawModeBytesList, debug, dryRun)
    print("...done writing read mouse config to the mouse")


def parseArgs(argv):
    description = __doc__
    parser = argparse.ArgumentParser(description=description, formatter_class=argparse.RawDescriptionHelpFormatter)
    if len(argv) == 1:
        argv.append('-h')

    parser.add_argument('SOURCE',
                        help='Configuration source, can be MOUSE for the mouse itself or a filename.',)
    parser.add_argument('DESTINATION', nargs='?', default=None,
                        help='Optional configuration destination, can be the MOUSE or filename.  If omitted, prints to stdout.',)
    parser.add_argument('-f', '--overwrite-file',
                        help='Normally, if the destination file already exists, the program will terminate.  This option forces the overwrite of DESTINATION even if it exists.',
                        action='store_true',)
    parser.add_argument('-n', '--dry-run',
                        help='For testing writes to the mouse, intended to be used in conjunction with debug, will do everything except for actually send the usb programming messages.',
                        action='store_true',)
    parser.add_argument('-d', '--debug',
                        help='Turn on debug printing.',
                        action='store_true',)
    parser.add_argument('--bytes',
                        help='Store output config in JSON byte array format.  This could be useful for moving betweeen versions of this app where the human readable JSON format changes.',
                        action='store_true',)

    cfg = parser.parse_args()
    return cfg

################################################################################
# usb read/write to the mouse control interface.
# Operates on a 3 element sequence where each element is a bytearray()
IDVENDOR = 0x046d
IDPRODUCT = 0xc24a
G600_CONTROL_INTERFACE = 1
G600_REPORT_IDS = (0x03f3, 0x03f4, 0x03f5)  # one for each of the three mouse "modes"

G600_READ_REQTYPE = 0xA1
G600_READ_REQ = 0x01
G600_READ_IDX = G600_CONTROL_INTERFACE
G600_READ_LENGTH = 154


def readUsbMouseMappingRawBytes(debug=False):
    """Returns three element list.
    One for each of the mouse "modes."
    Each list element is a bytearray() type.
    """
    if debug:
        print("About to read USB...")
    dev = usb.core.find(idVendor=IDVENDOR, idProduct=IDPRODUCT)
    if dev.is_kernel_driver_active(G600_CONTROL_INTERFACE) is True:
        # tell the kernel to detach
        dev.detach_kernel_driver(G600_CONTROL_INTERFACE)
        # claim the device
        usb.util.claim_interface(dev, G600_CONTROL_INTERFACE)
    modes = []
    for reportId in G600_REPORT_IDS:
        replyMsg = dev.ctrl_transfer(bmRequestType=G600_READ_REQTYPE,  # this means control
                                     bRequest=G600_READ_REQ,
                                     wValue=reportId,
                                     wIndex=G600_READ_IDX,
                                     data_or_wLength=G600_READ_LENGTH,
                                     timeout=None)
        if debug:
            print("for reportId=0x{:04x}, read these bytes: ".format(reportId),)
            print(" ".join("0x{:02x}".format(x) for x in replyMsg))
        modes.append(replyMsg)
    # release the device
    usb.util.release_interface(dev, G600_CONTROL_INTERFACE)
    # reattach the device to the OS kernel
    dev.attach_kernel_driver(G600_CONTROL_INTERFACE)
    # done
    if debug:
        print("...Done reading USB")
    return modes

G600_WRITE_REQTYPE = 0x21
G600_WRITE_REQ = 0x09
G600_WRITE_IDX = G600_CONTROL_INTERFACE


def writeUsbMouseMappingRawBytes(modes, debug=False, dryRun=True):
    """Argument should be a three element list.
    One for each of the mouse "modes."
    Each list element is a bytearray() type.
    """
    if debug:
        print("About to write USB...")
    dev = usb.core.find(idVendor=IDVENDOR, idProduct=IDPRODUCT)
    if dev.is_kernel_driver_active(G600_CONTROL_INTERFACE) is True:
        # tell the kernel to detach
        dev.detach_kernel_driver(G600_CONTROL_INTERFACE)
        # claim the device
        usb.util.claim_interface(dev, G600_CONTROL_INTERFACE)
    for reportId, rawBytes in zip(G600_REPORT_IDS, modes):
        if debug:
            print("for reportId=0x{:04x}, sending these bytes: ".format(reportId),)
            print(" ".join("0x{:02x}".format(x) for x in rawBytes))
        if dryRun:
            print("dryRun flag set, not sending usb config write message")
        else:
            l = dev.ctrl_transfer(bmRequestType=G600_WRITE_REQTYPE,  # this means control
                                  bRequest=G600_WRITE_REQ,
                                  wValue=reportId,
                                  wIndex=G600_WRITE_IDX,
                                  data_or_wLength=rawBytes,
                                  timeout=None)
            assert l == len(rawBytes)
            time.sleep(1.1)
    # release the device
    usb.util.release_interface(dev, G600_CONTROL_INTERFACE)
    # reattach the device to the OS kernel
    dev.attach_kernel_driver(G600_CONTROL_INTERFACE)
    if debug:
        print("...Done writing USB")
    # done
################################################################################

################################################################################
# raw scan code maps of known codes


def invMap(mapDict):
    return {v: k for k, v in mapDict.items()}

MOUSE_SCAN_CODES_DICT = {0x00: "NO_MOUSEBUT",
                         0x01: "LEFT_CLICK",
                         0x02: "RIGHT_CLICK",
                         0x03: "MIDDLE_CLICK",
                         0x04: "BACK",
                         0x05: "FORWARD",
                         0x06: "MOUSE6",
                         0x07: "MOUSE7",
                         0x11: "DPI_UP",
                         0x12: "DPI_DOWN",
                         0x13: "DPI_CYCLING",
                         0x14: "MODE_SWITCH",
                         0x15: "DPI_SHIFT",
                         0x16: "DPI_DEFAULT",
                         0x17: "GSHIFT",
                         }
MOUSE_SCAN_CODES_INVDICT = invMap(MOUSE_SCAN_CODES_DICT)  # for reverse lookup

KB_MODIFIER_BIT_CODES_DICT = {0: "LCTRL",
                              1: "LSHIFT",
                              2: "LALT",
                              3: "LGUI",
                              4: "RCTRL",
                              5: "RSHIFT",
                              6: "RALT",
                              7: "RGUI",
                              }
KB_MODIFIER_BIT_CODES_INVDICT = invMap(KB_MODIFIER_BIT_CODES_DICT)  # for reverse lookup

KB_SCAN_CODES_DICT = {0x00: "NOKEY",
                      0x01: "ROLL_OVER",
                      0x02: "POST_FAIL",
                      0x03: "ERRUNDEF",
                      0x04: "A",
                      0x05: "B",
                      0x06: "C",
                      0x07: "D",
                      0x08: "E",
                      0x09: "F",
                      0x0A: "G",
                      0x0B: "H",
                      0x0C: "I",
                      0x0D: "J",
                      0x0E: "K",
                      0x0F: "L",
                      0x10: "M",
                      0x11: "N",
                      0x12: "O",
                      0x13: "P",
                      0x14: "Q",
                      0x15: "R",
                      0x16: "S",
                      0x17: "T",
                      0x18: "U",
                      0x19: "V",
                      0x1A: "W",
                      0x1B: "X",
                      0x1C: "Y",
                      0x1D: "Z",
                      0x1E: "1",
                      0x1F: "2",
                      0x20: "3",
                      0x21: "4",
                      0x22: "5",
                      0x23: "6",
                      0x24: "7",
                      0x25: "8",
                      0x26: "9",
                      0x27: "0",
                      0x28: "ENTER",
                      0x29: "ESCAPE",
                      0x2A: "BSPACE",
                      0x2B: "TAB",
                      0x2C: "SPACE",
                      0x2D: "MINUS",
                      0x2E: "EQUAL",
                      0x2F: "LBRACKET",
                      0x30: "RBRACKET",
                      0x31: "BSLASH",
                      0x32: "NONUS_HASH",
                      0x33: "SCOLON",
                      0x34: "QUOTE",
                      0x35: "GRAVE",
                      0x36: "COMMA",
                      0x37: "DOT",
                      0x38: "SLASH",
                      0x39: "CAPSLOCK",
                      0x3A: "F1",
                      0x3B: "F2",
                      0x3C: "F3",
                      0x3D: "F4",
                      0x3E: "F5",
                      0x3F: "F6",
                      0x40: "F7",
                      0x41: "F8",
                      0x42: "F9",
                      0x43: "F10",
                      0x44: "F11",
                      0x45: "F12",
                      0x46: "PSCREEN",
                      0x47: "SCKLOCK",
                      0x48: "PAUSE",
                      0x49: "INSERT",
                      0x4A: "HOME",
                      0x4B: "PGUP",
                      0x4C: "DELETE",
                      0x4D: "END",
                      0x4E: "PGDOWN",
                      0x4F: "RIGHT",
                      0x50: "LEFT",
                      0x51: "DOWN",
                      0x52: "UP",
                      0x53: "NUMLOCK",
                      0x54: "KP_SLASH",
                      0x55: "KP_ASTERISK",
                      0x56: "KP_MINUS",
                      0x57: "KP_PLUS",
                      0x58: "KP_ENTER",
                      0x59: "KP_1",
                      0x5A: "KP_2",
                      0x5B: "KP_3",
                      0x5C: "KP_4",
                      0x5D: "KP_5",
                      0x5E: "KP_6",
                      0x5F: "KP_7",
                      0x60: "KP_8",
                      0x61: "KP_9",
                      0x62: "KP_0",
                      0x63: "KP_DOT",
                      0x64: "NON_US_BSLASH",
                      0x65: "APPLICATION",
                      0x66: "POWER",
                      0x67: "KP_EQUAL",
                      0x68: "F13",
                      0x69: "F14",
                      0x6A: "F15",
                      0x6B: "F16",
                      0x6C: "F17",
                      0x6D: "F18",
                      0x6E: "F19",
                      0x6F: "F20",
                      0x70: "F21",
                      0x71: "F22",
                      0x72: "F23",
                      0x73: "F24",
                      0x74: "EXECUTE",
                      0x75: "HELP",
                      0x76: "MENU",
                      0x77: "SELECT",
                      0x78: "STOP",
                      0x79: "AGAIN",
                      0x7A: "UNDO",
                      0x7B: "CUT",
                      0x7C: "COPY",
                      0x7D: "PASTE",
                      0x7E: "FIND",
                      0x7F: "MUTE",
                      0x80: "VOLUP",
                      0x81: "VOLDOWN",
                      0x82: "LOCKING_CAPS",
                      0x83: "LOCKING_NUM",
                      0x84: "LOCKING_SCROLL",
                      0x85: "KP_COMMA",
                      0x86: "KP_EQUAL_SIGN",
                      0x87: "INTERNATIONAL1",
                      0x88: "INTERNATIONAL2",
                      0x89: "INTERNATIONAL3",
                      0x8A: "INTERNATIONAL4",
                      0x8B: "INTERNATIONAL5",
                      0x8C: "INTERNATIONAL6",
                      0x8D: "INTERNATIONAL7",
                      0x8E: "INTERNATIONAL8",
                      0x8F: "INTERNATIONAL9",
                      0x90: "LANG1",
                      0x91: "LANG2",
                      0x92: "LANG3",
                      0x93: "LANG4",
                      0x94: "LANG5",
                      0x95: "LANG6",
                      0x96: "LANG7",
                      0x97: "LANG8",
                      0x98: "LANG9",
                      0x99: "ALT_ERASE",
                      0x9A: "SYSREQ",
                      0x9B: "CANCEL",
                      0x9C: "CLEAR",
                      0x9D: "PRIOR",
                      0x9E: "RETURN",
                      0x9F: "SEPARATOR",
                      0xA0: "OUT",
                      0xA1: "OPER",
                      0xA2: "CLEAR_AGAIN",
                      0xA3: "CRSEL",
                      0xA4: "EXSEL",
                      }
KB_SCAN_CODES_INVDICT = invMap(KB_SCAN_CODES_DICT)  # for reverse lookup

LIGHTING_EFFECT_DICT = {0x00: "NO_EFFECT",
                        0x01: "PULSE",
                        0x02: "RAINBOW",