#! /usr/bin/python3 -sP

# Copyright 2015-2016, 2020 Unrud <unrud@outlook.com>
#
# 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/>.

import argparse
import ctypes
import logging
import os
import platform
import re
import subprocess
import sys
import traceback

__all__ = ['init', 'set_mute', 'set_effect', 'recovery']

VERSION = '0.14'
SUPPORTED_I2C_BUS_NAMES = ['SMBus I801 adapter']
I2C_CLASS_PATH = '/sys/class/i2c-dev'
DEV_PATH = '/dev'
CMDLINE_PATH = '/proc/cmdline'
MODPROBE_EXE = '/usr/sbin/modprobe'
KERNEL_PARAMETER = 'acpi_enforce_resources=lax'
KERNEL_PARAMETER_BEFORE_VERSION = (4, 15)
REQUIRED_MODULES = ['i2c_dev', 'i2c_i801']
DEVICE_ADDRESS = 0x73
DATA_DISABLE_OUTPUT = [
    # CMD Value
    [0x00, 0x86],
]
DATA_ENABLE_OUTPUT = [
    # CMD Value
    [0x00, 0x82],
]
DATA_EFFECTS = [list(zip([0x04, 0x05, 0x07, 0x08, 0x09], x)) for x in [
    # Values for CMD
    #0x04  0x05  0x07  0x08  0x09
    [0x11, 0x02, 0x22, 0x82, 0x22],
    [0xee, 0x03, 0x40, 0x84, 0xff],
    [0xaa, 0x23, 0x40, 0x84, 0x00],
    [0xaa, 0x22, 0x33, 0x84, 0x00],
    [0x88, 0x03, 0x23, 0x82, 0x22],
    [0xaa, 0x23, 0x41, 0x84, 0x00],
    [0xaa, 0x02, 0x43, 0x82, 0x00],
]]
EFFECTS_INFO = {
    0: 'no change',
    1: 'bass boost',
    4: 'boost everything',
}
DATA_RECOVERY = [
    # CMD Value
    [0x0b, 0x82],
    [0x0b, 0x92],
]
DEFAULT_EFFECT = 1
I2C_SLAVE = 0x0703
I2C_SMBUS = 0x0720
I2C_SMBUS_BLOCK_MAX = 32
I2C_SMBUS_READ = 1
I2C_SMBUS_WRITE = 0
I2C_SMBUS_BYTE_DATA = 2


class i2c_smbus_data(ctypes.Union):
    _fields_ = [('byte', ctypes.c_ubyte),
                ('word', ctypes.c_ushort),
                ('block', ctypes.c_ubyte * (I2C_SMBUS_BLOCK_MAX + 2))]


class i2c_smbus_ioctl_data(ctypes.Structure):
    _fields_ = [('read_write', ctypes.c_ubyte),
                ('command', ctypes.c_ubyte),
                ('size', ctypes.c_uint),
                ('data', ctypes.POINTER(i2c_smbus_data))]


class SMBus(object):
    def __init__(self, path):
        self.__logger = logging.getLogger('SMBus')
        self.__logger.info('Opening I2C bus: %s', path)
        try:
            self.__libc = ctypes.cdll.LoadLibrary('libc.so.6')
        except OSError:
            self.__logger.error('Can\'t load library: %s', 'libc.so.6')
            raise
        try:
            self.__fd = os.open(path, os.O_RDWR)
        except OSError:
            self.__logger.error('Can\'t open file: %s', path)
            raise
        self.__address = None

    def __enter__(self):
        return self

    def __exit__(self, _type, value, _traceback):
        self.close()

    @property
    def address(self):
        return self.__address

    @address.setter
    def address(self, address):
        self.__logger.info('Setting I2C slave address: %d', address)
        if self.__address != address:
            err = self.__libc.ioctl(self.__fd, I2C_SLAVE, address)
            if err != 0:
                self.__logger.error('Can\'t set I2C slave address')
                raise OSError(err, os.strerror(err))
            self.__address = address

    def __access(self, read_write, device_cmd, size, data):
        if self.__address is None:
            self.__logger.error('No I2C slave address set')
            raise RuntimeError('No I2C slave address set')
        args = i2c_smbus_ioctl_data()
        args.read_write = read_write
        args.command = device_cmd
        args.size = size
        args.data = ctypes.pointer(data)
        err = self.__libc.ioctl(self.__fd, I2C_SMBUS, ctypes.byref(args))
        if err != 0:
            self.__logger.error('Can\'t transfer data on I2C bus')
            raise OSError(err, os.strerror(err))

    def write_byte_data(self, device_cmd, value):
        self.__logger.info('Writing byte data on I2C bus: '
                           '(device_cmd: 0x%x, value: 0x%x)',
                           device_cmd, value)
        data = i2c_smbus_data()
        data.byte = value
        self.__access(I2C_SMBUS_WRITE, device_cmd, I2C_SMBUS_BYTE_DATA, data)

    def read_byte_data(self, device_cmd):
        data = i2c_smbus_data()
        self.__access(I2C_SMBUS_READ, device_cmd, I2C_SMBUS_BYTE_DATA, data)
        result = data.byte & 0xff
        self.__logger.info('Read byte data on I2C bus: '
                           '(device_cmd: 0x%x, value: 0x%x)',
                           device_cmd, result)
        return result

    def close(self):
        self.__logger.info('Closing I2C bus')
        os.close(self.__fd)


def check_root():
    if os.geteuid() != 0:
        logging.warning('This program needs root privileges')


def check_cmdline():
    release = platform.release()
    try:
        kernel_version = tuple(map(int, re.split('[.-]', release)[:2]))
        if kernel_version >= KERNEL_PARAMETER_BEFORE_VERSION:
            return
    except ValueError:
        logging.warning('Can\'t find kernel version: %s', release)
    try:
        with open(CMDLINE_PATH, 'r') as f:
            cmdline_parameters = f.read().split()
    except IOError:
        logging.warning('Can\'t open file: %s', CMDLINE_PATH)
        return
    if KERNEL_PARAMETER not in cmdline_parameters:
        logging.warning('Kernel parameter is missing: %s', KERNEL_PARAMETER)


def check_modules():
    try:
        for module in REQUIRED_MODULES:
            logging.info('Trying to add module to the kernel: %s', module)
            if subprocess.call([MODPROBE_EXE, '--quiet', module]) != 0:
                logging.warning('Module is not loaded: %s', module)
    except OSError:
        logging.warning('modprobe not found')


def get_i2c_busses():
    busses = []
    try:
        i2c_directories = os.listdir(I2C_CLASS_PATH)
    except IOError:
        logging.error('Can\'t list directory: %s', I2C_CLASS_PATH)
        raise
    for i2c_dev in i2c_directories:
        path = os.path.join(I2C_CLASS_PATH, i2c_dev, 'name')
        try:
            with open(path) as f:
                i2c_bus_name = f.read().strip()
        except IOError:
            logging.warning('Can\'t open file: %s', path)
            continue
        busses.append((i2c_bus_name, i2c_dev))
    return busses


def do_checks_and_get_i2c_bus():
    check_root()
    check_cmdline()
    check_modules()
    i2c_busses = get_i2c_busses()
    selected_i2c_dev = None
    selected_i2c_bus_name = None
    logging.debug('Available i2c busses: %s', [x[0] for x in i2c_busses])
    logging.debug('Supported i2c bus names: %s', SUPPORTED_I2C_BUS_NAMES)
    for i2c_bus_name, i2c_dev in i2c_busses:
        for supported_name in SUPPORTED_I2C_BUS_NAMES:
            if supported_name in i2c_bus_name:
                selected_i2c_dev = i2c_dev
                selected_i2c_bus_name = i2c_bus_name
    if not selected_i2c_dev:
        logging.error('Can\'t find i2c bus')
        raise RuntimeError('Can\'t find i2c bus')
    logging.debug('Selected i2c bus: %s', selected_i2c_bus_name)
    return SMBus(os.path.join(DEV_PATH, selected_i2c_dev))


def write_prolog(i2c_bus):
    i2c_bus.write_byte_data(0x0a, 0x41)
    for device_cmd in [0x04, 0x09]:
        value = i2c_bus.read_byte_data(device_cmd)
        i2c_bus.write_byte_data(device_cmd, value)


def write_data_to_device(data):
    with do_checks_and_get_i2c_bus() as i2c_bus:
        i2c_bus.address = DEVICE_ADDRESS
        write_prolog(i2c_bus)
        for device_cmd, device_data in data:
            i2c_bus.write_byte_data(device_cmd, device_data)


def init():
    set_effect(DEFAULT_EFFECT)


def set_mute(value):
    if value:
        write_data_to_device(DATA_DISABLE_OUTPUT)
    else:
        write_data_to_device(DATA_ENABLE_OUTPUT)


def set_effect(i):
    if i < 0 or i >= len(DATA_EFFECTS):
        logging.error('Invalid effect')
        raise ValueError('Invalid effect')
    write_data_to_device(DATA_DISABLE_OUTPUT +
                         DATA_EFFECTS[i] +
                         DATA_ENABLE_OUTPUT)


def recovery():
    write_data_to_device(DATA_RECOVERY)


def parse_args():
    commands_help = {
        'init': 'initialize amplifier (with effect%d)' % DEFAULT_EFFECT,
        'mute': 'turn output off',
        'unmute': 'turn output on',
    }
    for i in EFFECTS_INFO:
        commands_help['effect%d' % i] = EFFECTS_INFO[i]
    commands = (['init'] + ['effect%d' % i for i in range(len(DATA_EFFECTS))] +
                ['mute', 'unmute', 'recovery'])
    epilog = '\navailable commands:\n'
    for command in commands:
        command_help = commands_help.get(command)
        if command_help:
            epilog += '  %- 15s%s\n' % (command, command_help)
        else:
            epilog += '  %s\n' % command
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description='Manage the headphone amplifier found in some '
                    'Clevo laptops',
        epilog=epilog)
    parser.add_argument('--version', action='version',
                        version='%%(prog)s %s' % VERSION)
    parser.add_argument('-v', '--verbose', help='increase output verbosity',
                        action='store_true')
    parser.add_argument('-f', '--force', action='store_true',
                        help=argparse.SUPPRESS)
    parser.add_argument('command', nargs='?', choices=commands,
                        metavar='command', default='init',
                        help='see the list of available commands below, '
                             'init is the default if the argument is omitted')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    logging.info('Version: %s', VERSION)
    if args.force:
        logging.warning('The -f, --force argument is deprecated')
    command = args.command
    if command == 'init':
        init()
    elif command.startswith('effect'):
        i = int(command[len('effect'):])
        set_effect(i)
    elif command == 'mute':
        set_mute(True)
    elif command == 'unmute':
        set_mute(False)
    elif command == 'recovery':
        recovery()


class ColorStreamHandler(logging.StreamHandler):
    def emit(self, record):
        try:
            FORMAT_SEQ = '\033[%dm'
            fmt = 0  # reset
            if record.levelno >= logging.ERROR:
                fmt = 31  # red
            elif record.levelno >= logging.WARNING:
                fmt = 33  # yellow
            msg = self.format(record)
            if hasattr(self.stream, 'isatty') and self.stream.isatty():
                fs = (FORMAT_SEQ + '%s' + FORMAT_SEQ + '\n') % (fmt, msg, 0)
            else:
                fs = '%s\n' % msg
            self.stream.write(fs)
            self.flush()
        except Exception:
            self.handleError(record)


def setup_logging():
    ch = ColorStreamHandler(sys.stderr)
    fmt = logging.Formatter('%(levelname)s:%(message)s')
    ch.setFormatter(fmt)
    logging.getLogger().addHandler(ch)


if __name__ == '__main__':
    setup_logging()
    try:
        main()
    except Exception:
        logging.error('Operation failed')
        logging.debug('Exception occurred:\n%s', traceback.format_exc())
        exit(1)
    except KeyboardInterrupt:
        logging.error('Operation cancelled')
        exit(1)
