#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# deploy.py
#
# Copyright (C) 2019-2022 Franco Masotti (franco \D\o\T masotti {-A-T-} tutanota \D\o\T com)
#
# 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/>.
r"""Deploy the unit files."""

import multiprocessing
import os
import pathlib
import shlex
import shutil
import subprocess

import fpyutils

SRC_DIR = '/home/jobs/services/by-user'
DST_DIR = '/etc/systemd/system'


class UserNotRoot(Exception):
    r"""The user running the script is not root."""


def get_unit_files(start_dir: str = '.', max_depth: int = 0) -> tuple:
    r"""Get the file names of the unit files."""
    unit_file: list = list()
    p = pathlib.Path(start_dir)

    # A file has a fake depth of 1 even if we are at level 0 of the directory tree.
    # The reason of this is how the depth is computed below.
    max_depth += 1

    # Match files and directories.
    for f in p.rglob('*'):
        if f.is_file() and len(
                pathlib.PurePath(f.relative_to(
                    pathlib.PurePath(start_dir))).parts) <= max_depth:
            if pathlib.PurePath(f).match('*.timer') or pathlib.PurePath(
                    f).match('*.service'):
                unit_file.append(f)

    return unit_file


def copy_unit_files(unit_files: list, dst_dir: str = DST_DIR):
    r"""Copy multiple unit files."""
    for f in unit_files:
        shutil.copyfile(
            str(f), str(pathlib.Path(shlex.quote(dst_dir),
                                     shlex.quote(f.name))))


def start_and_enable_unit(unit_name: str):
    r"""Start and enable units."""
    o1 = subprocess.run(shlex.split(
        'systemctl enable --full --quiet --no-reload --show-transaction --no-block '
        + shlex.quote(unit_name)),
                        check=True,
                        capture_output=True).stderr.decode('UTF-8').strip()
    if o1 != str():
        print(o1)

    # If `systemctl is-enabled` returns "disabled" its return value is 1.
    # For this reason check is set to false.
    status = subprocess.run(
        shlex.split('systemctl is-enabled ' + shlex.quote(unit_name)),
        check=False,
        capture_output=True).stdout.decode('UTF-8').strip()
    disable = True
    if status in ['enabled', 'enabled-runtime']:
        disable = False
    elif status in ['static']:
        # Completely disable units without the '[Install]' section.
        disable = True

    if disable:
        try:
            o2 = subprocess.run(shlex.split(
                'systemctl disable --full --quiet --no-reload --show-transaction --no-block --now '
                + shlex.quote(unit_name)),
                                check=True,
                                capture_output=True).stderr.decode(
                                    'UTF-8').strip()
            if o2 != str():
                print(o2)
        except subprocess.CalledProcessError:
            # A new template unit, the ones with 'name@.service' as filename,
            # without the '[Install]' section cannot be stopped nor disabled.
            # See https://wiki.archlinux.org/index.php/Systemd#Using_units
            pass
    else:
        o2 = subprocess.run(shlex.split(
            'systemctl start --full --quiet --no-reload --show-transaction --no-block '
            + shlex.quote(unit_name)),
                            check=True,
                            capture_output=True).stderr.decode(
                                'UTF-8').strip()
        if o2 != str():
            print(o2)


def start_and_enable_units(unit_files: list) -> int:
    r"""Start and enable all services and timers."""
    # Not all services have timer files but all timers have service files.
    # For these cases start and enable the service instead of the timer file.
    concurrent_workers: int = multiprocessing.cpu_count()
    final_list: list = list()
    timer_tmp: list = list()
    service_tmp: list = list()

    for unit in unit_files:
        if unit.name.endswith('.timer'):
            timer_tmp.append(pathlib.Path(unit.name).stem)
        elif unit.name.endswith('.service'):
            service_tmp.append(pathlib.Path(unit.name).stem)

    # O(n**2).
    for s in service_tmp:
        if s in timer_tmp:
            final_list.append(s + '.timer')
        else:
            final_list.append(s + '.service')

    args = iter(final_list)
    worker = start_and_enable_unit

    with multiprocessing.get_context('fork').Pool(
            processes=concurrent_workers) as pool:
        try:
            pool.imap_unordered(worker, args, 20)
        except (KeyboardInterrupt, InterruptedError):
            pool.terminate()
            pool.join()
        else:
            pool.close()
            pool.join()

    return len(final_list)


if __name__ == '__main__':
    if os.getuid() != 0:
        raise UserNotRoot

    unit_files: int = get_unit_files(SRC_DIR, 1)
    copy_unit_files(unit_files, DST_DIR)

    fpyutils.shell.execute_command_live_output('systemctl daemon-reload')
    fpyutils.shell.execute_command_live_output('systemctl reset-failed')

    total_unit_files: int = start_and_enable_units(unit_files)

    fpyutils.shell.execute_command_live_output('systemctl daemon-reload')
    fpyutils.shell.execute_command_live_output('systemctl reset-failed')

    print('\na total of ' + str(total_unit_files) +
          ' units have been handled by deploy.py')
    print('check units by running "systemctl"')