#!{{ swift_bin }}/python
# Copyright 2014, Rackspace US, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function
from optparse import OptionParser
from os.path import exists, dirname, join, basename

from swift.common.ring import RingBuilder
from swift.common.ring.utils import parse_builder_ring_filename_args

import json
import sys
import os
import time
from datetime import timedelta


USAGE = "usage: %prog -f <swift_ring.contents> -r <managed_region>"

DEVICE_KEY = "%(ip)s/%(device)s"
FULL_HOST_KEY = "%(ip)s:%(port)dR%(replication_ip)s:" \
                "%(replication_port)d/%(device)s"


class RingValidationError(Exception):
    pass


def update_host_in_ring(ringbuilder, new_host, old_host, old_host_idx,
                        validate=False):
    if new_host.get('zone', 0) != old_host['zone']:
        devstr = DEVICE_KEY % new_host
        raise RingValidationError('Cannot update zone on %s, this can only be '
                                  'done when the drive is added' % devstr)
    if new_host.get('region', 1) != old_host['region']:
        devstr = DEVICE_KEY % new_host
        raise RingValidationError('Cannot update region on %s, this can only '
                                  'be done when the drive is added' % devstr)

    try:
        old_host_str = FULL_HOST_KEY % old_host
        new_host_str = FULL_HOST_KEY % new_host

        new_weight = new_host.get('weight')
        old_weight = old_host.get('weight')

        if new_host_str != old_host_str:
            if not validate:
                ringbuilder.devs[old_host_idx].update(new_host)
                ringbuilder.devs_changed = True
                ringbuilder.version += 1
    except Exception as ex:
        raise RingValidationError(ex)

    if new_weight != old_weight and not validate:
        ringbuilder.set_dev_weight(ringbuilder.devs[old_host_idx]['id'],
                                   new_weight)


def add_host_to_ring(ringbuilder, host, validate=False):
    new_host = {'region': 1, 'zone': 0, 'meta': ''}
    new_host.update(host)
    try:
        if validate:
            ringbuilder.add_dev(new_host)
    except Exception as ex:
        raise RingValidationError(ex)


def build_ring(build_name, repl, min_part_hours, part_power, hosts,
               region=None, validate=False, reset_mph_clock=False):
    # Create the build file
    build_file = "%s.builder" % build_name
    if exists(build_file):
        ringbuilder = RingBuilder.load(build_file)
    else:
        ringbuilder = RingBuilder(part_power, repl, min_part_hours)

    # run some checks
    if repl != ringbuilder.replicas and not validate:
        ringbuilder.set_replicas(repl)

    if min_part_hours != ringbuilder.min_part_hours and not validate:
        ringbuilder.change_min_part_hours(min_part_hours)

    if part_power != ringbuilder.part_power:
        raise RingValidationError(
            'Part power cannot be changed! you must rebuild the ring if you '
            'need to change it.\nRing part power: %s Inventory part power: %s'
            % (ringbuilder.part_power, part_power))

    old_hosts = {}
    for i, dev in enumerate(ringbuilder.devs):
        if dev is not None:
            if region is None or int(region) == int(dev['region']):
                old_hosts[DEVICE_KEY % dev] = i
    for host in hosts:
        host_key = DEVICE_KEY % host
        if region is None or int(region) == int(host['region']):
            if host_key in old_hosts:
                old_host = ringbuilder.devs[old_hosts[host_key]]
                update_host_in_ring(ringbuilder, host, old_host,
                                    old_hosts[host_key], validate=validate)
                old_hosts.pop(host_key)
            else:
                add_host_to_ring(ringbuilder, host, validate=validate)

    if old_hosts and not validate:
        # There are still old hosts, these hosts must've been removed
        try:
            for host, idx in old_hosts.items():
                ringbuilder.remove_dev(ringbuilder.devs[idx]['id'])
        except Exception as ex:
            raise RingValidationError(ex)

    build_file, ring_file = parse_builder_ring_filename_args(('', build_file))
    # serialise to disk before we think about writing the ring
    backup_folder = join(dirname(build_file), 'backups')
    try:
        os.mkdir(backup_folder)
    except OSError:
        if not os.path.isdir(backup_folder):
            raise
    ts = time.time()

    ringbuilder.save(build_file)
    ringbuilder.save(join(backup_folder, '%d.' % ts + basename(build_file)))

    # Rebalance ring
    if not validate:
        if not hosts:
            ringdata = ringbuilder.get_ring()
            ringdata.save(join(backup_folder, '%d.' % ts +
                               basename(ring_file)))
            ringdata.save(ring_file)
        else:
            if reset_mph_clock:
                ringbuilder.pretend_min_part_hours_passed()
            if ringbuilder.min_part_seconds_left > 0:
                raise RingValidationError(
                    'The time between rebalances must be at least '
                    'min_part_hours: %s hours (%s remaining)' %
                    (ringbuilder.min_part_hours,
                     timedelta(seconds=ringbuilder.min_part_seconds_left)))
                exit(2)
            parts, balance, removed_devs = ringbuilder.rebalance()
            try:
                ringbuilder.validate()
            except Exception as ex:
                raise RingValidationError(ex)
            ringbuilder.save(join(backup_folder, '%d.' % ts +
                                  basename(build_file)))
            ringbuilder.save(build_file)
            ringdata = ringbuilder.get_ring()
            ringdata.save(join(backup_folder, '%d.' % ts +
                               basename(ring_file)))
            ringdata.save(ring_file)


def main(setup, region, reset_mph_clock):
    # load the json file
    try:
        with open(setup) as json_stream:
            _contents_file = json.load(json_stream)
    except Exception as ex:
        print("Failed to load json string %s" % ex)
        return 1

    hosts = _contents_file['drives']
    kargs = {'validate': True, 'hosts': hosts, 'region': region,
             'reset_mph_clock': reset_mph_clock}
    ring_call = [
        _contents_file['builder_file'],
        _contents_file['repl_number'],
        _contents_file['min_part_hours'],
        _contents_file['part_power']
    ]

    try:
        build_ring(*ring_call, **kargs)
    except RingValidationError as ex:
        print(ex)
        return 2

    # If the validation passes lets go ahead and build the rings.
    kargs.pop('validate')
    build_ring(*ring_call, **kargs)


if __name__ == "__main__":
    parser = OptionParser(USAGE)
    parser.add_option(
        "-f",
        "--file",
        dest="setup",
        help="Specify the swift ring contents file.",
        metavar="FILE"
    )
    parser.add_option(
        "-r",
        "--region",
        help="Specify the region to manage for the ring file.",
        dest="region",
        type='int',
        metavar="REGION"
    )
    parser.add_option(
        "-p",
        "--pretend_min_part_hours_passed",
        help="Reset the clock on the last time a rebalance happened.",
        dest="reset_mph_clock",
        action="store_true",
        default=False
    )

    options, _args = parser.parse_args(sys.argv[1:])
    if options.setup and not exists(options.setup):
        print("Swift ring contents file not found or doesn't exist")
        parser.print_help()
        sys.exit(1)

    sys.exit(main(options.setup, options.region, options.reset_mph_clock))