#!/usr/bin/env python3

from pathlib import Path
from argparse import ArgumentParser
import json
import os
import shutil
import subprocess
import time
import sys
import tempfile

eva_dir = os.environ.get('EVA_DIR')
if eva_dir is None:
    eva_dir = Path(__file__).absolute().parents[1]
else:
    eva_dir = Path(eva_dir)

mods_list = os.environ.get('MODS_LIST')
if mods_list is None:
    mods_list = eva_dir / 'install/mods.list'

is_tty = sys.stdout.isatty()

SLEEP_STEP = 0.05

CAROUSEL = ['-', '\\', '|', '/']

verbose = False

BLACK = '\033[30m'
RED = '\033[31m'
GREEN = '\033[32m'
YELLOW = '\033[33m'
BLUE = '\033[34m'
MAGENTA = '\033[35m'
CYAN = '\033[36m'
ENDC = '\033[0m'
BOLD = '\033[1m'

MODS_FIRST = ['pytest-runner', 'cffi', 'setuptools', 'setuptools-rust']


def split_mod_ver(module):
    try:
        mod, version = module.split('==', 1)
        version = '==' + version
    except:
        try:
            mod, version = module.split('>=', 1)
            version = '>=' + version
        except:
            try:
                mod, version = module.split('=>', 1)
                version = '>=' + version
            except:
                mod = module
                version = ''
    return (mod, version)


def cprint(text, colors='', **kwargs):
    if is_tty:
        print(colors + str(text) + ENDC, **kwargs)
    else:
        print(text, **kwargs)


def bs():
    print('\033[D', end='')


def l_startswith(s, arr, sfx=''):
    for a in arr:
        if s.startswith(a + sfx):
            return True
    return False


def get_config():
    p = subprocess.run(
        [f'{eva_dir}/sbin/eva-registry-cli', 'get', 'eva/config/python-venv'],
        stdout=subprocess.PIPE)
    if p.returncode:
        raise RuntimeError(
            f'unable to read the registry. the node is unavailable')
    return json.loads(p.stdout)


def set_config(config):
    payload = json.dumps(config).encode()
    p = subprocess.Popen([
        f'{eva_dir}/sbin/eva-registry-cli', 'set', 'eva/config/python-venv',
        '-', '-p', 'json'
    ],
                         stdin=subprocess.PIPE,
                         stdout=subprocess.PIPE)
    p.communicate(input=payload)
    if p.returncode:
        raise RuntimeError(
            f'unable to read the registry. the node is unavailable')


def xc(cmd, ps='working', warn=False):

    def print_st(st, color=''):
        data = st.read()
        if data:
            try:
                cprint(data.decode(), color)
            except:
                cprint(data, color)

    if verbose:
        print(ps)
        p = subprocess.run(cmd, shell=True)
    else:
        print(ps, end='  ' if is_tty else '... ', flush=True)
        p = subprocess.Popen(cmd,
                             shell=True,
                             stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)
        c_idx = 0
        while p.poll() is None:
            if is_tty:
                bs()
                print(CAROUSEL[c_idx], end='', flush=True)
                c_idx += 1
                if c_idx == len(CAROUSEL):
                    c_idx = 0
            time.sleep(SLEEP_STEP)
        if is_tty:
            bs()
        if p.returncode == 0:
            cprint('OK', GREEN, flush=True)
        elif warn:
            cprint('WARN', YELLOW + BOLD, flush=True)
        else:
            cprint('FAILED!', RED + BOLD, flush=True)
            print_st(p.stdout)
            print_st(p.stderr, RED)
    if p.returncode != 0:
        raise RuntimeError(f'process exited with the code {p.returncode}')


class App:

    def build(self, from_scratch=False):
        with open(mods_list) as fh:
            mods = [x for x in [x.strip() for x in fh.readlines()] if x]
        venv_dir = eva_dir / 'venv'
        if from_scratch and venv_dir.exists():
            cprint('Removing venv dir', YELLOW + BOLD)
            shutil.rmtree(venv_dir)
        config = get_config()
        python = config.get('python', 'python3')
        system_pip = config.get('use_system_pip', False)
        skip = config.get('skip', [])
        extra = config.get('extra', [])
        pip_xo = config.get('pip_extra_options', '')
        ssp = config.get('system_site_packages', True)
        mods += extra
        if not venv_dir.exists():
            cmd = f'{python} -m venv '
            if not system_pip:
                cmd += '--without-pip '
            if ssp:
                cmd += '--system-site-packages '
            cmd += f'{eva_dir}/venv'
            xc(cmd, ps='Preparing Python venv')
        pip3 = venv_dir / 'bin/pip3'
        if not pip3.exists():
            pip3_url = 'https://bootstrap.pypa.io/get-pip.py'
            if sys.version_info.major != 3:
                raise RuntimeError('Python 3 is required')
            if sys.version_info.minor < 9:
                pip3_url = f'https://bootstrap.pypa.io/pip/3.{sys.version_info.minor}/get-pip.py'
            xc(
                'curl --connect-timeout 5 '
                f'{pip3_url} |'
                f' "{eva_dir}/venv/bin/python"',
                ps='Downloading pip3')
        if not pip3.exists():
            raise RuntimeError('pip3 installation failed')
        mods = [x for x in mods if not l_startswith(x, skip, '=')]
        mf = [x for x in mods if l_startswith(x, MODS_FIRST, '=')]
        mods = [x for x in mods if not l_startswith(x, MODS_FIRST, '=')]
        for (m, p) in [(mf, 'Installing/updating setup modules'),
                       (mods, 'Installing/updating venv modules')]:
            with tempfile.NamedTemporaryFile('w') as fp:
                fp.write('\n'.join(m))
                fp.flush()
                xc(
                    f'[ -f ~/.cargo/env ] && . ~/.cargo/env ;'
                    f' {pip3} install -U {pip_xo} -r {fp.name}', p)
        xc('sync', ps='Flushing data to disk')
        yedb_sys = Path('/usr/local/bin/yedb')
        yedb_venv = venv_dir / 'bin/yedb'
        if yedb_sys.exists() and not yedb_venv.exists():
            os.symlink(yedb_sys, yedb_venv)
        print('VENV rebuilt successfully')
        print()

    def list(self):
        pip3 = eva_dir / 'venv/bin/pip3'
        if not pip3.exists():
            raise RuntimeError('venv not configured')
        p = subprocess.run([pip3, 'list'],
                           stdout=subprocess.PIPE,
                           stderr=subprocess.PIPE)
        if p.stdout:
            mods = p.stdout.decode().split('\n')
            if len(mods) > 1:
                cprint(mods[0], BLUE)
                cprint(mods[1], BLACK)
                for m in mods[2:]:
                    try:
                        name, value = m.split(' ', 1)
                        cprint(name, CYAN, end=' ')
                        cprint(value, YELLOW)
                    except:
                        print(m)
        if p.stderr:
            print(p.stderr.decode())
        if p.returncode != 0:
            raise RuntimeError(f'process exited with the code {p.returncode}')

    def edit(self):
        os.system(
            f'{eva_dir}/sbin/eva-registry-cli edit eva/config/python-venv')

    def update(self, modules, rebuild):
        self.add(modules, rebuild, force_install=True)

    def add(self, modules, rebuild, force_install=False):
        config = get_config()
        extra = config.get('extra', [])
        mods = {}
        added = set()
        for x in extra:
            (mod, version) = split_mod_ver(x)
            mods[mod] = version
        for m in modules:
            is_new = True
            (mod, version) = split_mod_ver(m)
            if not force_install:
                try:
                    if mods[mod] == version:
                        is_new = False
                except KeyError:
                    pass
            mods[mod] = version
            cprint(f'+ {mod}{version}', GREEN)
            if is_new:
                added.add(f'{mod}{version}')
        extra = set()
        for k, v in mods.items():
            extra.add(f'{k}{v}')
        extra = sorted(extra)
        config['extra'] = extra
        set_config(config)
        pip3 = eva_dir / 'venv/bin/pip3'
        pip_xo = config.get('pip_extra_options', '')
        if added and pip3.exists():
            mod_list = ' '.join(added)
            xc(f'{pip3} install -U {pip_xo} {mod_list}',
               ps='Adding module files')
            xc('sync', ps='Flushing data to disk')
        if rebuild:
            self.build()

    def remove(self, modules, rebuild):
        config = get_config()
        extra = config.get('extra', [])
        mods = {}
        removed = set()
        for x in extra:
            (mod, version) = split_mod_ver(x)
            mods[mod] = version
        for m in modules:
            (mod, version) = split_mod_ver(m)
            cprint(f'- {mod}{version}', YELLOW)
            try:
                del mods[mod]
                removed.add(mod)
            except KeyError:
                pass
        extra = set()
        for k, v in mods.items():
            extra.add(f'{k}{v}')
        extra = sorted(extra)
        config['extra'] = extra
        set_config(config)
        pip3 = eva_dir / 'venv/bin/pip3'
        if removed and pip3.exists():
            mod_list = ' '.join(removed)
            xc(f'yes | {pip3} uninstall {mod_list}', ps='Removing module files')
            xc('sync', ps='Flushing data to disk')
        if rebuild:
            self.build()

    def config(self):
        config = get_config()
        if is_tty:
            import shutil
            jq = shutil.which('jq')
            if jq:
                p = subprocess.Popen([jq], stdin=subprocess.PIPE)
                p.communicate(input=json.dumps(config).encode())
                return
        print(json.dumps(config, indent=4, sort_keys=True))

    def mirror_update(self, dest):
        import platform
        config = get_config()
        pip3 = eva_dir / 'venv/bin/pip3'
        mirror_dir = Path(dest)
        mirror_pypi_dir = mirror_dir / 'pypi'
        ppm = eva_dir / 'cli/pypi-mirror'
        if not pip3.exists():
            import shutil
            pip3 = shutil.which('pip3')
            if not pip3:
                raise RuntimeError('pip3 command not found')
        mirror_pypi_dir.mkdir(parents=True, exist_ok=True)
        with open(mods_list) as fh:
            mods = [x.strip() for x in fh.readlines()]
        mods += ['pip']
        print(f'Updating PyPi mirror (destination: {mirror_dir})')
        print()
        mods_skip = config.get('skip', [])
        mods_extra = config.get('extra', [])
        mirror_extra_python_versions = config.get('mirror_extra_versions', [])
        for m in mods.copy():
            if m in mods_skip or m.split('=', 1)[0] in mods_skip:
                cprint(f'- {m}', BLACK)
                mods.remove(m)
        for m in mods_extra:
            cprint(f'+ {m}', GREEN)
            mods.append(m)
        if mods_skip or mods_extra:
            print()
        print(f'Python version: '
              f'{sys.version_info.major}.{sys.version_info.minor}')
        print(f'CPU architecture: {platform.uname().machine}')
        if mirror_extra_python_versions:
            print(f'Extra Python versions: '
                  f'{", ".join(mirror_extra_python_versions)}')
        print()
        print(f'Modules: {len(mods)}')
        cprint('-' * 40, BLACK)
        # update modules
        for mod in mods:
            if mod:
                xc(
                    f'{ppm} download -p "{pip3}" -b '
                    f'-d "{mirror_pypi_dir}/downloads" "{mod}"',
                    ps=f'Downloading {mod}')
        # update compiled mods for extra Python versions
        cmods = Path(mirror_pypi_dir).glob(f'**/*-cp*.whl')
        xmods = set()
        srcs_req = set()
        for c in cmods:
            m = c.name.split('-')
            mod_name, mod_version = m[0], m[1]
            xmods.add(f'{mod_name}=={mod_version}')
        for pyver in mirror_extra_python_versions:
            if pyver != 'source':
                for xmod in xmods:
                    try:
                        xc(
                            f'{ppm} download -p "{pip3}" '
                            f'-b -d "{mirror_pypi_dir}/downloads" '
                            f'--python-version "{pyver}" "{xmod}"',
                            ps=f'Downloading {xmod} for Python {pyver}',
                            warn=True)
                    except KeyboardInterrupt:
                        raise
                    except:
                        cprint(
                            f'No binary package for {xmod}, '
                            'will download sources', YELLOW + BOLD)
                        srcs_req.add(xmod)
        # download modules sources for missing binary mods or for all
        for s in srcs_req \
                if 'source' not in mirror_extra_python_versions \
                else mods:
            try:
                xc(
                    f'{ppm} download -p "{pip3}" -d '
                    f'"{mirror_pypi_dir}/downloads" "{s}"',
                    ps=f'Downloading sources for {s}')
            except KeyboardInterrupt:
                raise
            except:
                self.print_warn(f'Unable to download sources for {s}')
        # update mirror index
        xc(
            f'{ppm} create -d "{mirror_pypi_dir}/downloads" '
            f'-m "{mirror_pypi_dir}/local"',
            ps='Updating indexes')

    def mirror_set(self, mirror_url):
        config = get_config()
        if mirror_url == 'default':
            config['pip_extra_options'] = ''
            msg = 'pypi.org (default)'
        elif not mirror_url.startswith('http://') and not mirror_url.startswith(
                'https://'):
            raise RuntimeError('invalid mirror URL')
        else:
            try:
                mirror_host = mirror_url.split('://',
                                               1)[1].split('/',
                                                           1)[0].split(':')[0]
            except:
                raise RuntimeError('Invalid mirror URL')
            while mirror_url.endswith('/'):
                mirror_url = mirror_url[:-1]
            config['pip_extra_options'] = (f'-i {mirror_url}/pypi/local'
                                           f' --trusted-host {mirror_host}')
            msg = mirror_url
        set_config(config)
        print(f'PyPi mirror has been set to {msg}')


ap = ArgumentParser(description='EVA ICS v4 Python VENV manager')

ap.add_argument('--verbose', action='store_true')

sp = ap.add_subparsers(dest='_cmd', help='command')

p = sp.add_parser('build')
p.add_argument('-S', '--from-scratch', action='store_true')

sp.add_parser('list')

sp.add_parser('edit')

sp.add_parser('config')

p = sp.add_parser('add')
p.add_argument('modules', metavar='MODULE', nargs='+')
p.add_argument('-B',
               '--rebuild',
               action='store_true',
               help='automatically rebuild venv')

p = sp.add_parser('remove')
p.add_argument('modules', metavar='MODULE', nargs='+')
p.add_argument('-B',
               '--rebuild',
               action='store_true',
               help='automatically rebuild venv')

p = sp.add_parser('update')
p.add_argument('modules', metavar='MODULE', nargs='+')
p.add_argument('-B',
               '--rebuild',
               action='store_true',
               help='automatically rebuild venv')

mirror_dir = eva_dir / 'mirror'
p = sp.add_parser('mirror-update', help='Create/update PyPi mirror')
p.add_argument('--dest',
               metavar='DIR',
               help=f'Mirror directory (default: {mirror_dir})',
               default=mirror_dir)

p = sp.add_parser('mirror-set', help='Set EVA ICS v4 mirror URL')
p.add_argument('mirror_url',
               metavar='URL',
               help=('EVA ICS v4 mirrer url as http://<ip/host>:port'
                     ' or "default" to restore the default settings'))

a = ap.parse_args()

if a._cmd is None:
    ap.print_help()
    exit(0)

app = App()

kw = vars(a)
cmd = kw.pop('_cmd').replace('-', '_')
verbose = kw.pop('verbose')

try:
    getattr(app, cmd)(**kw)
except KeyboardInterrupt:
    cprint('\nOperation aborted', YELLOW + BOLD)
except Exception as e:
    cprint(e, RED + BOLD)
    sys.exit(1)
