import sys
import argparse
import docker
import dockerpty
import io
import getpass
import pathlib
import platform
import os
import re
import xdg.BaseDirectory
import yaml
import datetime
import pytz
import subprocess
from pprint import pprint


def find_config_file(user_conf):
    conf_list = [
            'runon.yaml',
            '.runon.yaml',
            os.path.join(xdg.BaseDirectory.xdg_config_home, 'runon', 'runon.yaml'),
            os.path.join(xdg.BaseDirectory.xdg_config_home, 'runon', 'runon.default.yaml')
            ]
    if user_conf:
        conf_list = [ user_conf ]
    for conf in conf_list:
        if os.path.exists(conf):
            return conf
    return None


def read_yaml(conf_file):
    try:
        with open(conf_file, 'r') as file:
            conf = yaml.safe_load(file)
            conf['stamp'] = datetime.datetime.fromtimestamp(os.path.getmtime(conf_file), tz=pytz.UTC)
        return conf
    except yaml.YAMLError as e:
        print(f'ERROR: bad configuration file:')
        print(e)
        sys.exit(1)
    return conf

def list_osnames(conf_file):
    conf = read_yaml(conf_file)
    osnames = []
    for key in conf:
        if (type(conf[key]) is dict) and conf[key].get('image'):
            osnames.append(key)
    return osnames


def load_config(conf_file, osname):
    user_vars = {
            'osname': osname,
            'user': getpass.getuser(),
            'uid': os.getuid(),
            'home': pathlib.Path.home(),
            }
    conf = read_yaml(conf_file)
    osconf = conf.get(osname)
    if not osconf:
        print(f"ERROR: cannot find configuration for distribution {osname}")
        sys.exit(1)

    osconf['stamp'] = conf.get('stamp')
    osconf['osname'] = osname
    for k in [ 'dockerfile', 'packages', 'environment', 'binds' ]:
        if osconf.get(k):
            osconf[k] = [ s.format(**user_vars) for s in osconf[k]]
    return osconf


def make_osname_link(binpath, osname):
    link = os.path.join(os.path.dirname(binpath), osname)
    try:
        os.symlink('runon', link)
    except FileExistsError:
        pass

def make_image_name(osname):
    user = getpass.getuser()
    name = 'runon-{}-{}'.format(osname, user)
    return name


def build_image(client, conf, update, verbose):
    osname = conf.get('osname')
    image_name = 'runon-{}'.format(osname)
    tag = make_image_name(conf['osname'])
    cache_dir = os.path.join(xdg.BaseDirectory.xdg_cache_home, 'runon')
    cache_file = os.path.join(cache_dir, image_name)

    if not update and os.path.exists(cache_file):
        ts_image = datetime.datetime.fromtimestamp(os.path.getmtime(cache_file), tz=pytz.UTC)
        ts_conf = conf.get('stamp')
        if verbose:
            print('config: {}'.format(ts_conf))
            print('image:  {}'.format(ts_image))
        if ts_image and ts_image > ts_conf:
            if verbose:
                print('image: {} up-to-date'.format(image_name))
            image = client.images.get(tag)
            return image

    image = conf.get('image')
    packages = conf['packages']
    dockerfile = conf['dockerfile']
    pkginstall = conf['pkginstall']
    if image:
        dockerfile.insert(0, 'FROM {}'.format(image))
    for p in packages:
        dockerfile.append(pkginstall.format(p))
    try:
        if verbose:
            print('Building image {} ...'.format(tag))
        cmd = ['docker', 'build']
        if update:
            cmd.append('--no-cache')
        cmd += ['-t', tag, '-']
        ret = subprocess.run(cmd,
                             input='\n'.join(dockerfile).encode('utf-8'),
                             stderr=subprocess.STDOUT, check=True)
        image = client.images.get(tag)

        if not os.path.exists(cache_dir):
            os.mkdir(cache_dir)
        with open(cache_file, 'w') as file:
            if verbose:
                print('cache: {}'.format(cache_file))
            file.write('')

    except (docker.errors.BuildError, KeyboardInterrupt, subprocess.CalledProcessError, docker.errors.ImageNotFound) as e:
        print('Build Error: {}'.format(e))
        print()
        print('with dockerfile:')
        for line in dockerfile:
            print('  {}'.format(line))
        sys.exit(1)
    return image


def create_container(client, image, conf, command, verbose):
    volumes = {}
    environment = {}
    for mnt in conf['binds']:
        mnt = mnt.split(':')
        if mnt[-1] in ['ro','rw']:
            mode = mnt[-1]
            del mnt[-1]
        else:
            mode = 'rw'
        mnt = mnt[:2]
        bind = mnt[-1]
        vol = mnt[0]
        volumes[vol] = { 'bind': bind, 'mode': mode }
    hostname = platform.node()
    for v in conf['environment']:
        e = v.split('=')
        if len(e) == 1:
            e.append(os.getenv(e[0]))
        environment[e[0]] = e[1]
    #environment['debian_chroot']=conf['osname']
    user='{}:{}'.format(os.getuid(), os.getgid())
    pwd=os.getcwd()

    container = client.containers.create(image, command,
                                         detach=False, stdin_open=True, tty=True,
                                         auto_remove=True,
                                         hostname=hostname,
                                         volumes=volumes,
                                         environment=environment,
                                         user=user,
                                         network_mode='host',
                                         working_dir=pwd
                                         )
    return container


def run_container(client, container):
    try:
        dockerpty.start(client.api, container.id)
        container.reload()  # to update attrs fields
    except docker.errors.APIError as e:
        print('ERROR: {}'.format(e))
        sys.exit(1)
    ret = container.attrs['State']['ExitCode']
    return ret


def main():
    osname = None
    run_name = os.path.basename(os.sys.argv[0])
    if run_name == "runon":
        pass
    elif run_name.startswith("runon_"):
        osname = run_name[len("runon_"):]
    else:
        osname = run_name

    parser = argparse.ArgumentParser()
    if osname:
        parser.description = 'run commands on "{}" distribution'.format(osname)
    else:
        parser.description = 'run commands on any distribution'
        parser.add_argument('osname',
                            help = 'distribution name to run on, '
                            '"list" to dump all available distributions, '
                            '"edit" to open the current config file in a text editor.')
    parser.epilog = '(c) 2021 Gilles Grandou <gilles@grandou.net>'

    parser.add_argument('-v', '--verbose', action='store_true',
                        help='verbose output')
    parser.add_argument('-c', '--config',
                        help='specify config file')
    parser.add_argument('-u', '--update', action='store_true',
                        help='force image update')
    parser.add_argument('-l', '--link', action='store_true',
                        help='create a symlink to call "osname" as a shortcut to "runon osname"')
    parser.add_argument('command', nargs='*', default=None,
                        help = 'command to execute')

    args = parser.parse_args()
    if osname:
        args.osname = osname

    conf_file = find_config_file(args.config)
    if not conf_file:
        print('ERROR: config file not found')
        sys.exit(1)

    if args.osname == 'list':
        osnames = list_osnames(conf_file)
        print('Available distributions:')
        for o in sorted(osnames):
            print('  {}'.format(o))
        print()
        return 0
    elif args.osname == 'edit':
        cmd = [ 'xdg-open', conf_file ]
        ret = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL);
        return 0

    if args.link:
        make_osname_link(sys.argv[0], args.osname)

    conf = load_config(conf_file, args.osname)

    client = docker.from_env()
    image = build_image(client, conf, args.update, args.verbose)
    container = create_container(client, image, conf, args.command, args.verbose)
    ret = run_container(client, container)
    return ret