#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright 2018 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 subprocess
import re
import distro
import shutil

from scylla_util import *

if __name__ == '__main__':
    if os.getuid() > 0:
        print('Requires root permission.')
        sys.exit(1)

    parser = argparse.ArgumentParser(description='Optimize NTP client setting for Scylla.')
    if is_redhat_variant():
        parser.add_argument('--subdomain',
                            help='specify subdomain of pool.ntp.org (ex: centos, fedora or amazon)')
    args = parser.parse_args()

    if is_debian_variant():
        if not shutil.which('ntpd') or not shutil.which('ntpdate'):
            apt_install('ntp ntpdate')
        with open('/etc/ntp.conf') as f:
            conf = f.read()
        if not re.match(r'^server ', conf, flags=re.MULTILINE):
            conf2 = re.findall(r'^(?!pool).*\n', conf, flags=re.MULTILINE)
            with open('/etc/ntp.conf', 'w') as f:
                f.writelines(conf2)
                f.write('server ntp.ubuntu.com iburst')
        ntp = systemd_unit('ntp.service')
        ntp.stop()
        # ignore error, ntpd may able to adjust clock later
        run('ntpdate ntp.ubuntu.com', exception=False)
        ntp.start()

    if is_gentoo_variant():
        if not shutil.which('ntpd'):
            emerge_install('net-misc/ntp')
        with open('/etc/ntp.conf') as f:
            conf = f.read()
        match = re.search(r'^server\s+(\S*)(\s+\S+)?', conf, flags=re.MULTILINE)
        server = match.group(1)
        ntpd = systemd_unit('ntpd.service')
        ntpd.stop()
        run('ntpdate {}'.format(server))
        ntpd.enable()
        ntpd.start()

    if is_redhat_variant():
        if int(distro.major_version()) >= 8:
            if not shutil.which('chronyd'):
                yum_install('chrony')
            with open('/etc/chrony.conf') as f:
                conf = f.read()
            if args.subdomain:
                conf2 = re.sub(r'pool\s+([0-9]+)\.(\S+)\.pool\.ntp\.org', 'pool \\1.{}.pool.ntp.org'.format(args.subdomain), conf, flags=re.MULTILINE)
                with open('/etc/chrony.conf', 'w') as f:
                    f.write(conf2)
                conf = conf2
            chronyd = systemd_unit('chronyd.service')
            chronyd.enable()
            chronyd.restart()
            run('chronyc makestep')
        else:
            if not shutil.which('ntpd') or not shutil.which('ntpdate'):
                yum_install('ntp ntpdate')
            with open('/etc/ntp.conf') as f:
                conf = f.read()
            if args.subdomain:
                conf2 = re.sub(r'server\s+([0-9]+)\.(\S+)\.pool\.ntp\.org', 'server \\1.{}.pool.ntp.org'.format(args.subdomain), conf, flags=re.MULTILINE)
                with open('/etc/ntp.conf', 'w') as f:
                    f.write(conf2)
                conf = conf2
            match = re.search(r'^server\s+(\S*)(\s+\S+)?', conf, flags=re.MULTILINE)
            server = match.group(1)
            ntpd = systemd_unit('ntpd.service')
            ntpd.stop()
            # ignore error, ntpd may able to adjust clock later
            run('ntpdate {}'.format(server), exception=False)
            ntpd.enable()
            ntpd.start()
