#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright 2019-present ScyllaDB
#

#
# SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1

import os
import sys
import shlex
import argparse
import psutil
from pathlib import Path
from scylla_util import swap_exists, out, systemd_unit
from subprocess import run

def GB(n):
    return n * 1024 * 1024 * 1024

def to_GB(n):
    return '{:.2f}'.format(n / 1024 / 1024 / 1024)

def find_mount_point(path):
    path = path.absolute()
    while not path.is_mount():
        path = path.parent
    return path

def get_fs_type(path):
    mnt = find_mount_point(path)
    for part in psutil.disk_partitions():
        if part.mountpoint == str(mnt):
            return part.fstype
    return None

if __name__ == '__main__':
    if os.getuid() > 0:
        print('Requires root permission.')
        sys.exit(1)
    parser = argparse.ArgumentParser(description='Configure swap for Scylla.')
    parser.add_argument('--swap-directory',
                        help='specify swapfile directory', default='/')
    parser.add_argument('--swap-size', type=int,
                        help='specify swapfile size in GB')
    parser.add_argument('--swap-size-bytes', type=int,
                        help='specify swapfile size in bytes')
    args = parser.parse_args()

    if swap_exists():
        print('swap already configured, exiting setup')
        sys.exit(1)

    if args.swap_size and args.swap_size_bytes:
        print("Cannot specify both --swap-size and --swap-size-bytes")
        sys.exit(1)

    swap_directory = Path(args.swap_directory)
    swapfile =  swap_directory / 'swapfile'
    if swapfile.exists():
        print('swapfile {} already exists'.format(swapfile))
        sys.exit(1)

    swapunit_bn = out('systemd-escape -p --suffix=swap {}'.format(swapfile))
    swapunit = Path('/etc/systemd/system/{}'.format(swapunit_bn))
    if swapunit.exists():
        print('swap unit {} already exists'.format(swapunit))
        sys.exit(1)

    diskfree = psutil.disk_usage(args.swap_directory).free
    if args.swap_size or args.swap_size_bytes:
        if args.swap_size:
            swapsize = GB(args.swap_size)
        else:
            swapsize = args.swap_size_bytes
        if swapsize > diskfree:
            print('swap directory {} does not have enough disk space. {}GB space required.'.format(args.swap_directory, to_GB(swapsize)))
            sys.exit(1)
    else:
        memtotal = psutil.virtual_memory().total

        # Scylla document says 'swap size should be set to either total_mem/3 or
        # 16GB - lower of the two', so we need to compare 16g vs memtotal/3 and
        # choose lower one
        # see: https://docs.scylladb.com/faq/#do-i-need-to-configure-swap-on-a-scylla-node
        swapsize = GB(16) if GB(16) < int(memtotal / 3) else int(memtotal / 3)

        # We should not fill entire disk space with swapfile, it's safer to limit
        # swap size 50% of diskfree
        half_of_diskfree = int(diskfree / 2)
        if swapsize > half_of_diskfree:
            # out of disk space, abort setup
            if half_of_diskfree <= GB(1):
                print('swap directory {} does not have enough disk space.')
                sys.exit(1)
            swapsize = half_of_diskfree

    swapsize_mb = int(swapsize / 1024 / 1024)
    fs_type = get_fs_type(swap_directory)
    if fs_type == 'ext4':
        run(f'fallocate -l {swapsize_mb}MiB {swapfile}', shell=True, check=True)
    else:
        run('dd if=/dev/zero of={} bs=1M count={}'.format(swapfile, swapsize_mb), shell=True, check=True)
    swapfile.chmod(0o600)
    run('mkswap -f {}'.format(swapfile), shell=True, check=True)

    mount_point = find_mount_point(swap_directory)
    mount_unit = out(f'systemd-escape -p --suffix=mount {shlex.quote(str(mount_point))}')

    # Add DefaultDependencies=no to the swap unit to avoid getting the default
    # Before=swap.target dependency. We apply this to all clouds, but the
    # requirement came from Azure:
    #
    # On Azure, the swap directory is on the Azure ephemeral disk (mounted on /mnt).
    # However, cloud-init makes this mount (i.e., the mnt.mount unit) depend on
    # the network (After=network-online.target). By extension, this means that
    # the swap unit depends on the network. If we didn't use DefaultDependencies=no,
    # then the swap unit would be part of the swap.target which other services
    # assume to be a local boot target, so we would end up with dependency cycles
    # such as:
    #
    # swap.target -> mnt-swapfile.swap -> mnt.mount -> network-online.target -> network.target -> systemd-resolved.service -> tmp.mount -> swap.target
    #
    # By removing the automatic Before=swap.target, the swap unit is no longer
    # part of swap.target, avoiding such cycles. The swap will still be
    # activated via WantedBy=multi-user.target.
    unit_data = '''
[Unit]
Description=swapfile
DefaultDependencies=no
After={}
Conflicts=umount.target
Before=umount.target

[Swap]
What={}

[Install]
WantedBy=multi-user.target
'''[1:-1].format(mount_unit, swapfile)
    with swapunit.open('w') as f:
        f.write(unit_data)
    systemd_unit.reload()
    swap = systemd_unit(swapunit_bn)
    swap.enable()
    swap.start()
