#!/usr/bin/python3

import sys
import os
import logging
import subprocess
import glob
import stat
import shutil
import argparse
import datetime

curdir = os.path.dirname (__file__)
sys.path.insert (0, curdir)
from boottools import utils, apt, btog

os.chdir (curdir)

def _logging (lvl='INFO', logfile='/tmp/boot-tools_installation.log'):
    numeric_level = getattr (logging, lvl.upper(), None)
    if numeric_level is None:
        numeric_level = getattr (logging, 'INFO')

    logging.basicConfig (
        format='%(levelname)s %(asctime)s (%(funcName)s) %(message)s',
        level=numeric_level,
        handlers = [
            logging.FileHandler (logfile),
            logging.StreamHandler (sys.stdout),
        ],
    )
    return logging.getLogger ('boottools')

def _mount_rootfs (btrootfsimg, btrootfsmnt):
    try: utils.mount (btrootfsimg, btrootfsmnt, opts=['-o', 'loop,offset=32256'])
    except:
        logger.error ('mount failed')
        sys.exit (1)

def _get_pxepkg():
    pxepkg = None
    cache = apt.cache_search (['gpxe', 'ipxe'])
    if cache['gpxe']: pxepkg = 'gpxe'
    if cache['ipxe']: pxepkg = 'ipxe'
    if pxepkg is None:
        logger.error ('neither gpxe nor ipxe found in apt cache')
        sys.exit (1)
    logger.info (f'PXE package is "{pxepkg}"')
    return pxepkg

def _mkrootfs (btrootfsimg, btrootfsmnt, btrootfsimglabel, btvirtualdisksize, bttargetdir, osarch):
    logger.info ('Stage 1.1 - create, partition and format the rootfs')
    rc = subprocess.run (f'file "{btrootfsimg}" |grep -q "partition 1 *: ID=0x83"', shell=True).returncode
    if (rc):   ## 'file|grep' failed
        try: btog.mkrootfs (btrootfsimg, btrootfsimglabel, btrootfsmnt, btvirtualdisksize, bttargetdir, osarch)
        except Exception as e:
            logger.error (str (e))
            sys.exit (1)

def _debootstrap (btrootfsimg, btrootfsmnt, osarch, oscodename, oshttp):
    logger.info ('Stage 1.2 - debootstrap system')
    _mount_rootfs (btrootfsimg, btrootfsmnt)
    try: os.stat (os.path.join (btrootfsmnt, 'etc'))
    except:
        logger.debug (f'stat failed, calling btog.debootstrap()')
        try: btog.debootstrap (btrootfsimg, btrootfsmnt, osarch, oscodename, oshttp)
        except Exception as e:
            utils.umount (btrootfsmnt)
            logger.error (str (e))
            sys.exit (1)
    utils.umount (btrootfsmnt)

def _initramfs_version (gitrelease, osrelease, curdir):
    utils.run (['sed', '-i', f'1 s/$/ {gitrelease} ({osrelease})/', f'{curdir}/includes/etc/initramfs-tools/scripts/VERSION.txt'])

def _copy_og_files (btrootfsimg, btrootfsmnt, osdistrib, oscodename):
    _mount_rootfs (btrootfsimg, btrootfsmnt)
    builder   = '/tmp/opengnsys/oglive_builder'
    og_shared = '/tmp/opengnsys/shared'
    btog.copy_og_files (builder, og_shared, btrootfsmnt, osdistrib, oscodename)
    utils.umount (btrootfsmnt)

def _chroot_tasks (cfgfile, curdir, osrelease, osarch):
    logger.debug (f'running \'schroot --chroot IMGogclient -- {curdir}/chroot-tasks.py --osrelease "{osrelease}" --osarch "{osarch}"\'')
    stdout, _ = utils.run (['schroot', '--chroot', 'IMGogclient', '--', f'{curdir}/chroot-tasks.py', '--osrelease', osrelease, '--osarch', osarch, '--config', cfgfile])
    ## this leaves initrd.img-6.8.0-31-generic and vmlinuz-6.8.0-31-generic in /tmp

def _ssh_stuff (btrootfsimg, btrootfsmnt):
    _mount_rootfs (btrootfsimg, btrootfsmnt)
    btog.sysctl (btrootfsmnt)
    btog.ssh_server (btrootfsmnt)
    btog.ssh_client (btrootfsmnt)
    utils.umount (btrootfsmnt)
    ## the end result is:
    ## - there's a new key pair in the VM (or docker container), in /root/.ssh
    ## - there's another new key pair in the rootfs, in /var/lib/tftpboot/ogclient/ogclientmount/root/.ssh
    ## - the two pubkeys (one of each pair) end up being authorised in the rootfs, in /var/lib/tftpboot/ogclient/ogclientmount/root/.ssh/authorized_keys

def _mkinitrd_squashfs_isofs (bttargetdir, osrelease, btrootfsimg, btrootfsmnt, pxepkg, isolinux_tpl, nameisoclient):
    logger.info ('Stage 4.1 - Put initrd in place')
    _mount_rootfs (btrootfsimg, btrootfsmnt)
    btog.move_initrd (bttargetdir, osrelease)

    logger.info ('Stage 4.2 - make squash filesystem')
    btog.mksquashfs (bttargetdir, btrootfsmnt)
    utils.umount (btrootfsmnt)

    logger.info ('Stage 4.3 - make iso filesystem')
    btog.mkisofs (pxepkg, isolinux_tpl, bttargetdir, nameisoclient)

def _main (cfgfile, config, type_client):
    isolinux_tpl     = config['General'].get ('isolinux_template')
    btrootfsimglabel = config['General'].get ('rootfs_image_label')

    logger.info ('OpenGnsys CLIENT installation begins')

    fd = open (f'{curdir}/gitrelease', 'r')     ## per the Dockerfile
    gitrelease = fd.readline().strip()
    fd.close()

    osdistrib, oscodename, osrelease, osarch, oshttp = btog.GetOsInfo (type_client)
    if osdistrib is None:
        logger.error ('GetOsInfo() failed')
        sys.exit (1)
    bttargetdir, btrootfsimg, btrootfsmnt, btvirtualdisksize = btog.GetVar (osarch)
    logger.info (':'.join ([osdistrib, oscodename, osrelease, osarch, oshttp]))

    ## this is convenient in case the previous run failed and we want to run this program again
    try: utils.umount (btrootfsmnt)
    except: pass

    logger.info ('STAGE 1 - create and bootstrap rootfs')
    _mkrootfs (btrootfsimg, btrootfsmnt, btrootfsimglabel, btvirtualdisksize, bttargetdir, osarch)
    _debootstrap (btrootfsimg, btrootfsmnt, osarch, oscodename, oshttp)

    logger.info ('STAGE 2 - copy files to the rootfs')
    today = datetime.datetime.now(datetime.timezone.utc).strftime ('%Y%m%d')
    _initramfs_version (gitrelease+'_'+today, osrelease, curdir)
    _copy_og_files (btrootfsimg, btrootfsmnt, osdistrib, oscodename)

    logger.info ('STAGE 3 - perform tasks within the chroot')
    _chroot_tasks (cfgfile, curdir, osrelease, osarch)

    _ssh_stuff (btrootfsimg, btrootfsmnt)

    logger.info ('STAGE 4 - generate distribution files')
    pxepkg = _get_pxepkg()
    nameisoclient  = '-'.join (['ogLive', oscodename, osrelease, osarch, gitrelease+'_'+today])
    _mkinitrd_squashfs_isofs (bttargetdir, osrelease, btrootfsimg, btrootfsmnt, pxepkg, isolinux_tpl, nameisoclient)

    logger.info ('OpenGnsys installation finished')


if __name__ == '__main__':
    if os.getuid():
        print ('ERROR: this program must run under root privileges!!', file=sys.stderr)
        sys.exit (1)

    parser = argparse.ArgumentParser()
    parser.add_argument ('--loglevel', help='Log level',                  action='store')
    parser.add_argument ('--codename', help='OS codename',                action='store')
    parser.add_argument ('--config',   help='Path to configuration file', action='store')
    args = parser.parse_args()

    cfgfile = args.config or 'mkoglive.cfg'
    config = utils.read_config (cfgfile)
    if config is None:
        print ('ERROR: no configuration found', file=sys.stderr)
        sys.exit (1)
    cfg_loglevel = config['General'].get ('logging_level')

    logger = _logging (args.loglevel or cfg_loglevel)

    _main (cfgfile, config, args.codename)