#!/usr/bin/env python3

import argparse
from loguru import logger
import os
import platform
import subprocess


__version__ = '1.0.9'
IMAGE_NAME = 'gasparka/spectrogram{}'.format(":arm" if os.uname()[4].startswith("arm") or os.uname()[4].startswith("aarch")  else "")


def warn_temperature():
    logger.warning(
        'Cool your LimeSDR or risk damage!')


def docker_pull():
    logger.info('Pulling latest docker image...')
    try:
        subprocess.run('docker pull {}'.format(IMAGE_NAME),
                       shell=True,
                       check=True)
    except:  # e.g. when there is no network connection
        pass


def docker_installed():
    try:
        subprocess.run('docker --version', shell=True, check=True, stdout=subprocess.PIPE)
        return True
    except:
        return False


def docker_image_running(name):
    try:
        res = subprocess.run('docker inspect --format="{{.State.Running}}" ' + name, shell=True, check=True,
                             stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        return bool(res)
    except:
        # assume that error is because image does not exist i.e. is not running
        return False


def docker_start_gqrx():
    logger.info('Starting gqrx...')
    warn_temperature()
    subprocess.run('xhost +local:docker', shell=True, check=False)
    subprocess.run(
        'docker run -i --net=host --privileged --env="DISPLAY" --volume="$HOME/.Xauthority:/root/.Xauthority:rw" --name gqrx --rm {} gqrx'.format(
            IMAGE_NAME),
        shell=True,
        check=True)


def docker_start_server_daemon():
    logger.info('Starting server daemon...')
    subprocess.run('docker kill spectrogram_server', shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    subprocess.run(
        'docker run -d --name spectrogram_server --rm --privileged --net=host  {} SoapySDRServer --bind'.format(
            IMAGE_NAME),
        shell=True,
        check=True)


def docker_program_fpga():
    logger.info('Programming FPGA, takes ~20 sec...')
    subprocess.run(
        'docker run -i --privileged {} LimeUtil --fpga=/LimeSDR-Mini_lms7_trx_HW_1.2_auto.rpd'.format(IMAGE_NAME),
        shell=True,
        check=True)


def docker_restore_fpga():
    logger.info('Restoring FPGA, takes ~20 sec...')
    subprocess.run('docker run -i --privileged {} LimeUtil --update'.format(IMAGE_NAME),
                   shell=True,
                   check=True)


def probe_devices():
    logger.info('Probing for LimeSDR-Mini devices...')
    res = subprocess.run(
        'docker run -i --privileged --net=host {} SoapySDRUtil --probe="driver=remote"'.format(IMAGE_NAME),
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE)
    output = res.stdout.decode()

    remote_available = 'LimeSDR-Mini' in output and 'gatewareVersion=1.0' in output
    if remote_available:
        return remote_available, False, False

    # probe for local device
    res = subprocess.run(
        'docker run -i --privileged --net=host {} SoapySDRUtil --probe="driver=lime"'.format(IMAGE_NAME),
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE)
    output = res.stdout.decode()

    local_available = 'LimeSDR-Mini' in output
    local_fpga_ok = 'gatewareVersion=1.0' in output
    logger.info('remote_available={}, local_available={}, local_fpga_ok={}'.format(remote_available, local_available,
                                                                                local_fpga_ok))
    return remote_available, local_available, local_fpga_ok


def main():
    warn_temperature()
    parser = argparse.ArgumentParser(description='Spectrogram (80MHz bandwidth) accelerator for LimeSDR')
    parser.add_argument('--server_only', help='Used for remote applications.', action='store_true')
    parser.add_argument('--fpga_restore', help='Runs `LimeUtil --update`.', action='store_true')
    parser.add_argument('--force_fpga', help='In case automatic detection fails.', action='store_true')
    args = vars(parser.parse_args())

    if platform.system() != 'Linux':
        logger.error('Sorry, this currently only works on Linux. Data rates are low, so you could try Virtual Machine.')
        return

    if not docker_installed():
        logger.error('Docker is not installed, cant continue. Fix: \n'
                  '$ curl -fsSL https://get.docker.com | sh \n'
                  '$ sudo groupadd docker \n'
                  '$ sudo usermod -aG docker $USER \n'
                  'Log out and log back in so that your group membership is re-evaluated.')
        return

    docker_pull()

    remote_available, local_available, local_fpga_ok = probe_devices()

    if args['force_fpga']:
        local_fpga_ok = False

    if args['server_only']:
        if not local_available:
            logger.info('No local device found, cant start the server')
            return
        if not local_fpga_ok:
            docker_program_fpga()
        docker_start_server_daemon()
        logger.info('Server is started!')
        return

    if args['fpga_restore']:
        if not local_available:
            logger.info('No local device found, cant restore the FPGA.')
            return
        docker_restore_fpga()
        return

    # default behaviour (no args)
    if remote_available:
        docker_start_gqrx()
    elif local_available:
        if not local_fpga_ok:
            docker_program_fpga()

        docker_start_server_daemon()
        try:
            docker_start_gqrx()
        except:
            pass  # yolo

        logger.info('Killing local server...')
        subprocess.run('docker kill spectrogram_server', shell=True)

    subprocess.run('xhost -local:docker', shell=True, check=False)
    logger.info('Done!')


if __name__ == '__main__':
    main()
