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

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

import os
import sys
import argparse
import psutil
from pathlib import Path
from scylla_util import *
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)
    unit_data = '''
[Unit]
Description=swapfile

[Swap]
What={}

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