#!/usr/bin/python3

# Copyright (C) 2017 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 re
import scylla_util
import subprocess
import argparse
import yaml
import logging
import sys
import scylla_blocktune as blocktune

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='IO Setup script for Scylla.')
    parser.add_argument('--ami', dest='ami', action='store_true',
                        help='configure AWS AMI')
    args = parser.parse_args()

    cpudata = scylla_util.scylla_cpuinfo()
    if not scylla_util.is_developer_mode():
        if args.ami:
            idata = scylla_util.aws_instance()

            disk_properties = {}
            disk_properties["mountpoint"] = "/var/lib/scylla"
            nr_disks = len(idata.ephemeral_disks())
            ## both i3 and i2 can run with 1 I/O Queue per shard
            if idata.instance() == "i3.large":
                disk_properties["read_iops"] = 111000
                disk_properties["read_bandwidth"] = 653925080
                disk_properties["write_iops"] = 36800
                disk_properties["write_bandwidth"] = 215066473
            elif idata.instance() == "i3.xlarge":
                disk_properties["read_iops"] = 200800
                disk_properties["read_bandwidth"] = 1185106376
                disk_properties["write_iops"] = 53180
                disk_properties["write_bandwidth"] = 423621267
            elif idata.instance_class() == "i3":
                disk_properties["read_iops"] = 411200 * nr_disks
                disk_properties["read_bandwidth"] = 2015342735 * nr_disks
                disk_properties["write_iops"] = 181500 * nr_disks
                disk_properties["write_bandwidth"] = 808775652 * nr_disks
            elif idata.instance_class() == "i2":
                disk_properties["read_iops"] = 64000 * nr_disks
                disk_properties["read_bandwidth"] = 507338935 * nr_disks
                disk_properties["write_iops"] = 57100 * nr_disks
                disk_properties["write_bandwidth"] = 483141731 * nr_disks
            properties_file = open("/etc/scylla.d/io_properties.yaml", "w")
            yaml.dump({ "disks": [ disk_properties ] }, properties_file,  default_flow_style=False)
            ioconf = open("/etc/scylla.d/io.conf", "w")
            ioconf.write("SEASTAR_IO=\"--io-properties-file={}\"\n".format(properties_file.name))
        else:
            if "SCYLLA_CONF" in os.environ:
                conf_dir = os.environ["SCYLLA_CONF"]
            else:
                conf_dir = "/etc/scylla"
            cfg = yaml.load(open(os.path.join(conf_dir, "scylla.yaml")))
            data_dirs = cfg["data_file_directories"]
            if len(data_dirs) > 1:
                logging.warn("%d data directories found. scylla_io_setup currently lacks support for it, and only %s will be evaluated",
                             len(data_dirs), data_dirs[0])

            data_dir = data_dirs[0]
            if os.path.exists(data_dir) == False:
                logging.error("%s is not found. please check the configuration and run scylla_io_setup again.\n", data_dir)
                sys.exit(1)
            if os.path.isdir(data_dir) == False:
                logging.error("%s is not a directory. please check the configuration and run scylla_io_setup again.\n", data_dir)
                sys.exit(1)
            iotune_args = []
            if cpudata.cpuset():
                iotune_args += [ "--cpuset", ",".join(map(str, cpudata.cpuset())) ]
            elif cpudata.smp():
                iotune_args += [ "--smp", cpudata.smp() ]

            blocktune.tune_fs(data_dir, '2')
            try:
                subprocess.check_call(["iotune",
                                       "--evaluation-directory", data_dir,
                                       "--format", "envfile",
                                       "--options-file", "/etc/scylla.d/io.conf",
                                       "--properties-file", "/etc/scylla.d/io_properties.yaml"] + iotune_args)
            except Exception:
                logging.error("%s did not pass validation tests, it may not be on XFS and/or has limited disk space.\n"
                              "This is a non-supported setup, and performance is expected to be very bad.\n"
                              "For better performance, placing your data on XFS-formatted directories is required.\n"
                              "To override this error, enable developer mode as follow:\n"
                              "sudo /usr/lib/scylla/scylla_dev_mode_setup --developer-mode 1", data_dir)
                sys.exit(1)

