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

#
# This file is part of Scylla.
#
# Scylla is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Scylla 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 Scylla.  If not, see <http://www.gnu.org/licenses/>.

import os
import sys
import argparse
import psutil
from pathlib import Path
from scylla_util import *

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

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

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')
    args = parser.parse_args()

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

    diskfree = psutil.disk_usage(args.swap_directory).free
    if args.swap_size:
        swapsize = GB(args.swap_size)
        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)
    swapfile = Path(args.swap_directory) / 'swapfile'
    if swapfile.exists():
        print('swapfile {} already exists'.format(swapfile))
        sys.exit(1)
    run('dd if=/dev/zero of={} bs=1M count={}'.format(swapfile, swapsize_mb))
    swapfile.chmod(0o600)
    run('mkswap -f {}'.format(swapfile))
    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)
    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()
