#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# 2020 - Marco Trevisan
#
# Initializer for ThinkPad's validity sensors 138a:0090 and 138a:0097 and 06cb:009a
#
# 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 os
import subprocess
import sys
import tempfile
import urllib.request
import shutil

from enum import Enum, auto
from time import sleep
from usb import core as usb_core

python_validity_data='/usr/share/python-validity/'

# FIXME: supported usb ids are duplicated in usb.py
class VFS(Enum):
    DEV_90 = (0x138a, 0x0090)
    DEV_97 = (0x138a, 0x0097)
    DEV_9a = (0x06cb, 0x009a)

DEFAULT_URIS = {
    VFS.DEV_90: {
        'driver': 'https://download.lenovo.com/pccbbs/mobiles/n1cgn08w.exe',
        'referral': 'https://support.lenovo.com/us/en/downloads/DS120491',
    },
    VFS.DEV_97: {
        'driver': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe',
        'referral': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe'
    },
    VFS.DEV_9a: {
        'driver': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe',
        'referral': 'https://download.lenovo.com/pccbbs/mobiles/nz3gf07w.exe'
    }
}

# FIXME: filenames are duplicated in upload_fwext.py
DEFAULT_FW_NAMES = {
    VFS.DEV_90: '6_07f_Lenovo.xpfwext',
    VFS.DEV_97: '6_07f_lenovo_mis_qm.xpfwext',
    VFS.DEV_9a: '6_07f_lenovo_mis_qm.xpfwext'
}


def download_and_extract_fw(dev_type, fwdir, fwuri=None):
    fwuri = fwuri if fwuri else DEFAULT_URIS[dev_type]['driver']
    fwarchive = os.path.join(fwdir, 'fwinstaller.exe')
    fwname = DEFAULT_FW_NAMES[dev_type]

    print('Downloading {} to extract {}'.format(fwuri, fwname))

    req = urllib.request.Request(fwuri)
    req.add_header('Referer', DEFAULT_URIS[dev_type].get('referral', ''))
    req.add_header('User-Agent', 'Mozilla/5.0 (X11; U; Linux)')

    with urllib.request.urlopen(req) as response:
        with open(fwarchive, 'wb') as out_file:
            out_file.write(response.read())

    subprocess.check_call(['innoextract',
        '--output-dir', fwdir,
        '--include', fwname,
        '--collisions', 'overwrite',
        fwarchive
    ])

    fwpath = subprocess.check_output([
        'find', fwdir, '-name', fwname]).decode('utf-8').strip()
    print('Found firmware at {}'.format(fwpath))

    if not fwpath:
        raise Exception('No {} found in the archive'.format(fwname))

    return fwpath

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--driver-uri')

    args = parser.parse_args()

    if os.geteuid() != 0:
        raise Exception('This script needs to be executed as root')

    dev_type = None
    for d in VFS:
        dev = usb_core.find(idVendor=d.value[0], idProduct=d.value[1])
        if dev:
            dev_type = d

    if not dev_type:
        raise Exception('No supported validity device found')

    try:
        subprocess.check_call(['innoextract', '--version'],
            stdout=subprocess.DEVNULL)
    except Exception as e:
        print('Impossible to run innoextract: {}'.format(e))
        sys.exit(1)

    with tempfile.TemporaryDirectory() as fwdir:
        fwpath = download_and_extract_fw(dev_type, fwdir, fwuri=args.driver_uri)
        shutil.copy(fwpath, python_validity_data)